diff --git a/src/maxtext/__init__.py b/src/maxtext/__init__.py
index 25895b51a9..646fbf7caa 100644
--- a/src/maxtext/__init__.py
+++ b/src/maxtext/__init__.py
@@ -42,3 +42,4 @@
 
 Transformer = models.Transformer
 transformer_as_linen = models.transformer_as_linen
+from_config = model_creation_utils.from_config
diff --git a/tests/assets/logits_generation/generate_grpo_golden_logits.py b/tests/assets/logits_generation/generate_grpo_golden_logits.py
index cae8b9e4d3..9dc78f9df0 100644
--- a/tests/assets/logits_generation/generate_grpo_golden_logits.py
+++ b/tests/assets/logits_generation/generate_grpo_golden_logits.py
@@ -26,6 +26,7 @@
 
 from datasets import load_dataset
 from flax import linen as nn
+from flax import nnx
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
@@ -33,11 +34,12 @@
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
 from maxtext.common.common_types import Array, MODEL_MODE_TRAIN
-from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, generate_completions, grpo_loss_fn
-from maxtext.experimental.rl.grpo_utils import compute_log_probs
+from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, generate_completions, grpo_loss_fn, grpo_loss_fn_nnx
+from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
 from maxtext.inference.maxengine import maxengine
 from maxtext.models import models
 from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
 from tests.post_training.integration.grpo_trainer_correctness_test import prepare_maxtext_inputs
 import numpy as np
 import torch
@@ -46,6 +48,43 @@
 from trl import GRPOConfig, GRPOTrainer
 
 
+def _setup_model(config, mesh, rng):
+  """Builds the model, and for NNX a frozen reference clone, dispatching on pure_nnx.
+
+  Returns (model, reference_model, state). For NNX the model carries its own params
+  (from_pretrained loads the checkpoint or inits) and state is None; for Linen the
+  model is a ToLinen module with a separate decode state.
+  """
+  if config.pure_nnx:
+    model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=rng)
+    return model, nnx.clone(model), None
+  model = models.transformer_as_linen(config=config, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN)
+  init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, config, False, rng)
+  state, state_mesh_annotations = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
+  return model, None, (state, state_mesh_annotations)
+
+
+def _logps(config, model, state, ids, pos, seg, comp_seg):
+  """Policy per-token log-probs, dispatching between NNX and Linen."""
+  if config.pure_nnx:
+    return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False)
+  return compute_log_probs(model, state.params, ids, pos, seg, comp_seg, config, is_train=False)
+
+
+def _reference_logps(config, model, reference_model, reference_params, ids, pos, seg, comp_seg):
+  """Reference per-token log-probs. NNX uses the cloned reference model; Linen uses the saved params."""
+  if config.pure_nnx:
+    return compute_log_probs_nnx(reference_model, ids, pos, seg, comp_seg, config, is_train=False)
+  return compute_log_probs(model, {"params": reference_params}, ids, pos, seg, comp_seg, config, is_train=False)
+
+
+def _grpo_loss(config, model, reference_model, state, reference_params, data, rng):
+  """GRPO loss, dispatching between NNX (reference model) and Linen (reference params)."""
+  if config.pure_nnx:
+    return grpo_loss_fn_nnx(model, config, data, rng, None, reference_model)
+  return grpo_loss_fn(model, config, data, rng, state.params, reference_params)
+
+
 class GRPOTest(unittest.TestCase):
 
   def setUp(self):
@@ -72,28 +111,21 @@ def setUp(self):
     self.rng = jax.random.key(self.cfg.init_weights_seed)
     devices_array = maxtext_utils.create_device_mesh(self.cfg)
     mesh = Mesh(devices_array, self.cfg.mesh_axes)
+    self.mesh = mesh
     # With checkpoint
+    self.model, self.reference_model, linen_state = _setup_model(self.cfg, mesh, self.rng)
     if self.cfg.pure_nnx:
-      # NNX has a different function to init the training state.
-      raise NotImplementedError("Pure NNX support has not been implemented yet.")
+      self.state = None
+      self.state_mesh_shardings = None  # NNX param shardings are derived in the generation step.
     else:
-      self.model = models.transformer_as_linen(config=self.cfg, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN)
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.cfg, False, self.rng)
-    self.state, state_mesh_annotations = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn)
-    self.state_mesh_shardings = nn.logical_to_mesh_sharding(state_mesh_annotations, mesh, self.cfg.logical_axis_rules)
+      self.state, state_mesh_annotations = linen_state
+      self.state_mesh_shardings = nn.logical_to_mesh_sharding(state_mesh_annotations, mesh, self.cfg.logical_axis_rules)
     self.data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
     # Without checkpoint
-    if self.cfg_no_ckpt_loading.pure_nnx:
-      # NNX has a different function to init the training state.
-      raise NotImplementedError("Pure NNX support has not been implemented yet.")
-    else:
-      self.model_no_ckpt_loading = models.transformer_as_linen(
-          config=self.cfg_no_ckpt_loading, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN
-      )
-      init_state_fn = functools.partial(
-          maxtext_utils.init_initial_state, self.model_no_ckpt_loading, None, self.cfg_no_ckpt_loading, False, self.rng
-      )
-    self.state_no_ckpt_loading, _ = maxtext_utils.setup_decode_state(self.cfg_no_ckpt_loading, mesh, None, init_state_fn)
+    self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading, linen_state_no_ckpt = _setup_model(
+        self.cfg_no_ckpt_loading, mesh, self.rng
+    )
+    self.state_no_ckpt_loading = None if self.cfg_no_ckpt_loading.pure_nnx else linen_state_no_ckpt[0]
 
     self.tokenizer_model = transformers.AutoTokenizer.from_pretrained(
         "meta-llama/Llama-3.1-8B",
@@ -181,22 +213,20 @@ def test_w_trl_and_write_golden_data(self):
     input_ids, input_segmentation, input_position, completion_segmentation = prepare_maxtext_inputs(
         self.cfg.prompt, self.tokenizer_model
     )
-    maxtext_per_token_logps, _ = compute_log_probs(
-        self.model,
-        self.state.params,
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps, _ = _logps(
+        self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation
     )
 
-    reference_params = jax.tree.map(jnp.copy, self.state.params["params"])
-    self.state = _merge_grpo_state(self.state, reference_params)
-
-    reference_params_no_ckpt_loading = jax.tree.map(jnp.copy, self.state_no_ckpt_loading.params["params"])
-    self.state_no_ckpt_loading = _merge_grpo_state(self.state_no_ckpt_loading, reference_params_no_ckpt_loading)
+    # The reference is a frozen copy of the step-0 policy. NNX holds it as a cloned
+    # model (built in setUp); Linen snapshots the params and merges them into the state.
+    reference_params = None
+    reference_params_no_ckpt_loading = None
+    if not self.cfg.pure_nnx:
+      reference_params = jax.tree.map(jnp.copy, self.state.params["params"])
+      self.state = _merge_grpo_state(self.state, reference_params)
+    if not self.cfg_no_ckpt_loading.pure_nnx:
+      reference_params_no_ckpt_loading = jax.tree.map(jnp.copy, self.state_no_ckpt_loading.params["params"])
+      self.state_no_ckpt_loading = _merge_grpo_state(self.state_no_ckpt_loading, reference_params_no_ckpt_loading)
 
     data = {
         "prompt_completions": input_ids,
@@ -204,32 +234,27 @@ def test_w_trl_and_write_golden_data(self):
         "prompt_completions_segmentation": input_segmentation,
         "ar_completions_segmentation": completion_segmentation,
     }
-    maxtext_loss, aux = grpo_loss_fn(self.model, self.cfg, data, self.rng, self.state.params, reference_params)
+    maxtext_loss, aux = _grpo_loss(
+        self.cfg, self.model, self.reference_model, self.state, reference_params, data, self.rng
+    )
     # pylint: disable=protected-access
     self.assertEqual(self.trainer._metrics["train"]["kl"][0], aux.avg_kl.tolist())
     self.assertEqual(hf_loss.item(), maxtext_loss.tolist())
     # since this is on-policy
     self.assertEqual(aux.avg_advantage.tolist(), 0.0)
     # since we are at step 0
-    maxtext_per_token_logps, _ = compute_log_probs(
-        self.model,
-        self.state.params,
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps, _ = _logps(
+        self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation
     )
-    maxtext_per_token_logps_ref, _ = compute_log_probs(
+    maxtext_per_token_logps_ref, _ = _reference_logps(
+        self.cfg,
         self.model,
-        {"params": reference_params},
+        self.reference_model,
+        reference_params,
         input_ids,
         input_position,
         input_segmentation,
         completion_segmentation,
-        self.cfg,
-        is_train=False,
     )
     self.assertTrue(
         jax.numpy.allclose(
@@ -243,25 +268,24 @@ def test_w_trl_and_write_golden_data(self):
     # Now that we have ensured that the MaxText implementation is correct
     # let us create a MaxText model without the checkpoint and save the logits
 
-    maxtext_per_token_logps_no_ckpt_loading, _ = compute_log_probs(
+    maxtext_per_token_logps_no_ckpt_loading, _ = _logps(
+        self.cfg_no_ckpt_loading,
         self.model_no_ckpt_loading,
-        self.state_no_ckpt_loading.params,
+        self.state_no_ckpt_loading,
         input_ids,
         input_position,
         input_segmentation,
         completion_segmentation,
-        self.cfg_no_ckpt_loading,
-        is_train=False,
-        rngs=self.rng,
     )
 
-    maxtext_loss, aux = grpo_loss_fn(
-        self.model_no_ckpt_loading,
+    maxtext_loss, aux = _grpo_loss(
         self.cfg_no_ckpt_loading,
+        self.model_no_ckpt_loading,
+        self.reference_model_no_ckpt_loading,
+        self.state_no_ckpt_loading,
+        reference_params_no_ckpt_loading,
         data,
         self.rng,
-        self.state_no_ckpt_loading.params,
-        reference_params_no_ckpt_loading,
     )
 
     engine = maxengine.MaxEngine(self.cfg_no_ckpt_loading_inference)
@@ -274,14 +298,21 @@ def test_w_trl_and_write_golden_data(self):
     )
     prompt_true_length = jnp.array([len(prompt_tokens)] * 4)
     engine_data = {"prompt": prompt, "prompt_true_length": prompt_true_length}
+    if self.cfg_no_ckpt_loading.pure_nnx:
+      # NNX params live on the model; the inference engine is NNX-aware (config.pure_nnx).
+      gen_params = nnx.state(self.model_no_ckpt_loading, nnx.Param)
+      gen_param_shardings = jax.tree.map(lambda _: jax.NamedSharding(self.mesh, jax.sharding.PartitionSpec()), gen_params)
+    else:
+      gen_params = {"params": self.state_no_ckpt_loading.params["params"]}
+      gen_param_shardings = self.state_mesh_shardings.params
     p_generate_completions: Callable[[dict, dict, Array], Array] = jax.jit(
         functools.partial(generate_completions, self.cfg, self.tokenizer_model, engine),
-        in_shardings=(self.data_sharding, self.state_mesh_shardings.params, None),
+        in_shardings=(self.data_sharding, gen_param_shardings, None),
         out_shardings=self.data_sharding,
         donate_argnums=(0,),
     )
     # pylint: disable=not-callable
-    engine_data = p_generate_completions(engine_data, {"params": self.state_no_ckpt_loading.params["params"]}, self.rng)
+    engine_data = p_generate_completions(engine_data, gen_params, self.rng)
     data_to_save = {
         "maxtext_loss": maxtext_loss.tolist(),
         "input_ids": input_ids[0].tolist(),
diff --git a/tests/post_training/integration/grpo_correctness.py b/tests/post_training/integration/grpo_correctness.py
index adefc03a7e..84ffdc85e0 100644
--- a/tests/post_training/integration/grpo_correctness.py
+++ b/tests/post_training/integration/grpo_correctness.py
@@ -18,16 +18,18 @@
 import unittest
 
 from datasets import load_dataset
+from flax import nnx
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from maxtext.configs import pyconfig
 from maxtext.common.common_types import MODEL_MODE_TRAIN
-from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn
-from maxtext.experimental.rl.grpo_utils import compute_log_probs
+from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, grpo_loss_fn_nnx
+from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.models import models
 from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
 import numpy as np
 import pytest
 import torch
@@ -62,12 +64,16 @@ def setUp(self):
     devices_array = maxtext_utils.create_device_mesh(self.cfg)
     mesh = Mesh(devices_array, self.cfg.mesh_axes)
     if self.cfg.pure_nnx:
-      # NNX has a different function to init the training state.
-      raise NotImplementedError("Pure NNX support has not been implemented yet.")
+      # NNX: from_pretrained loads the checkpoint (or inits) into the model, which
+      # carries its own params. The frozen reference is a clone of the policy.
+      self.model = model_creation_utils.from_pretrained(self.cfg, mesh=mesh, rng_key=self.rng)
+      self.reference_model = nnx.clone(self.model)
+      self.state = None
     else:
       self.model = models.transformer_as_linen(config=self.cfg, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN)
       init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.cfg, False, self.rng)
-    self.state, _ = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn)
+      self.reference_model = None
+      self.state, _ = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn)
     self.tokenizer_model = transformers.AutoTokenizer.from_pretrained(
         "meta-llama/Llama-3.1-8B",
         add_bos_token=False,
@@ -140,15 +146,15 @@ def _prepare_trl_inputs(self):
     logits_to_keep = tokenized_inputs["input_ids"].size()[1]
     return input_ids, attention_mask, logits_to_keep
 
-  def test_logits(self):
-    def _prepare_inputs():
-      input_ids = jnp.tile(jnp.array(self.tokenizer_model.encode(self.input_str)), (4, 1))
-      input_segmentation = (input_ids > 0).astype(jnp.int32)
-      input_position = jnp.tile(jnp.arange(input_ids.shape[1]), (4, 1))
-
-      return input_ids, input_segmentation, input_position
-
-    inputs, inputs_segmentation, inputs_position = _prepare_inputs()
+  def _maxtext_logits(self, inputs, inputs_position, inputs_segmentation):
+    """Forward pass logits, dispatching between the NNX and Linen models."""
+    if self.cfg.pure_nnx:
+      return self.model(
+          decoder_input_tokens=inputs,
+          decoder_positions=inputs_position,
+          decoder_segment_ids=inputs_segmentation,
+          enable_dropout=False,
+      )
     logits, _ = self.model.apply(
         self.state.params,
         inputs,
@@ -158,6 +164,64 @@ def _prepare_inputs():
         rngs=self.rng,
         mutable="intermediates",
     )
+    return logits
+
+  def _policy_logps(self, input_ids, input_position, input_segmentation, completion_segmentation):
+    """Policy per-token log-probs, dispatching between NNX and Linen."""
+    if self.cfg.pure_nnx:
+      return compute_log_probs_nnx(
+          self.model, input_ids, input_position, input_segmentation, completion_segmentation, self.cfg, is_train=False
+      )
+    return compute_log_probs(
+        self.model,
+        self.state.params,
+        input_ids,
+        input_position,
+        input_segmentation,
+        completion_segmentation,
+        self.cfg,
+        is_train=False,
+    )
+
+  def _reference_logps(self, input_ids, input_position, input_segmentation, completion_segmentation, reference_params):
+    """Reference per-token log-probs. NNX uses the cloned reference model; Linen uses the saved params."""
+    if self.cfg.pure_nnx:
+      return compute_log_probs_nnx(
+          self.reference_model,
+          input_ids,
+          input_position,
+          input_segmentation,
+          completion_segmentation,
+          self.cfg,
+          is_train=False,
+      )
+    return compute_log_probs(
+        self.model,
+        {"params": reference_params},
+        input_ids,
+        input_position,
+        input_segmentation,
+        completion_segmentation,
+        self.cfg,
+        is_train=False,
+    )
+
+  def _grpo_loss(self, data, reference_params):
+    """GRPO loss, dispatching between NNX (reference model) and Linen (reference params)."""
+    if self.cfg.pure_nnx:
+      return grpo_loss_fn_nnx(self.model, self.cfg, data, self.rng, None, self.reference_model)
+    return grpo_loss_fn(self.model, self.cfg, data, self.rng, self.state.params, reference_params)
+
+  def test_logits(self):
+    def _prepare_inputs():
+      input_ids = jnp.tile(jnp.array(self.tokenizer_model.encode(self.input_str)), (4, 1))
+      input_segmentation = (input_ids > 0).astype(jnp.int32)
+      input_position = jnp.tile(jnp.arange(input_ids.shape[1]), (4, 1))
+
+      return input_ids, input_segmentation, input_position
+
+    inputs, inputs_segmentation, inputs_position = _prepare_inputs()
+    logits = self._maxtext_logits(inputs, inputs_position, inputs_segmentation)
 
     hf_logits = (
         self.hf_model(
@@ -173,15 +237,8 @@ def _prepare_inputs():
   def test_logps(self):
 
     input_ids, input_segmentation, input_position, completion_segmentation = self._prepare_maxtext_inputs()
-    maxtext_per_token_logps, _ = compute_log_probs(
-        self.model,
-        self.state.params,
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps, _ = self._policy_logps(
+        input_ids, input_position, input_segmentation, completion_segmentation
     )
     hf_input_ids, attention_mask, logits_to_keep = self._prepare_trl_inputs()
     with torch.no_grad():
@@ -245,57 +302,33 @@ def test_loss_kl_div(self):
     self.trainer._get_per_token_logps(self.hf_model, hf_input_ids, attention_mask, logits_to_keep)  # pylint: disable=protected-access
 
     input_ids, input_segmentation, input_position, completion_segmentation = self._prepare_maxtext_inputs()
-    maxtext_per_token_logps, _ = compute_log_probs(
-        self.model,
-        self.state.params,
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps, _ = self._policy_logps(
+        input_ids, input_position, input_segmentation, completion_segmentation
     )
 
-    reference_params = jax.tree.map(jnp.copy, self.state.params["params"])
-    self.state = _merge_grpo_state(self.state, reference_params)
+    # The reference is a frozen copy of the step-0 policy. NNX holds it as a cloned
+    # model (built in setUp); Linen snapshots the params and merges them into the state.
+    reference_params = None
+    if not self.cfg.pure_nnx:
+      reference_params = jax.tree.map(jnp.copy, self.state.params["params"])
+      self.state = _merge_grpo_state(self.state, reference_params)
     data = {
         "prompt_completions": input_ids,
         "prompt_completions_position": input_position,
         "prompt_completions_segmentation": input_segmentation,
         "ar_completions_segmentation": completion_segmentation,
     }
-    maxtext_loss, aux = grpo_loss_fn(
-        self.model,
-        self.cfg,
-        data,
-        self.rng,
-        self.state.params,
-        reference_params,
-    )
+    maxtext_loss, aux = self._grpo_loss(data, reference_params)
     self.assertEqual(self.trainer._metrics["train"]["kl"][0], aux.avg_kl.tolist())  # pylint: disable=protected-access
     self.assertEqual(hf_loss.item(), maxtext_loss.tolist())
     # since this is on-policy
     self.assertEqual(aux.avg_advantage.tolist(), 0.0)
     # since we are at step 0
-    maxtext_per_token_logps, _ = compute_log_probs(
-        self.model,
-        self.state.params,
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps, _ = self._policy_logps(
+        input_ids, input_position, input_segmentation, completion_segmentation
     )
-    maxtext_per_token_logps_ref, _ = compute_log_probs(
-        self.model,
-        {"params": reference_params},
-        input_ids,
-        input_position,
-        input_segmentation,
-        completion_segmentation,
-        self.cfg,
-        is_train=False,
+    maxtext_per_token_logps_ref, _ = self._reference_logps(
+        input_ids, input_position, input_segmentation, completion_segmentation, reference_params
     )
     self.assertTrue(
         jax.numpy.allclose(
diff --git a/tests/post_training/integration/grpo_trainer_correctness_test.py b/tests/post_training/integration/grpo_trainer_correctness_test.py
index b880a0e678..c5a4392c98 100644
--- a/tests/post_training/integration/grpo_trainer_correctness_test.py
+++ b/tests/post_training/integration/grpo_trainer_correctness_test.py
@@ -39,15 +39,17 @@
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
 from maxtext.common.common_types import MODEL_MODE_TRAIN
+from flax import nnx
 from maxtext.experimental.rl import grpo_utils
-from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, setup_train_loop
-from maxtext.experimental.rl.grpo_utils import compute_log_probs
+from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, grpo_loss_fn_nnx, setup_train_loop
+from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
 from maxtext.inference import offline_engine
 from maxtext.inference.maxengine import maxengine
 from maxtext.inference.offline_engine import InputData
 from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
 import numpy as np
 import pytest
 import transformers
@@ -69,19 +71,27 @@ def get_golden_data(config):
 
 
 def setup_maxtext_model(config, mesh):
-  """setup maxtext model"""
+  """Sets up the MaxText model.
+
+  Returns (model, state, reference, init_rng, state_mesh_shardings, data_sharding). On
+  the NNX path the model carries its own params (state is None) and reference is a
+  cloned frozen model; on Linen, state is the decode state and reference is a copy of
+  the params merged into it.
+  """
   init_rng = jax.random.PRNGKey(config.init_weights_seed)
-  quant = quantizations.configure_quantization(config)
+  data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
 
   if config.pure_nnx:
-    # NNX has a different function to init the training state.
-    raise NotImplementedError("Pure NNX support has not been implemented yet.")
-  else:
-    maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-    init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng)
+    maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=init_rng)
+    reference = nnx.clone(maxtext_model)
+    # state_mesh_shardings is unused by the live correctness test on the NNX path.
+    return maxtext_model, None, reference, init_rng, None, data_sharding
+
+  quant = quantizations.configure_quantization(config)
+  maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+  init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng)
   state, state_mesh_annotations = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
   state_mesh_shardings = nn.logical_to_mesh_sharding(state_mesh_annotations, mesh, config.logical_axis_rules)
-  data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
   reference_params = jax.tree.map(jnp.copy, state.params["params"])
   state = _merge_grpo_state(state, reference_params)
   return (
@@ -94,6 +104,20 @@ def setup_maxtext_model(config, mesh):
   )
 
 
+def _logps(config, model, state, ids, pos, seg, comp_seg, rngs=None):
+  """Per-token log-probs, dispatching between the NNX and Linen models."""
+  if config.pure_nnx:
+    return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False)
+  return compute_log_probs(model, state.params, ids, pos, seg, comp_seg, config, is_train=False, rngs=rngs)
+
+
+def _grpo_loss(config, model, state, reference, data, rng):
+  """GRPO loss. On NNX `reference` is the frozen reference model; on Linen it is the reference params."""
+  if config.pure_nnx:
+    return grpo_loss_fn_nnx(model, config, data, rng, None, reference)
+  return grpo_loss_fn(model, config, data, rng, state.params, reference)
+
+
 def prepare_maxtext_inputs(input_str, tokenizer_model):
   """prepare maxtext inputs"""
   prompt = tokenizer_model.encode(input_str)
@@ -184,15 +208,14 @@ def test_grpo_trainer_correctness(self):
         self.config.prompt, self.tokenizer_model
     )
     # Obtain per-token logits.
-    maxtext_per_token_logps, _ = compute_log_probs(
+    maxtext_per_token_logps, _ = _logps(
+        self.config,
         maxtext_model,
-        state.params,
+        state,
         input_ids,
         input_position,
         input_segmentation,
         completion_segmentation,
-        self.config,
-        is_train=False,
         rngs=self.rng,
     )
     jax.debug.print(
@@ -225,7 +248,7 @@ def test_grpo_trainer_correctness(self):
         "ar_completions_segmentation": completion_segmentation,
     }
     # Compute the loss and auxiliary values.
-    maxtext_loss, aux = grpo_loss_fn(maxtext_model, self.config, data, rng, state.params, reference_params)
+    maxtext_loss, aux = _grpo_loss(self.config, maxtext_model, state, reference_params, data, rng)
 
     # Assert that the computed loss and auxiliary averages match the golden data.
     self.assertEqual(maxtext_loss.tolist(), golden_data["maxtext_loss"])
diff --git a/tests/post_training/integration/sft_trainer_correctness_test.py b/tests/post_training/integration/sft_trainer_correctness_test.py
index 89ac19d0f3..a94932a182 100644
--- a/tests/post_training/integration/sft_trainer_correctness_test.py
+++ b/tests/post_training/integration/sft_trainer_correctness_test.py
@@ -30,6 +30,7 @@
 import sys
 import unittest
 
+from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
@@ -43,6 +44,8 @@
 from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.utils import maxtext_utils
+from maxtext.utils import maxtext_utils_nnx
+from maxtext.utils import model_creation_utils
 import numpy as np
 import pytest
 from transformers import AutoTokenizer
@@ -115,15 +118,17 @@ def prepare_maxtext_inputs(maxtext_data, config):
 def setup_maxtext_model(config):
   """Setup MaxText model."""
   init_rng = jax.random.PRNGKey(config.init_weights_seed)
-  quant = quantizations.configure_quantization(config)
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = Mesh(devices_array, config.mesh_axes)
   if config.pure_nnx:
-    # NNX has a different function to init the training state.
-    raise NotImplementedError("Pure NNX support has not been implemented yet.")
-  else:
-    maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-    init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng)
+    # NNX model: params live on the module, so there is no separate train state.
+    rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng)
+    with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      maxtext_model = model_creation_utils.from_config(config, mesh=mesh, rngs=rngs)
+    return maxtext_model, None, init_rng
+  quant = quantizations.configure_quantization(config)
+  maxtext_model = models.transformer_as_linen(config=config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+  init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, init_rng)
   state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
   return maxtext_model, state, init_rng
 
@@ -131,6 +136,14 @@ def setup_maxtext_model(config):
 def get_maxtext_logits(config, maxtext_data):
   """Get logits generated by MaxText."""
   maxtext_model, state, rng = setup_maxtext_model(config)
+  if config.pure_nnx:
+    # NNX forward: the model carries its own params and rng state.
+    return maxtext_model(
+        decoder_input_tokens=maxtext_data["inputs"],
+        decoder_positions=maxtext_data["inputs_position"],
+        decoder_segment_ids=maxtext_data["inputs_segmentation"],
+        enable_dropout=False,
+    )
   maxtext_logits, _ = maxtext_model.apply(
       state.params,
       maxtext_data["inputs"],
diff --git a/tests/unit/correctness_tests_nnx_dispatch_test.py b/tests/unit/correctness_tests_nnx_dispatch_test.py
new file mode 100644
index 0000000000..c0177cfb14
--- /dev/null
+++ b/tests/unit/correctness_tests_nnx_dispatch_test.py
@@ -0,0 +1,94 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU unit coverage for the pure-NNX dispatch in the post-training correctness tests.
+
+The correctness tests themselves are TPU-only and skipped (b/425997645), so this
+exercises the changed dispatch code that otherwise has no CPU coverage:
+  - `mt.from_config` is exported (the GRPO trainer calls it)
+  - the SFT correctness test's `setup_maxtext_model` / `get_maxtext_logits` run on
+    both paths (pure_nnx=True -> NNX, pure_nnx=False -> Linen) and stay finite
+
+The GRPO NNX building blocks the other dispatch helpers call
+(`compute_log_probs_nnx`, `grpo_loss_fn_nnx`) are already covered by grpo_nnx_test.
+"""
+
+import os
+import sys
+import unittest
+
+import jax
+import jax.numpy as jnp
+
+import maxtext as mt
+from maxtext.configs import pyconfig
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
+from tests.post_training.integration import sft_trainer_correctness_test as sft
+
+# Small model so the test stays fast and runs on a single CPU device.
+_SMALL = {
+    "base_emb_dim": 16,
+    "base_num_decoder_layers": 2,
+    "base_num_query_heads": 2,
+    "base_num_kv_heads": 2,
+    "base_mlp_dim": 32,
+    "head_dim": 8,
+    "vocab_size": 32,
+    "max_target_length": 8,
+    "per_device_batch_size": 1,
+}
+
+
+def _sft_config(pure_nnx):
+  return pyconfig.initialize(
+      [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")],
+      run_name=f"unit-sft-{pure_nnx}",
+      model_name="default",
+      enable_checkpointing=False,
+      pure_nnx=pure_nnx,
+      enable_nnx=pure_nnx,
+      pure_nnx_decoder=pure_nnx,
+      **_SMALL,
+  )
+
+
+def _fake_data(config):
+  b = jax.device_count() * config.per_device_batch_size
+  length = config.max_target_length
+  return {
+      "inputs": jnp.ones((b, length), jnp.int32),
+      "inputs_position": jnp.broadcast_to(jnp.arange(length, dtype=jnp.int32), (b, length)),
+      "inputs_segmentation": jnp.ones((b, length), jnp.int32),
+  }
+
+
+class CorrectnessTestNNXDispatchTest(unittest.TestCase):
+
+  def test_from_config_is_exported(self):
+    # grpo_trainer calls mt.from_config; the package must export it.
+    self.assertTrue(hasattr(mt, "from_config"))
+
+  def test_sft_logits_nnx_path(self):
+    config = _sft_config(pure_nnx=True)
+    logits = sft.get_maxtext_logits(config, _fake_data(config))
+    self.assertTrue(bool(jnp.isfinite(logits).all()))
+
+  def test_sft_logits_linen_path(self):
+    config = _sft_config(pure_nnx=False)
+    logits = sft.get_maxtext_logits(config, _fake_data(config))
+    self.assertTrue(bool(jnp.isfinite(logits).all()))
+
+
+if __name__ == "__main__":
+  unittest.main()