[NNX] Implement pure-NNX path for post-training correctness tests

[ecnal-cienet]

Description

Implement the pure-NNX path for the post-training correctness / golden-generation tests, which raised NotImplementedError under pure_nnx=True (the default since PR #3526). Done as a dispatch that keeps the Linen path working (Linen stays available for a while).

NNX branch per site:

• Model built with model_creation_utils.from_pretrained(config, mesh, rng_key) (loads a checkpoint when load_parameters_path is set, else inits); frozen reference via nnx.clone.
• Forward = native model(...); per-token log-probs = compute_log_probs_nnx; GRPO loss = grpo_loss_fn_nnx.
• Linen branches unchanged (transformer_as_linen + state.params + compute_log_probs / grpo_loss_fn / _merge_grpo_state).

Also exports from_config from the maxtext package — grpo_trainer calls mt.from_config, which previously did not resolve.

Tests

CPU unit test — tests/unit/correctness_tests_nnx_dispatch_test.py covers the changed code that has no other CPU coverage: the mt.from_config export and the SFT correctness test's own setup_maxtext_model/get_maxtext_logits on both paths (pure_nnx=True/False). The GRPO NNX building blocks (compute_log_probs_nnx, grpo_loss_fn_nnx) are already covered by grpo_nnx_test.py.

JAX_PLATFORMS=cpu python3 -m pytest tests/unit/correctness_tests_nnx_dispatch_test.py -v
# 3 passed

Full tests (TPU host + HF auth + golden data): they're @pytest.mark.skip (b/425997645) + @pytest.mark.tpu_only, so the skip must be lifted first.

export PYTHONPATH=$PWD/src HF_TOKEN=<token>

# SFT — comment the @pytest.mark.skip line, then run on TPU:
sed -i '/@pytest.mark.skip(reason="Logit output test fragile/s/^/# /' \
  tests/post_training/integration/sft_trainer_correctness_test.py
python3 -m pytest -v -s \
  tests/post_training/integration/sft_trainer_correctness_test.py::SFTTrainerCorrectnessTest::test_sft_trainer_correctness
git checkout tests/post_training/integration/sft_trainer_correctness_test.py   # restore the skip

# GRPO trainer correctness — same skip-lift pattern:
sed -i '/@pytest.mark.skip(reason="Logit output test fragile/s/^/# /' \
  tests/post_training/integration/grpo_trainer_correctness_test.py
python3 -m pytest -v -s \
  tests/post_training/integration/grpo_trainer_correctness_test.py::GRPOTrainerCorrectnessTest::test_grpo_trainer_correctness
git checkout tests/post_training/integration/grpo_trainer_correctness_test.py

# GRPO correctness — torch/TRL, no method skip (use the post-train venv):
python3 -m pytest -v -s tests/post_training/integration/grpo_correctness.py

# Golden generator (script, not a pytest test):
python3 -m tests.assets.logits_generation.generate_grpo_golden_logits

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @ecnal-cienet, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @ecnal-cienet, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 I'm sorry @ecnal-cienet, but I was unable to process your request. Please see the logs for more details.