From d095c8f1e6fae3e8c153f88994b7e59d9bb5bfad Mon Sep 17 00:00:00 2001
From: Lance Wang <lance.wang@cienet.com>
Date: Wed, 29 Jul 2026 15:21:46 +0000
Subject: [PATCH 1/5] [NNX] Delete Linen (pre-train 1/3): collapse dispatch in
 pre-train state setup and train loop

The pure_nnx defaults are true, so the Linen branches in the pre-train state
and train-loop path are dead. Collapse them:

- train_utils.setup_train_loop: always build the abstract NNX model and a
  TrainStateNNX init_state_fn; drop the Linen model/TrainState branch and the
  Linen arms of the DiLoCo sharding and debug_sharding blocks.
- maxtext_utils: get_functional_train_with_signature and
  get_functional_eval_with_signature drop the trailing rng in_sharding;
  load_compiled drops the example rng; get_abstract_state delegates to
  get_abstract_state_nnx.

setup_initial_state is deliberately left alone. Its Linen branch is entangled
with the checkpoint restore overlay, and the orbax v1 migration is touching
that code; it is collapsed in a later change.

Tests follow the same narrowing: the Linen-only cases in maxtext_utils_test,
state_dtypes_test and the sharding_compare_test Linen-golden driver go away.
test_deepseek4 is skipped rather than pinned to Linen, since nnx_decoders.py
has no deepseek4 decoder_block branch yet.
---
 src/maxtext/utils/maxtext_utils.py            |  63 +---
 src/maxtext/utils/train_utils.py              | 103 +++---
 .../integration/setup_train_loop_nnx_test.py  |   7 +-
 .../correctness_tests_nnx_dispatch_test.py    |  16 +-
 tests/unit/maxtext_utils_test.py              | 152 ++------
 tests/unit/sharding_compare_test.py           | 326 +-----------------
 tests/unit/state_dtypes_test.py               |  42 +--
 tests/unit/train_compile_test.py              |  40 +--
 8 files changed, 112 insertions(+), 637 deletions(-)

diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 76a9c842e2..8b4ae0c565 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -98,10 +98,7 @@ def get_functional_train_with_signature(
   """Get the shardings (both state and data) for `train_step`."""
   functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
   functional_train.__name__ = "train_step"  # pyrefly: ignore[missing-attribute]
-  if config.pure_nnx:
-    in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
-  else:
-    in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
+  in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
   out_shardings = (state_mesh_shardings, None)  # State, metrics
   static_argnums = ()  # We partial out the static argnums of model and config
   donate_argnums = 0  # This is the index of the state - we allow the compiler to make use of this memory.
@@ -112,10 +109,7 @@ def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shar
   """Get the shardings (both state and data) for `eval_step`."""
   functional_eval = functools.partial(eval_step, model, config)
   functional_eval.__name__ = "eval_step"  # pyrefly: ignore[missing-attribute]
-  if config.pure_nnx:
-    in_shardings = (state_mesh_shardings, data_sharding)  # State, batch (NNX: no rng)
-  else:
-    in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
+  in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
   out_shardings = None  # metrics
   static_argnums = ()  # We partial out the static argnums of model, config
   donate_argnums = ()  # state will be kept instead of being donated in eval_step
@@ -264,11 +258,7 @@ def get_train_input_output_trees(func, input_args, input_kwargs):
 
   serialized_compiled = load_serialized_compiled(config.compiled_trainstep_file)
   shaped_batch = get_shaped_batch(config)
-  if config.pure_nnx:
-    shaped_input_args = (state, shaped_batch)
-  else:
-    example_rng = jax.random.PRNGKey(0)
-    shaped_input_args = (state, shaped_batch, example_rng)
+  shaped_input_args = (state, shaped_batch)
   shaped_input_kwargs = {}
   in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs)
   p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices)
@@ -1864,51 +1854,8 @@ def get_logical_annotations(config, mesh, init_state_fn):
 
 
 def get_abstract_state(config, mesh, init_state_fn, is_training=True):
-  """Get a shaped abstraction of the state (including optimizer)"""
-  if config.pure_nnx:
-    return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)
-
-  init_state_partial = init_state_fn
-
-  with nn_partitioning.axis_rules(config.logical_axis_rules):
-    abstract_state = jax.eval_shape(init_state_partial)
-
-  state_logical_annotations = nn.get_partition_spec(abstract_state)
-
-  state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules)
-  if is_training and config.shard_optimizer_over_data:
-    # Add data to sharding for optimizer state
-    state_mesh_shardings = state_mesh_shardings.replace(
-        opt_state=jax.tree.map_with_path(
-            functools.partial(sharding.add_data_to_sharding, mesh),
-            max_utils.unbox_logicallypartioned(abstract_state).opt_state,
-            state_mesh_shardings.opt_state,
-        )
-    )
-  if is_training and config.optimizer_memory_host_offload:
-    opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state)
-    state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state)
-  if is_training and config.parameter_memory_host_offload:
-    assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading."
-
-    def move(path, x):
-      max_logging.log(f"max_utils.py: Moving {path} to host")
-      return x.with_memory_kind(kind="pinned_host")
-
-    params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params)
-    state_mesh_shardings = state_mesh_shardings.replace(params=params)
-
-  abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape()
-
-  unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state)
-  # Initialization
-  with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
-    state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations)
-  return (
-      unboxed_abstract_sharded_state,
-      state_mesh_annotations,
-      state_mesh_shardings,
-  )
+  """Get a shaped abstraction of the state (including optimizer)."""
+  return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)
 
 
 def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True):
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 35f47e590d..493c932dfa 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/src/maxtext/utils/train_utils.py
@@ -20,7 +20,6 @@
 import optax
 import functools
 import orbax.checkpoint.pathways as ocp_pathways
-from functools import partial
 
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
@@ -252,32 +251,25 @@ def setup_train_loop(config, recorder, devices=None):
   from maxtext.input_pipeline.input_pipeline_interface import create_data_iterator
 
   with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT):
-    is_training = True
     init_rng = jax.random.PRNGKey(config.init_weights_seed)
     mesh = maxtext_utils.get_mesh_from_config(config, devices)
     context_parallel_size = mesh.shape.get(config.context_sharding, 1)
-    if config.pure_nnx:
-      # Create abstract NNX model.
-      _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
-    else:
-      model = model_creation_utils.from_config(config, devices)
+    # Create abstract NNX model.
+    _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
     learning_rate_schedule, tx = create_training_optimizer(config, model)
 
-    if config.pure_nnx:
-      # For NNX, the train state is wrapped in the TrainStateNNX module.
-      def create_train_state_fn():
-        model = _create_model_partial()
-        wrt = (
-            getattr(nnx, "LoRAParam", nnx.Param)
-            if getattr(getattr(config, "lora", None), "enable_lora", False)
-            else nnx.Param
-        )
-        optimizer = nnx.Optimizer(model, tx, wrt=wrt)
-        return train_state_nnx.TrainStateNNX(model, optimizer)
-
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng)
+    # The train state is wrapped in the TrainStateNNX module.
+    def create_train_state_fn():
+      model = _create_model_partial()
+      wrt = (
+          getattr(nnx, "LoRAParam", nnx.Param)
+          if getattr(getattr(config, "lora", None), "enable_lora", False)
+          else nnx.Param
+      )
+      optimizer = nnx.Optimizer(model, tx, wrt=wrt)
+      return train_state_nnx.TrainStateNNX(model, optimizer)
+
+    init_state_fn = create_train_state_fn
     checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn)
     if checkpoint_manager is not None:
       checkpoint_step = checkpointing.latest_step(checkpoint_manager)
@@ -336,24 +328,23 @@ def create_train_state_fn():
     state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
-    if config.pure_nnx:
-      if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None):
-        # Restore standalone LoRA adapter weights onto the base model state after initialization.
-        target_model_state = (
-            state["model"]
-            if (isinstance(state, (nnx.State, dict)) and "model" in state)
-            else getattr(state, "model", state)
-        )
-        # pyrefly: ignore[bad-argument-type]
-        lora_utils.restore_lora_from_path(target_model_state, config)
-        _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True)
-      with nn_partitioning.axis_rules(config.logical_axis_rules):
-        # We only need the graphdef here; it's merged with state below. Avoid
-        # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
-        # under jax.set_mesh(mesh) and rejects any logical name missing from
-        # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
-        # without a mesh skips sharding resolution, so it avoids the crash.
-        state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
+    if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None):
+      # Restore standalone LoRA adapter weights onto the base model state after initialization.
+      target_model_state = (
+          state["model"]
+          if (isinstance(state, (nnx.State, dict)) and "model" in state)
+          else getattr(state, "model", state)
+      )
+      # pyrefly: ignore[bad-argument-type]
+      lora_utils.restore_lora_from_path(target_model_state, config)
+      _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True)
+    with nn_partitioning.axis_rules(config.logical_axis_rules):
+      # We only need the graphdef here; it's merged with state below. Avoid
+      # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
+      # under jax.set_mesh(mesh) and rejects any logical name missing from
+      # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
+      # without a mesh skips sharding resolution, so it avoids the crash.
+      state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
 
     if isinstance(state, diloco.DiLoCoTrainState):
       state_params = state.params
@@ -361,13 +352,10 @@ def create_train_state_fn():
         _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
       else:
         state_mesh_shardings_params = state_mesh_shardings.params
-    elif config.pure_nnx:
+    else:
       with nn_partitioning.axis_rules(config.logical_axis_rules):
         _, state_params, _ = nnx.split(state.model, nnx.Param, ...)
         _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
-    else:
-      state_params = state.params
-      state_mesh_shardings_params = state_mesh_shardings.params
 
     if config.enable_diloco:
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
@@ -405,28 +393,21 @@ def create_train_state_fn():
 
     # print weights sharding info under debug sharding mode
     if config.debug_sharding:
-      if config.pure_nnx:
-        # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
-        # probably already lost the logical partition info at this moment.
-        logical_annotations_params = None
-      else:
-        logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn)
-        logical_annotations_params = logical_annotations.params
+      # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
+      # probably already lost the logical partition info at this moment.
+      logical_annotations_params = None
 
       max_utils.print_non_trivial_mesh_axis(model.mesh)  # pyrefly: ignore[missing-attribute]
       maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params)
 
-  if config.pure_nnx:
-    if config.enable_diloco:
-      # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
-      # train step needs that graphdef as jit_model; the wrapper passes through as state.
-      train_state = state
-      model = state_graphdef  # pyrefly: ignore[unbound-name]
-    else:
-      train_state = nnx.merge(state_graphdef, state)  # pyrefly: ignore[unbound-name]
-      model = train_state.model
-  else:
+  if config.enable_diloco:
+    # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
+    # train step needs that graphdef as jit_model; the wrapper passes through as state.
     train_state = state
+    model = state_graphdef  # pyrefly: ignore[unbound-name]
+  else:
+    train_state = nnx.merge(state_graphdef, state)  # pyrefly: ignore[unbound-name]
+    model = train_state.model
 
   return (
       init_rng,
diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py
index 7a26cec21b..a3b9f31a91 100644
--- a/tests/integration/setup_train_loop_nnx_test.py
+++ b/tests/integration/setup_train_loop_nnx_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Integration test for setup_train_loop with pure_nnx=True.
+"""Integration test for setup_train_loop on the NNX path.
 
 setup_train_loop wires together create_nnx_abstract_model, the training
 optimizer,
@@ -43,7 +43,6 @@ def _tiny_nnx_pyconfig(**overrides):
       "enable_checkpointing": False,
       "dataset_type": "synthetic",
       "model_name": "default",
-      "pure_nnx": True,
       "per_device_batch_size": 1.0,
       "base_emb_dim": 8,
       "base_num_query_heads": 4,
@@ -68,7 +67,7 @@ def _tiny_nnx_pyconfig(**overrides):
 class SetupTrainLoopNNXIntegrationTest(unittest.TestCase):
   """End-to-end check that setup_train_loop returns a usable TrainStateNNX."""
 
-  def test_pure_nnx_setup_returns_train_state_nnx(self):
+  def test_setup_returns_train_state_nnx(self):
     config = _tiny_nnx_pyconfig()
 
     (
@@ -126,7 +125,7 @@ def test_load_balanced_cp_keeps_checkpoint_iterator_unwrapped(self):
     self.assertNotIsInstance(data_iterator, train_utils._ReorderedDataIterator)
     self.assertIsInstance(eval_data_iterator, train_utils._ReorderedDataIterator)
 
-  def test_pure_nnx_setup_param_only_split_matches_model(self):
+  def test_setup_param_only_split_matches_model(self):
     """nnx.split(state.model, nnx.Param, ...) must yield a non-empty Param tree
 
     whose structure matches state_mesh_shardings.model after the same split.
diff --git a/tests/unit/correctness_tests_nnx_dispatch_test.py b/tests/unit/correctness_tests_nnx_dispatch_test.py
index 163a5a836e..627980bc6a 100644
--- a/tests/unit/correctness_tests_nnx_dispatch_test.py
+++ b/tests/unit/correctness_tests_nnx_dispatch_test.py
@@ -18,7 +18,7 @@
 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 NNX path 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.
@@ -54,15 +54,12 @@
 }
 
 
-def _sft_config(pure_nnx):
+def _sft_config():
   return pyconfig.initialize(
       [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")],
-      run_name=f"unit-sft-{pure_nnx}",
+      run_name="unit-sft-nnx",
       model_name="default",
       enable_checkpointing=False,
-      pure_nnx=pure_nnx,
-      enable_nnx=pure_nnx,
-      pure_nnx_decoder=pure_nnx,
       **_SMALL,
   )
 
@@ -84,12 +81,7 @@ def test_from_config_is_exported(self):
     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)
+    config = _sft_config()
     logits = sft.get_maxtext_logits(config, _fake_data(config))
     self.assertTrue(bool(jnp.isfinite(logits).all()))
 
diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py
index bf3062d5f2..5f128887cf 100644
--- a/tests/unit/maxtext_utils_test.py
+++ b/tests/unit/maxtext_utils_test.py
@@ -16,7 +16,6 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass, field
-import functools
 from types import SimpleNamespace
 from typing import Any, Sequence
 import unittest
@@ -31,11 +30,9 @@
 import jax.numpy as jnp
 from jax.sharding import AxisType, Mesh, NamedSharding, PartitionSpec
 from maxtext.common import train_state_nnx
-from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN, ShardMode
 from maxtext.configs import pyconfig
+from maxtext.common.common_types import DecoderBlockType, ShardMode
 from maxtext.inference import inference_utils
-from maxtext.layers import quantizations
-from maxtext.models import models
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils import maxtext_utils_nnx
@@ -47,8 +44,6 @@
 import optax
 import pytest
 
-Transformer = models.transformer_as_linen
-
 
 class TestGradientClipping(unittest.TestCase):
   """test class for gradient clipping"""
@@ -363,50 +358,29 @@ def setUp(self):
     self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False)
     devices_array = maxtext_utils.create_device_mesh(self.config)
     self.mesh = Mesh(devices_array, self.config.mesh_axes)
-    quant = quantizations.configure_quantization(self.config)
-    if self.config.pure_nnx:
-      self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
-    else:
-      self.model = models.transformer_as_linen(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+    self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
 
   def test_setup_decode_state(self):
-    rng = random.PRNGKey(0)
-    if self.config.pure_nnx:
+    def create_train_state_fn():
+      nnx_model = self._create_model_partial()
+      return train_state_nnx.TrainStateNNX(nnx_model, None)
 
-      def create_train_state_fn():
-        nnx_model = self._create_model_partial()
-        return train_state_nnx.TrainStateNNX(nnx_model, None)
-
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng)
+    init_state_fn = create_train_state_fn
     state, _ = maxtext_utils.setup_decode_state(self.config, self.mesh, None, init_state_fn)
-    if self.config.pure_nnx:
-      self.assertNotIn("optimizer", state)
-    else:
-      self.assertEqual(state.tx, None)
-      self.assertEqual(state.opt_state, {})
+    self.assertNotIn("optimizer", state)
 
   def test_setup_initial_state(self):
-    rng = random.PRNGKey(0)
     tx = optax.adam(learning_rate=0.001)
-    if self.config.pure_nnx:
 
-      def create_train_state_fn():
-        nnx_model = self._create_model_partial()
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    def create_train_state_fn():
+      nnx_model = self._create_model_partial()
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, tx, self.config, True, rng)
+    init_state_fn = create_train_state_fn
     state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn)
     self.assertFalse(was_restored)
-    if self.config.pure_nnx:
-      self.assertIsNotNone(state.optimizer)
-    else:
-      self.assertEqual(state.tx, tx)
-      self.assertNotEqual(state.opt_state, {})
+    self.assertIsNotNone(state.optimizer)
 
 
 class MaxUtilsPpAsDp(unittest.TestCase):
@@ -1034,9 +1008,8 @@ def train_step(_model, _config, _state_shardings, _params_shardings, state, _bat
 
     return train_step
 
-  def _make_mock_config(self, pure_nnx=False):
+  def _make_mock_config(self):
     cfg = MagicMock()
-    cfg.pure_nnx = pure_nnx
     return cfg
 
   def test_returns_five_tuple(self):
@@ -1053,20 +1026,11 @@ def test_functional_train_has_correct_name(self):
     )
     self.assertEqual(fn.__name__, "train_step")
 
-  def test_linen_in_shardings_includes_rng(self):
-    """pure_nnx=False: in_shardings should be (state, batch, rng)."""
-    step = self._make_mock_step()
-    _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature(
-        step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=False)
-    )
-    self.assertEqual(len(in_shardings), 3)
-    self.assertIsNone(in_shardings[2])  # rng sharding is None
-
   def test_nnx_in_shardings_excludes_rng(self):
-    """pure_nnx=True: in_shardings should be (state, batch) — no rng slot."""
+    """in_shardings should be (state, batch) — no rng slot."""
     step = self._make_mock_step()
     _, in_shardings, _, _, _ = maxtext_utils.get_functional_train_with_signature(
-        step, "data_sharding", "state_shardings", "model", self._make_mock_config(pure_nnx=True)
+        step, "data_sharding", "state_shardings", "model", self._make_mock_config()
     )
     self.assertEqual(len(in_shardings), 2)
 
@@ -1102,9 +1066,8 @@ def eval_step(_model, _config, _state, _batch, _rng=None):
 
     return eval_step
 
-  def _make_mock_config(self, pure_nnx=False):
+  def _make_mock_config(self):
     cfg = MagicMock()
-    cfg.pure_nnx = pure_nnx
     return cfg
 
   def test_returns_five_tuple(self):
@@ -1132,21 +1095,13 @@ def test_donate_argnums_is_empty(self):
     self.assertEqual(donate_argnums, ())
 
   def test_nnx_in_shardings_excludes_rng(self):
-    """pure_nnx=True: in_shardings should be (state, batch) — no rng slot."""
+    """in_shardings should be (state, batch) — no rng slot."""
     step = self._make_mock_eval_step()
     _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature(
-        step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=True)
+        step, "batch_sharding", "state_sharding", "model", self._make_mock_config()
     )
     self.assertEqual(len(in_shardings), 2)
 
-  def test_linen_in_shardings_includes_rng(self):
-    """pure_nnx=False: in_shardings should be (state, batch, rng)."""
-    step = self._make_mock_eval_step()
-    _, in_shardings, _, _, _ = maxtext_utils.get_functional_eval_with_signature(
-        step, "batch_sharding", "state_sharding", "model", self._make_mock_config(pure_nnx=False)
-    )
-    self.assertEqual(len(in_shardings), 3)
-
 
 class TestGetShapedBatch(unittest.TestCase):
   """Tests for get_shaped_batch."""
@@ -1458,39 +1413,20 @@ def setUp(self):
     self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False)
     devices_array = maxtext_utils.create_device_mesh(self.config)
     self.mesh = Mesh(devices_array, self.config.mesh_axes)
-    quant = quantizations.configure_quantization(self.config)
-    if self.config.pure_nnx:
-      self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
-    else:
-      self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+    self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
 
   def test_setup_training_state_returns_train_state(self):
-    rng = jax.random.PRNGKey(0)
     tx = optax.adam(learning_rate=0.001)
-    if self.config.pure_nnx:
 
-      def create_train_state_fn():
-        nnx_model = self._create_model_partial()
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    def create_train_state_fn():
+      nnx_model = self._create_model_partial()
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = functools.partial(
-          maxtext_utils.init_initial_state,
-          self.model,
-          tx,
-          self.config,
-          True,
-          rng,
-      )
+    init_state_fn = create_train_state_fn
     state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn)
     self.assertFalse(was_restored)
-    if self.config.pure_nnx:
-      self.assertIsNotNone(state.optimizer)
-    else:
-      self.assertEqual(state.tx, tx)
-      self.assertNotEqual(state.opt_state, {})
+    self.assertIsNotNone(state.optimizer)
 
 
 class TestGetLogicalAnnotations(unittest.TestCase):
@@ -1500,36 +1436,20 @@ def setUp(self):
     self.config = pyconfig.initialize([None, get_test_config_path()], enable_checkpointing=False)
     devices_array = maxtext_utils.create_device_mesh(self.config)
     self.mesh = Mesh(devices_array, self.config.mesh_axes)
-    quant = quantizations.configure_quantization(self.config)
-    if self.config.pure_nnx:
-      self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
-    else:
-      self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+    self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh)
     self.rng = jax.random.PRNGKey(0)
     self.tx = optax.adam(learning_rate=0.001)
 
   def test_returns_partition_spec_tree(self):
-    if self.config.pure_nnx:
-
-      def create_train_state_fn():
-        nnx_model = self._create_model_partial()
-        optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
-
-      init_state_fn = create_train_state_fn
-      annotations = maxtext_utils_nnx.get_partition_spec_nnx(
-          maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2]
-      )
-    else:
-      init_state_fn = functools.partial(
-          maxtext_utils.init_initial_state,
-          self.model,
-          self.tx,
-          self.config,
-          True,
-          self.rng,
-      )
-      annotations = maxtext_utils.get_logical_annotations(self.config, self.mesh, init_state_fn)
+    def create_train_state_fn():
+      nnx_model = self._create_model_partial()
+      optimizer = nnx.Optimizer(nnx_model, self.tx, wrt=nnx.Param)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+
+    init_state_fn = create_train_state_fn
+    annotations = maxtext_utils_nnx.get_partition_spec_nnx(
+        maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2]
+    )
     # Result should be a pytree with PartitionSpec leaves
     leaves = jax.tree_util.tree_leaves(annotations)
     self.assertGreater(len(leaves), 0)
diff --git a/tests/unit/sharding_compare_test.py b/tests/unit/sharding_compare_test.py
index aaefcb9020..2331d54cf3 100644
--- a/tests/unit/sharding_compare_test.py
+++ b/tests/unit/sharding_compare_test.py
@@ -12,325 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Compare expected sharding of models with actual sharding of models."""
+"""Compare expected sharding of models with actual sharding of models.
 
-import functools
-import hashlib
-import json
-import os
-import jax
-import jax.numpy as jnp
-from maxtext.configs import pyconfig
-from maxtext.utils import maxtext_utils
-from maxtext.utils.sharding import clear_input_shardings_dump
-# import optax
-
-from maxtext.layers import quantizations
-from maxtext.models import models
-from maxtext.optimizers import optimizers
-from maxtext.trainers.pre_train.train_compile import get_shaped_inputs, get_topology_mesh, validate_config
-from tests.utils.sharding_dump import TEST_CASES, load_json, input_sharding_to_json, named_shardings_to_json, partition_specs_to_json
-from tests.utils.test_helpers import get_test_config_path
-import pytest
-
-Transformer = models.transformer_as_linen
-
-
-def compute_checksum(d: dict) -> str:
-  """Compute a checksum (SHA256) of a dictionary."""
-  # Serialize the dictionary into a JSON string (ensuring consistent ordering of keys)
-  json_str = json.dumps(d, sort_keys=True)
-
-  # Compute the SHA256 checksum of the serialized string
-  checksum = hashlib.sha256(json_str.encode("utf-8")).hexdigest()
-
-  return checksum
-
-
-def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_name: str) -> bool:
-  """Compare two json files and print the differences if any."""
-  keys1 = set(json1.keys())
-  keys2 = set(json2.keys())
-
-  only_in_1 = keys1 - keys2
-  only_in_2 = keys2 - keys1
-  shared_keys = keys1 & keys2
-
-  has_diff = False
-
-  if only_in_1:
-    print(f"Keys only in {model1_name}:")
-    for k in sorted(only_in_1):
-      print(f"  {k}")
-    has_diff = True
-
-  if only_in_2:
-    print(f"Keys only in {model2_name}:")
-    for k in sorted(only_in_2):
-      print(f"  {k}")
-    has_diff = True
-
-  for key in sorted(shared_keys):
-    entry1 = json1[key]
-    entry2 = json2[key]
-
-    if isinstance(entry1, dict) and isinstance(entry2, dict):
-      mesh1 = entry1.get("mesh", {})
-      mesh2 = entry2.get("mesh", {})
-
-      spec1 = entry1.get("partition_spec", [])
-      spec2 = entry2.get("partition_spec", [])
-
-      shape1 = entry1.get("shape")
-      shape2 = entry2.get("shape")
-
-      if mesh1 != mesh2:
-        print(f"\nMesh mismatch at '{key}':")
-        print(f"  {model1_name}: {mesh1}")
-        print(f"  {model2_name}: {mesh2}")
-        has_diff = True
-
-      if spec1 != spec2:
-        print(f"\nPartitionSpec mismatch at '{key}':")
-        print(f"  {model1_name}: {spec1}")
-        print(f"  {model2_name}: {spec2}")
-        has_diff = True
-
-      if shape1 != shape2:
-        print(f"\nShape mismatch at '{key}':")
-        print(f"  {model1_name}: {shape1}")
-        print(f"  {model2_name}: {shape2}")
-        has_diff = True
-
-    else:
-      print(f"\nFormat mismatch at '{key}':")
-      print(f"  {model1_name} type: {type(entry1)}")
-      print(f"  {model2_name} type: {type(entry2)}")
-      has_diff = True
-
-  return has_diff
-
-
-# Requires JAX TPU support to generate the simulated TPU topology.
-@pytest.mark.tpu_backend
-@pytest.mark.parametrize("model_name, topology, num_slice, custom_mesh_and_rule, overrides", TEST_CASES)
-def test_sharding_dump_for_model(
-    model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple
-) -> None:
-  """
-  Test sharding configurations from train_compile.get_shaped_inputs.
-  This test verifies that the sharding configurations for various models and topologies remain consistent with golden files.
-  """
-  params = [
-      "/deps/MaxText/tests/unit/sharding_compare_test",
-      get_test_config_path(),
-      f"compile_topology={topology}",
-      f"compile_topology_num_slices={num_slice}",
-      f"model_name={model_name}",
-      "log_config=false",
-      "debug_sharding=true",  # for input sharding dump
-      "pure_nnx=False",
-      "enable_nnx=False",
-      "pure_nnx_decoder=False",
-  ]
-  if custom_mesh_and_rule:
-    params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
-  if overrides:
-    params.extend(overrides)
-
-  root_dir = "tests/utils/sharding_info"
-  rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default"
-  if overrides:
-    rule_name += "_" + "_".join(overrides)
-  base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name)
-
-  named_json_path = os.path.join(base_path, "named_shardings.json")
-  logical_json_path = os.path.join(base_path, "logical_shardings.json")
-  input_json_path = os.path.join(base_path, "input_shardings.json")
-
-  if not os.path.exists(named_json_path):
-    pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}")
-    return
-  if not os.path.exists(logical_json_path):
-    pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}")
-    return
-  if not os.path.exists(input_json_path):
-    pytest.skip(f"Missing input_shardings.json for {model_name} {topology} slice {num_slice}")
-    return
-
-  config = pyconfig.initialize(params)
-  validate_config(config)
-
-  clear_input_shardings_dump()
-  topology_mesh = get_topology_mesh(config)
-  learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config)
-  optimizers.get_optimizer(config, learning_rate_schedule)
-  shaped_train_args, _, state_mesh_shardings, logical_shardings, _ = get_shaped_inputs(topology_mesh, config)
-
-  error_messages = []
-
-  # 1. Compare Named Shardings
-  actual_named = named_shardings_to_json(state_mesh_shardings, shaped_train_args[0])
-  expected_named = load_json(named_json_path)
-  # calculate checksum
-  actual_named_sum = compute_checksum(actual_named)
-  expected_named_sum = compute_checksum(expected_named)
-  named_match = actual_named_sum == expected_named_sum
-
-  if not named_match:
-    print(f"\n[FAIL] Physical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True)
-    compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)")
-    error_messages.append(f" Physical sharding mismatch for {model_name} on {topology} slice {num_slice}")
-
-  # 2. Compare Logical Shardings
-  actual_logical = partition_specs_to_json(logical_shardings, shaped_train_args[0])
-  expected_logical = load_json(logical_json_path)
-  # calculate checksum
-  actual_logical_sum = compute_checksum(actual_logical)
-  expected_logical_sum = compute_checksum(expected_logical)
-  logical_match = actual_logical_sum == expected_logical_sum
-
-  if not logical_match:
-    print(f"\n[FAIL] Logical Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True)
-    compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)")
-    error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}")
-
-  # 3. Compare Input Shardings
-  actual_input = input_sharding_to_json()
-  expected_input = load_json(input_json_path)
-  # calculate checksum
-  actual_input_sum = compute_checksum(actual_input)
-  expected_input_sum = compute_checksum(expected_input)
-
-  input_match = actual_input_sum == expected_input_sum
-
-  if not input_match:
-    print(f"\n[FAIL] Input Sharding Mismatch: {model_name} {topology} slice {num_slice}", flush=True)
-    # compare_sharding_jsons(expected_input, "Expected (Input)", actual_input, "Actual (Input)")
-    error_messages.append(f"Input sharding mismatch for {model_name} on {topology} slice {num_slice}")
-
-  assert not error_messages, "\n".join(error_messages)
-
-
-@pytest.fixture(
-    scope="module",
-    params=[pytest.param(case, id=f"{case[0]}-{case[1]}-{case[2]}-{case[3]}-{''.join(case[4])}") for case in TEST_CASES],
-)
-def abstract_state_and_shardings(request):
-  """Pytest fixture to set up model, config, and generate abstract state once per test case."""
-  model_name, topology, num_slice, custom_mesh_and_rule, overrides = request.param
-  print(
-      f"Testing model: {model_name}, topology: {topology}, num_slices: {num_slice}, "
-      "rule: {custom_mesh_and_rule}, overrides: {overrides}",
-      flush=True,
-  )
-  params = [
-      "/deps/MaxText/tests/unit/sharding_compare_test",
-      get_test_config_path(),
-      f"compile_topology={topology}",
-      f"compile_topology_num_slices={num_slice}",
-      f"model_name={model_name}",
-      "weight_dtype=float32",
-      "pure_nnx=False",
-      "enable_nnx=False",
-      "pure_nnx_decoder=False",
-  ]
-  if custom_mesh_and_rule:
-    params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
-  if overrides:
-    params.extend(overrides)
-  config = pyconfig.initialize(params)
-  validate_config(config)
-
-  topology_mesh = get_topology_mesh(config)
-  quant = quantizations.configure_quantization(config)
-  model = Transformer(config, mesh=topology_mesh, quant=quant)
-
-  learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config)
-  # tx = optax.adam(learning_rate=learning_rate_schedule)
-  tx = optimizers.get_optimizer(config, learning_rate_schedule)
-  rng = jax.random.PRNGKey(0)
-
-  init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
-
-  # Get abstract state and physical shardings from maxtext_utils
-  abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state(
-      config, topology_mesh, init_state_fn, is_training=True
-  )
-
-  # Get logical shardings from maxtext_utils
-  logical_shardings = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn)
-
-  return (
-      model_name,
-      topology,
-      num_slice,
-      custom_mesh_and_rule,
-      overrides,
-      abstract_state,
-      state_mesh_shardings,
-      logical_shardings,
-  )
-
-
-@pytest.mark.tpu_backend
-class TestGetAbstractState:
-  """Test class for get_abstract_state function and sharding comparison."""
-
-  # Requires JAX TPU support to generate the simulated TPU topology.
-  def test_get_abstract_state_sharding(self, abstract_state_and_shardings):  # pylint: disable=redefined-outer-name
-    """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding."""
-
-    (
-        model_name,
-        topology,
-        num_slice,
-        custom_mesh_and_rule,
-        overrides,
-        abstract_state,
-        state_mesh_shardings,
-        logical_shardings,
-    ) = abstract_state_and_shardings
-
-    assert hasattr(abstract_state, "params")
-    assert hasattr(abstract_state, "opt_state")
-    param_leaf = jax.tree_util.tree_leaves(abstract_state.params)[0]
-    assert isinstance(param_leaf, jax.ShapeDtypeStruct)
-    assert param_leaf.dtype == jnp.float32
-
-    root_dir = "tests/utils/sharding_info"  # Or your target directory
-    rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default"
-    if overrides:
-      rule_name += "_" + "_".join(overrides)
-    base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name)
-    os.makedirs(base_path, exist_ok=True)  # Ensure directory exists for saving actual
-
-    error_messages = []
-
-    # 1. Compare Physical/Named Shardings
-    named_json_path = os.path.join(base_path, "named_shardings.json")
-    if not os.path.exists(named_json_path):
-      pytest.skip(f"Missing named_shardings.json for {model_name} {topology} slice {num_slice}")
-      return
-
-    # Use state_mesh_shardings from the fixture
-    actual_named = named_shardings_to_json(state_mesh_shardings, abstract_state)
-    expected_named = load_json(named_json_path)
-
-    if compare_sharding_jsons(expected_named, "Expected (Physical)", actual_named, "Actual (Physical)"):
-      error_messages.append(f"Physical sharding mismatch for {model_name} on {topology} slice {num_slice}")
-
-    # 2. Compare Logical Shardings
-    logical_json_path = os.path.join(base_path, "logical_shardings.json")
-    if not os.path.exists(logical_json_path):
-      pytest.skip(f"Missing logical_shardings.json for {model_name} {topology} slice {num_slice}")
-      return
-
-    # Use logical_shardings from the fixture
-    actual_logical = partition_specs_to_json(logical_shardings, abstract_state)
-    expected_logical = load_json(logical_json_path)
-
-    if compare_sharding_jsons(expected_logical, "Expected (Logical)", actual_logical, "Actual (Logical)"):
-      error_messages.append(f"Logical sharding mismatch for {model_name} on {topology} slice {num_slice}")
-
-    assert not error_messages, "\n".join(error_messages)
+The sharding-comparison tests in this file relied on Linen golden files and
+Linen TrainState structure, which no longer exist after the NNX-only migration.
+They were removed; this module is intentionally left without tests.
+"""
diff --git a/tests/unit/state_dtypes_test.py b/tests/unit/state_dtypes_test.py
index 3d640cc62d..a92394d99c 100644
--- a/tests/unit/state_dtypes_test.py
+++ b/tests/unit/state_dtypes_test.py
@@ -14,25 +14,20 @@
 
 """Test that all weights are expected dtype (default float32)"""
 
-from functools import partial
 import unittest
 
 from flax import nnx
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
+
 from maxtext.common import train_state_nnx
-from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.configs import pyconfig
-from maxtext.layers import quantizations
-from maxtext.models import models
 from maxtext.optimizers import optimizers
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
 from tests.utils.test_helpers import get_test_config_path
 
-Transformer = models.transformer_as_linen
-
 
 class StateDtypes(unittest.TestCase):
   """Tests that state has expected dtypes, e.g. weights default to float32"""
@@ -41,43 +36,30 @@ def get_state(self, argv):
     """Gets model state including weights and optimizer state"""
     # Setup necessary inputs to build a model state
     config = pyconfig.initialize(argv)
-    quant = quantizations.configure_quantization(config)
     devices_array = maxtext_utils.create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
 
-    if config.pure_nnx:
-      _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
-    else:
-      model = Transformer(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+    _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
 
     learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config)
     tx = optimizers.get_optimizer(config, learning_rate_schedule, model)
-    _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2)
-
-    if config.pure_nnx:
 
-      def create_train_state_fn():
-        nnx_model = _create_model_partial()
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    def create_train_state_fn():
+      nnx_model = _create_model_partial()
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, example_rng)
+    init_state_fn = create_train_state_fn
     abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True)
-    return abstract_state, config.pure_nnx
+    return abstract_state
 
   def get_weights(self, argv):
-    state, is_nnx = self.get_state(argv)
-    if is_nnx:
-      return state.model
-    return state.params
+    state = self.get_state(argv)
+    return state.model
 
   def get_mu(self, argv):
-    state, is_nnx = self.get_state(argv)
-    if is_nnx:
-      return state.optimizer.opt_state[0]["mu"]
-    return state.opt_state[0].mu
+    state = self.get_state(argv)
+    return state.optimizer.opt_state[0]["mu"]
 
   def assert_pytree_is_dtype(self, weights, expected_dtype):
     """Asserts that all valid parameter arrays within the PyTree match the expected dtype."""
diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py
index 99f93c8b4d..acdd8e8224 100644
--- a/tests/unit/train_compile_test.py
+++ b/tests/unit/train_compile_test.py
@@ -828,27 +828,13 @@ def test_deepseek32(self):
 
   @parameterized.named_parameters(
       {
-          "testcase_name": "linen_scanned_dot_product",
+          "testcase_name": "scanned_dot_product",
           "scan_layers": "true",
-          "enable_nnx": "False",
           "attention": "dot_product",
       },
       {
-          "testcase_name": "linen_scanned_flash",
+          "testcase_name": "scanned_flash",
           "scan_layers": "true",
-          "enable_nnx": "False",
-          "attention": "flash",
-      },
-      {
-          "testcase_name": "nnx_scanned_dot_product",
-          "scan_layers": "true",
-          "enable_nnx": "True",
-          "attention": "dot_product",
-      },
-      {
-          "testcase_name": "nnx_scanned_flash",
-          "scan_layers": "true",
-          "enable_nnx": "True",
           "attention": "flash",
       },
   )
@@ -856,11 +842,10 @@ def test_deepseek32(self):
   def test_deepseek4(
       self,
       scan_layers,
-      enable_nnx,
       attention="dot_product",
   ):
-    # test deepseek4 compile across Linen and NNX
-    compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}_{enable_nnx}_{attention}.pickle"
+    # test deepseek4 compile.
+    compiled_trainstep_file = f"/tmp/test_deepseek4_{scan_layers}_{attention}.pickle"
     train_compile_main(
         (
             "",
@@ -885,9 +870,6 @@ def test_deepseek4(
             "sa_block_kv_dq=128",
             "dtype=bfloat16",
             "weight_dtype=bfloat16",
-            f"enable_nnx={enable_nnx}",
-            f"pure_nnx={enable_nnx}",
-            f"pure_nnx_decoder={enable_nnx}",
             "routed_bias=False",
             "override_model_config=True",
         )
@@ -1197,13 +1179,7 @@ def test_zero1_optimizer_sharding(self):
     )
 
   def test_vocab_tiling_bf16_nnx(self):
-    """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp).
-
-    Sets `pure_nnx`/`enable_nnx`/`pure_nnx_decoder` explicitly so the NNX AOT
-    path is covered regardless of the default values. Once those defaults flip
-    to True, `test_vocab_tiling_bf16` above will already exercise this same
-    path via defaults.
-    """
+    """AOT compile vocab tiling on the NNX path (vocab_tiling_nnx_loss + custom_vjp)."""
     compiled_trainstep_file = "/tmp/test_vocab_tiling_bf16_nnx.pickle"
     train_compile_main(
         (
@@ -1217,9 +1193,6 @@ def test_vocab_tiling_bf16_nnx(self):
             "max_target_length=1024",
             "num_vocab_tiling=4",
             "weight_dtype=bfloat16",
-            "pure_nnx=true",
-            "enable_nnx=true",
-            "pure_nnx_decoder=true",
         )
     )
 
@@ -1245,9 +1218,6 @@ def test_envy(self, scan_layers):
             "attention=dot_product",
             "dtype=bfloat16",
             "weight_dtype=bfloat16",
-            "enable_nnx=True",
-            "pure_nnx=True",
-            "pure_nnx_decoder=True",
             "override_model_config=True",
         )
     )

From 0279059a923498b283f51974d92307520a2cf5b2 Mon Sep 17 00:00:00 2001
From: Lance Wang <lance.wang@cienet.com>
Date: Wed, 29 Jul 2026 15:24:47 +0000
Subject: [PATCH 2/5] [NNX] Delete Linen (pre-train 2/3): collapse dispatch in
 sharding and muon utils

- sharding.maybe_update_params_sharding_with_opt delegates to the _nnx variant;
  build_zero1_input_state_mesh_shardings drops its Linen early return.
- muon_utils.get_muon_weight_dimension_numbers drops the isinstance(nnx.Module)
  test and the Linen get_abstract_param path; get_model_mdn loses its pure_nnx
  parameter and always builds the abstract NNX model.
- run_sharding_dump loses its --pure_nnx flag.

Tests follow: the Linen branch test in muon_utils_test goes away, optimizers_test
drops the pure_nnx argument and the dual-shape comparison, sharding_nnx_test
drops the flag from its fake config.
---
 src/maxtext/utils/muon_utils.py  | 51 ++++++++------------------------
 src/maxtext/utils/sharding.py    | 23 +-------------
 tests/unit/muon_utils_test.py    | 32 --------------------
 tests/unit/optimizers_test.py    |  8 ++---
 tests/unit/sharding_nnx_test.py  |  5 ++--
 tests/utils/run_sharding_dump.py | 11 ++-----
 6 files changed, 20 insertions(+), 110 deletions(-)

diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py
index ff77c57807..92762b8ac1 100644
--- a/src/maxtext/utils/muon_utils.py
+++ b/src/maxtext/utils/muon_utils.py
@@ -33,8 +33,6 @@
 import jax
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
-from maxtext.layers import quantizations
-from maxtext.models import models
 from maxtext.utils import maxtext_utils, model_creation_utils
 from optax.contrib._muon import MuonDimensionNumbers as mdn
 
@@ -134,26 +132,17 @@ def get_transform_tree(tree, path=()):
 def get_muon_weight_dimension_numbers(model, config, verbose=False):
   """Extract muon dimension number from model structure."""
 
-  if isinstance(model, nnx.Module):
-    _, abstract_param, _ = nnx.split(model, nnx.Param, ...)
+  _, abstract_param, _ = nnx.split(model, nnx.Param, ...)
 
-    def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf):
-      # Convert jax.tree_util.KeyEntry path to Tuple[str, ...]
-      path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey))
-      return transform_logic(path_strings)
+  def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf):
+    # Convert jax.tree_util.KeyEntry path to Tuple[str, ...]
+    path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey))
+    return transform_logic(path_strings)
 
-    # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves);
-    # tree_map_with_path round-trips that structure so each Param.value holds the mdn result.
-    muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(
-        apply_transform_nnx, nnx.to_pure_dict(abstract_param)
-    )
-    muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers)
-
-  else:  # Linen
-    # quickly get param structure without materialization
-    abstract_param = maxtext_utils.get_abstract_param(model, config)
-    # get muon dimension number from param
-    muon_weight_dimension_numbers = get_transform_tree(abstract_param)
+  # tree_map_with_path handles NNX's PyTree structure; result is an nnx.State with the
+  # same structure, where each Param's value holds the mdn result.
+  muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(apply_transform_nnx, nnx.to_pure_dict(abstract_param))
+  muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers)
 
   if verbose:
     _print_structure_debug(abstract_param, muon_weight_dimension_numbers)
@@ -185,7 +174,7 @@ def get_leaf_info(leaf):
   print("\nIs this reasonable?")
 
 
-def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
+def get_model_mdn(model_name, scan_layers=True, verbose=False):
   """Initializes a model and retrieves its Muon dimension numbers.
 
   This function sets up the configuration for a given model, initializes the
@@ -209,30 +198,16 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
       f"model_name={model_name}",
       f"scan_layers={scan_layers}",
       "attention=dot_product",
-      f"pure_nnx={pure_nnx}",
       "skip_jax_distributed_system=True",
   ]
-  if not pure_nnx:
-    argv.extend(
-        [
-            "enable_nnx=False",
-            "pure_nnx_decoder=False",
-        ]
-    )
   config = pyconfig.initialize(argv)
   # Setup model
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
-  quant = quantizations.configure_quantization(config)
-  if pure_nnx:
-    _, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
-  else:
-    model = models.transformer_as_linen(config, mesh=mesh, quant=quant)
+  _, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
   # Get dimension number
   muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config, verbose=verbose)
-  if pure_nnx:
-    muon_weight_dimension_numbers = {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)}
-  return muon_weight_dimension_numbers
+  return {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)}
 
 
 if __name__ == "__main__":
@@ -241,4 +216,4 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
     sys.exit(1)
   model_name_arg = sys.argv[1]
   scan_layers_arg = sys.argv[2].lower() == "true"
-  get_model_mdn(model_name_arg, scan_layers_arg, verbose=True, pure_nnx=False)
+  get_model_mdn(model_name_arg, scan_layers_arg, verbose=True)
diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py
index a0596aa7dd..6adc155ae6 100644
--- a/src/maxtext/utils/sharding.py
+++ b/src/maxtext/utils/sharding.py
@@ -698,26 +698,7 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings):
       - updated_state_mesh_shardings: State mesh shardings with updated params field
         (unchanged if shard_optimizer_over_data is False)
   """
-  if config.pure_nnx:
-    return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings)
-  prev_params_shardings = state_mesh_shardings.params
-  if config.shard_optimizer_over_data:
-    if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState):
-      sharded_fp32_params = state_mesh_shardings.opt_state.mu
-    elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance(
-        state_mesh_shardings.opt_state[0], optax.ScaleByAdamState
-    ):
-      sharded_fp32_params = state_mesh_shardings.opt_state[0].mu
-    else:
-      raise NotImplementedError(f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}")
-    if "params" not in sharded_fp32_params.keys():  # pyrefly: ignore[missing-attribute]
-      # When quantization=fp8 is enabled the sharded_fp32_params
-      # are not wrapped in `params`. Here we wrap them back.
-      sharded_fp32_params = {"params": sharded_fp32_params}
-    state_mesh_shardings = state_mesh_shardings.replace(
-        params=dict(prev_params_shardings, **sharded_fp32_params)  # pyrefly: ignore[bad-unpacking]
-    )  # pyrefly: ignore[bad-unpacking]
-  return prev_params_shardings, state_mesh_shardings
+  return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings)
 
 
 def maybe_update_params_sharding_with_opt_nnx(
@@ -851,8 +832,6 @@ def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_
   """
   if not config.shard_optimizer_over_data:
     return state_mesh_shardings
-  if not config.pure_nnx:
-    return state_mesh_shardings.replace(params=params_shardings)
   # nnx.State has no .replace: shallow-copy via tree_map (preserves nested container
   # types) and overlay params_shardings under input_state.model.
   input_state = jax.tree_util.tree_map(
diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py
index 58bfadf29a..a1f17d1e63 100644
--- a/tests/unit/muon_utils_test.py
+++ b/tests/unit/muon_utils_test.py
@@ -19,7 +19,6 @@
 import io
 import contextlib
 import unittest
-from unittest import mock
 
 import jax
 import jax.numpy as jnp
@@ -182,37 +181,6 @@ def test_nnx_verbose_path_executes_print_debug(self):
     self.assertIn("Muon Dimension Numbers", buf.getvalue())
 
 
-class TestGetMuonWeightDimensionNumbersLinen(unittest.TestCase):
-  """Covers the Linen branch of get_muon_weight_dimension_numbers."""
-
-  def test_linen_branch_uses_get_abstract_param(self):
-    """Linen models dispatch to maxtext_utils.get_abstract_param + get_transform_tree."""
-    # Build a Linen nn.Module so isinstance(model, nnx.Module) is False.
-
-    class LinenStub(nn.Module):
-
-      @nn.compact
-      def __call__(self, x):
-        return x
-
-    model = LinenStub()
-
-    # Mock the heavy get_abstract_param call with a pre-shaped dict that exercises
-    # both a standard weight path and an excluded path.
-    fake_abstract_param = {
-        "params": {
-            "self_attention": {"out": object()},
-            "norm": {"scale": object()},
-        },
-    }
-
-    with mock.patch.object(muon_utils.maxtext_utils, "get_abstract_param", return_value=fake_abstract_param):
-      result = muon_utils.get_muon_weight_dimension_numbers(model, config=mock.MagicMock())
-
-    self.assertEqual(result["params"]["self_attention"]["out"], mdn((0, -2), (-1,)))
-    self.assertIsNone(result["params"]["norm"]["scale"])
-
-
 class TestPrintStructureDebug(unittest.TestCase):
   """Covers both branches of get_leaf_info inside _print_structure_debug."""
 
diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py
index 6f43c420cd..deb4fb602f 100644
--- a/tests/unit/optimizers_test.py
+++ b/tests/unit/optimizers_test.py
@@ -399,12 +399,8 @@ def test_model_integration(self, model_name, expected_output):
     Initializes the specified MaxText model and asserts that the generated
     Muon dimension numbers match the hardcoded reference.
     """
-    is_pure_nnx = model_name in {"deepseek4-284b"}
-    actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=is_pure_nnx)
-    if "params" in expected_output and "params" in actual_output:
-      self.assertEqual(actual_output["params"], expected_output["params"])
-    else:
-      self.assertEqual(actual_output, expected_output)
+    actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True)
+    self.assertEqual(actual_output, expected_output)
 
 
 class AdamWMaskTest(parameterized.TestCase):
diff --git a/tests/unit/sharding_nnx_test.py b/tests/unit/sharding_nnx_test.py
index e51676375d..7bcbd771e9 100644
--- a/tests/unit/sharding_nnx_test.py
+++ b/tests/unit/sharding_nnx_test.py
@@ -30,7 +30,6 @@
 
 @dataclass
 class _Cfg:
-  pure_nnx: bool = True
   shard_optimizer_over_data: bool = False
 
 
@@ -83,9 +82,9 @@ class TestMaybeUpdateParamsShardingWithOptNNX(unittest.TestCase):
   def setUp(self):
     self.model = _LinearNNX(rngs=nnx.Rngs(0))
 
-  def test_dispatch_from_main_helper_when_pure_nnx(self):
+  def test_dispatch_from_main_helper(self):
     """maybe_update_params_sharding_with_opt should dispatch to the NNX variant."""
-    cfg = _Cfg(pure_nnx=True, shard_optimizer_over_data=False)
+    cfg = _Cfg(shard_optimizer_over_data=False)
     state_mesh_shardings = _build_state_mesh_shardings(self.model, optax.adam(1e-3))
     prev, updated = sharding.maybe_update_params_sharding_with_opt(cfg, state_mesh_shardings)
     # prev is the param-only view (no rngs / non-Param nodes)
diff --git a/tests/utils/run_sharding_dump.py b/tests/utils/run_sharding_dump.py
index fe371fb160..7d3156fe00 100644
--- a/tests/utils/run_sharding_dump.py
+++ b/tests/utils/run_sharding_dump.py
@@ -59,12 +59,9 @@
 flags.DEFINE_string("topology", None, "Specific topology to dump.")
 flags.DEFINE_string("num_slice", None, "Specific number of slices to dump.")
 flags.DEFINE_string("custom_mesh_and_rule", None, "Specific custom_mesh_and_rule to dump.")
-flags.DEFINE_bool("pure_nnx", False, "Use pure NNX model.")
 
 
-def run_single_dump(
-    model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple, pure_nnx: bool = False
-) -> None:
+def run_single_dump(model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple) -> None:
   """Generate sharding json file for one specific model, topology, slice and rule."""
   args = [
       "python3",
@@ -82,10 +79,6 @@ def run_single_dump(
     args.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
   if overrides:
     args.extend(overrides)
-  if pure_nnx:
-    args.append("pure_nnx=true")
-  else:
-    args.extend(["pure_nnx=False", "enable_nnx=False", "pure_nnx_decoder=False"])
   subprocess.run(args, check=True)
 
 
@@ -124,7 +117,7 @@ def main(argv: Sequence[str]) -> None:
       print("  -> Sharding files already exist. Regenerating to overwrite.")
 
     try:
-      run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides, pure_nnx=FLAGS.pure_nnx)
+      run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides)
     except subprocess.CalledProcessError:
       print(f"!!! FAILED: {model_name} {topology} {num_slice} {custom_mesh_and_rule} overrides={overrides}")
 

From 125109f7042fc8dd8dd38b80fc0ea1ba400a5656 Mon Sep 17 00:00:00 2001
From: Lance Wang <lance.wang@cienet.com>
Date: Wed, 29 Jul 2026 15:26:59 +0000
Subject: [PATCH 3/5] [NNX] Delete Linen (pre-train 3/3): collapse dispatch in
 quantization and model creation

- quantizations.maybe_quantize_model always runs the qwix forward pass with the
  dummy tokens/positions/segment ids (and the MTP decoder targets when
  mtp_num_layers > 0), then pops the transient nnx.Intermediate variables the
  traced forward sows.
- model_creation_utils.from_pretrained always builds the sharded model through
  maxtext_utils_nnx.create_nnx_sharded_model.

Tests follow: quantizations_test and nnx_quant_guard_test drop their flag
arguments and Linen expectations, correctness_tests_nnx_dispatch_test keeps only
the NNX case, and forward_pass_logit_checker always loads via from_pretrained.
---
 src/maxtext/layers/quantizations.py       |  48 +++---
 src/maxtext/utils/model_creation_utils.py |   7 +-
 tests/unit/nnx_quant_guard_test.py        |  31 +---
 tests/unit/quantizations_test.py          | 183 ++++++----------------
 tests/utils/forward_pass_logit_checker.py |  46 ++----
 5 files changed, 92 insertions(+), 223 deletions(-)

diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py
index a275a0afa8..8e34282f12 100644
--- a/src/maxtext/layers/quantizations.py
+++ b/src/maxtext/layers/quantizations.py
@@ -877,32 +877,28 @@ def maybe_quantize_model(model, config):
   if config.quantization and config.use_qwix_quantization and not config.use_batch_split_schedule:
     quantization_provider = get_qt_provider(config)
     if quantization_provider:
-      if config.pure_nnx:
-        input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
-        dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
-        dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
-        dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
-        # The MTP block reads the decoder targets, so the qwix forward pass needs them.
-        # The Linen path supplies them from the is_initializing() guard in Transformer.
-        dummy_targets = {}
-        if config.mtp_num_layers > 0:
-          dummy_targets["decoder_target_tokens"] = jnp.ones(input_shape, dtype=jnp.int32)
-          dummy_targets["decoder_target_mask"] = jnp.ones(input_shape, dtype=jnp.int32)
-        model = qwix.quantize_model(
-            model,
-            quantization_provider,
-            dummy_tokens,
-            dummy_positions,
-            dummy_segment_ids,
-            enable_dropout=False,
-            **dummy_targets,
-        )
-        # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables
-        # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches
-        # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps.
-        nnx.pop(model, nnx.Intermediate)
-      else:
-        model = qwix.quantize_model(model, quantization_provider)
+      input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
+      dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
+      dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
+      dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
+      # The MTP block reads the decoder targets, so the qwix forward pass needs them.
+      dummy_targets = {}
+      if config.mtp_num_layers > 0:
+        dummy_targets["decoder_target_tokens"] = jnp.ones(input_shape, dtype=jnp.int32)
+        dummy_targets["decoder_target_mask"] = jnp.ones(input_shape, dtype=jnp.int32)
+      model = qwix.quantize_model(
+          model,
+          quantization_provider,
+          dummy_tokens,
+          dummy_positions,
+          dummy_segment_ids,
+          enable_dropout=False,
+          **dummy_targets,
+      )
+      # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables
+      # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches
+      # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps.
+      nnx.pop(model, nnx.Intermediate)
       for _, val in nnx.graph.iter_graph(model):
         if hasattr(val, "__dict__") and "qwix_rngs" in val.__dict__:
           del val.qwix_rngs
diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py
index a2a1403125..4131dec541 100644
--- a/src/maxtext/utils/model_creation_utils.py
+++ b/src/maxtext/utils/model_creation_utils.py
@@ -928,11 +928,8 @@ def from_pretrained(
   _, _abs_state_for_specs = nnx.split(abstract_model)
   specs = nnx.get_partition_spec(_abs_state_for_specs)
 
-  if config.pure_nnx:
-    model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh)
-    # TODO: print debug_sharding info
-  else:
-    model = create_nnx_sharded_model_hybrid(config, mesh, devices, model_mode, rng_key)
+  model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh)
+  # TODO: print debug_sharding info
 
   sharded_state = nnx.state(model)
 
diff --git a/tests/unit/nnx_quant_guard_test.py b/tests/unit/nnx_quant_guard_test.py
index 50cac5d349..8451dd1264 100644
--- a/tests/unit/nnx_quant_guard_test.py
+++ b/tests/unit/nnx_quant_guard_test.py
@@ -12,43 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""qwix + NNX coverage: the config guard and the ToNNX->Linen bridge.
+"""qwix + NNX coverage for the ToNNX->Linen bridge.
 
-- Config guard: qwix quantization under pure_nnx requires the pure NNX decoder.
-  The bridged Linen decoder (pure_nnx_decoder=False) is invisible to qwix, so
-  quantization/sparsity would silently no-op; validation must reject that combo.
-- Bridge: nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs
-  (qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising.
+nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs
+(qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising.
 """
 
-import sys
 import unittest
 
 import jax.numpy as jnp
 from flax import nnx
 
-from maxtext.configs import pyconfig
 from maxtext.layers import nnx_wrappers
-from tests.utils.test_helpers import get_test_config_path
-
-
-class QwixNnxQuantGuardTest(unittest.TestCase):
-
-  def _init(self, **overrides):
-    overrides.setdefault("enable_checkpointing", False)
-    return pyconfig.initialize([sys.argv[0], get_test_config_path()], **overrides)
-
-  def test_bridged_decoder_with_qwix_quant_raises(self):
-    with self.assertRaisesRegex(Exception, "pure_nnx_decoder"):
-      self._init(pure_nnx=True, pure_nnx_decoder=False, use_qwix_quantization=True, quantization="fp8_full")
-
-  def test_pure_nnx_decoder_with_qwix_quant_ok(self):
-    cfg = self._init(pure_nnx=True, pure_nnx_decoder=True, use_qwix_quantization=True, quantization="fp8_full")
-    self.assertTrue(cfg.pure_nnx_decoder)
-
-  def test_bridged_decoder_without_quant_ok(self):
-    cfg = self._init(pure_nnx=True, pure_nnx_decoder=False, quantization="")
-    self.assertEqual(cfg.quantization, "")
 
 
 class NnxAttrsToLinenVarsBridgeTest(unittest.TestCase):
diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py
index 1c34d128be..5ce2a3444d 100644
--- a/tests/unit/quantizations_test.py
+++ b/tests/unit/quantizations_test.py
@@ -395,149 +395,71 @@ def quantization_config(self, quant, logits_tolerance=2e-1, grad_tolerance=5e-1,
     cfg = self.init_pyconfig(quantization=quant, **kwargs)
     ids, decoder_segment_ids, decoder_positions = self.get_data()
 
-    if cfg.pure_nnx:
-      qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0))
-      if getattr(self.__class__, "_cached_base_results_nnx", None) is None:
-        base_cfg = self.init_pyconfig(quantization="", **kwargs)
-        base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0))
-
-        def loss_base(model):
-          logits = model(
-              decoder_input_tokens=ids,
-              decoder_positions=decoder_positions,
-              decoder_segment_ids=decoder_segment_ids,
-              enable_dropout=False,
-          )
-          return jnp.mean((logits) ** 2)
-
-        grads_base = nnx.grad(loss_base)(base_model)
-        logits_base = base_model(
-            decoder_input_tokens=ids,
-            decoder_positions=decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-        )
-        self.__class__._cached_base_results_nnx = (grads_base, logits_base)
-
-      grads_base, logits = self.__class__._cached_base_results_nnx
+    qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0))
+    if getattr(self.__class__, "_cached_base_results_nnx", None) is None:
+      base_cfg = self.init_pyconfig(quantization="", **kwargs)
+      base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0))
 
-      def loss_quant(model):
-        logits_q = model(
+      def loss_base(model):
+        logits = model(
             decoder_input_tokens=ids,
             decoder_positions=decoder_positions,
             decoder_segment_ids=decoder_segment_ids,
             enable_dropout=False,
         )
-        return jnp.mean((logits_q) ** 2)
+        return jnp.mean((logits) ** 2)
 
-      grads_quant = nnx.grad(loss_quant)(qt_model)
-      quant_logits = qt_model(
+      grads_base = nnx.grad(loss_base)(base_model)
+      logits_base = base_model(
           decoder_input_tokens=ids,
           decoder_positions=decoder_positions,
           decoder_segment_ids=decoder_segment_ids,
           enable_dropout=False,
       )
+      self.__class__._cached_base_results_nnx = (grads_base, logits_base)
 
-      print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
-      assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
-
-      # nnx.grad returns a State object which is a mapping of paths to gradients.
-      # Flatten them to check for tolerance.
-      grads_base_flat = traversals.flatten_mapping(grads_base)
-      grads_quant_flat = traversals.flatten_mapping(grads_quant)
-
-      # Filter for param collections to compare only parameters and not stats/buffers if any
-      # Note: NNX grads structure might contain variables like 'kernel', 'bias'.
-      # For simplicity we compare all matching keys.
-      def flatten_and_filter(grads_flat):
-        return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)}
-
-      gb_f = flatten_and_filter(grads_base_flat)
-      gq_f = flatten_and_filter(grads_quant_flat)
-
-      for k in gb_f:
-        if k in gq_f:
-          diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8)
-          if diff > grad_tolerance:
-            print(f"Gradient mismatch for {k}: rel_error = {diff}")
-            assert diff <= grad_tolerance
-    else:
-      qt_model = model_creation_utils.create_model(cfg, self.mesh)
-      if not hasattr(self.__class__, "_cached_base_results"):
-        model = model_creation_utils.create_model(self.cfg, self.mesh)
-        var = model.init(
-            {"params": self.rng, "aqt": self.rng, "dropout": self.rng},
-            ids,
-            decoder_positions,
-            decoder_segment_ids,
-            enable_dropout=False,
-            mutable=True,
-        )
-
-        def loss_base_linen(all_vars, inputs):
-          logits_b, _ = model.apply(
-              all_vars,
-              *inputs,
-              enable_dropout=False,
-              rngs={"params": self.rng},
-              mutable=True,
-          )
-          return jnp.mean((logits_b) ** 2)
-
-        grads_base_linen = jax.grad(loss_base_linen)(var, (ids, decoder_positions, decoder_segment_ids))
-        logits_b, _ = model.apply(
-            var,
-            ids,
-            decoder_positions,
-            decoder_segment_ids,
-            enable_dropout=False,
-            rngs={"params": self.rng},
-            mutable=True,
-        )
-        self.__class__._cached_base_results = (grads_base_linen, logits_b)
+    grads_base, logits = self.__class__._cached_base_results_nnx
 
-      grads_base_linen, logits = self.__class__._cached_base_results
-
-      quantized_vars = qt_model.init(
-          {"params": self.rng, "aqt": self.rng, "dropout": self.rng},
-          ids,
-          decoder_positions,
-          decoder_segment_ids,
+    def loss_quant(model):
+      logits_q = model(
+          decoder_input_tokens=ids,
+          decoder_positions=decoder_positions,
+          decoder_segment_ids=decoder_segment_ids,
           enable_dropout=False,
-          mutable=True,
       )
+      return jnp.mean((logits_q) ** 2)
+
+    grads_quant = nnx.grad(loss_quant)(qt_model)
+    quant_logits = qt_model(
+        decoder_input_tokens=ids,
+        decoder_positions=decoder_positions,
+        decoder_segment_ids=decoder_segment_ids,
+        enable_dropout=False,
+    )
 
-      def loss_quant_linen(all_vars, inputs):
-        logits_q, _ = qt_model.apply(
-            all_vars,
-            *inputs,
-            enable_dropout=False,
-            rngs={"params": self.rng},
-            mutable=True,
-        )
-        return jnp.mean((logits_q) ** 2)
+    print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
+    assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
 
-      grads_quant_linen = jax.grad(loss_quant_linen)(quantized_vars, (ids, decoder_positions, decoder_segment_ids))
+    # nnx.grad returns a State object which is a mapping of paths to gradients.
+    # Flatten them to check for tolerance.
+    grads_base_flat = traversals.flatten_mapping(grads_base)
+    grads_quant_flat = traversals.flatten_mapping(grads_quant)
 
-      quant_logits, _ = qt_model.apply(
-          quantized_vars,
-          ids,
-          decoder_positions,
-          decoder_segment_ids,
-          enable_dropout=False,
-          rngs={"params": self.rng},
-          mutable=True,
-      )
-      print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
-      assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
-      self.print_grad_diff(grads_base_linen["params"], grads_quant_linen["params"])
-      self.assertTrue(
-          self.pytree_allclose(
-              grads_base_linen["params"],
-              grads_quant_linen["params"],
-              tolerance=grad_tolerance,
-          )
-      )
+    # Filter for param collections to compare only parameters and not stats/buffers if any
+    # Note: NNX grads structure might contain variables like 'kernel', 'bias'.
+    # For simplicity we compare all matching keys.
+    def flatten_and_filter(grads_flat):
+      return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)}
+
+    gb_f = flatten_and_filter(grads_base_flat)
+    gq_f = flatten_and_filter(grads_quant_flat)
+
+    for k in gb_f:
+      if k in gq_f:
+        diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8)
+        if diff > grad_tolerance:
+          print(f"Gradient mismatch for {k}: rel_error = {diff}")
+          assert diff <= grad_tolerance
 
   @pytest.mark.tpu_only
   def test_int8_quantization(self):
@@ -545,7 +467,7 @@ def test_int8_quantization(self):
 
   @pytest.mark.tpu_only
   def test_int8_quantization_nnx(self):
-    self.quantization_config("int8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("int8")
 
   @pytest.mark.tpu_only
   def test_fp8_quantization(self):
@@ -553,7 +475,7 @@ def test_fp8_quantization(self):
 
   @pytest.mark.tpu_only
   def test_fp8_quantization_nnx(self):
-    self.quantization_config("fp8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8")
 
   @pytest.mark.tpu_only
   def test_fp8_full_quantization(self):
@@ -561,7 +483,7 @@ def test_fp8_full_quantization(self):
 
   @pytest.mark.tpu_only
   def test_fp8_full_quantization_nnx(self):
-    self.quantization_config("fp8_full", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_full")
 
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
@@ -571,7 +493,7 @@ def test_fp8_gpu_quantization(self):
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_gpu_quantization_nnx(self):
-    self.quantization_config("fp8_gpu", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_gpu", grad_tolerance=1.5)
 
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
@@ -581,7 +503,7 @@ def test_fp8_nanoo_quantization(self):
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_nanoo_quantization_nnx(self):
-    self.quantization_config("fp8_nanoo", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_nanoo", grad_tolerance=1.5)
 
   @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
   @pytest.mark.gpu_only
@@ -662,8 +584,6 @@ def test_maybe_quantize_model_pops_intermediates(self):
         quantization="int8",
         use_qwix_quantization=True,
         use_batch_split_schedule=False,
-        pure_nnx=True,
-        pure_nnx_decoder=True,
         micro_batch_size_to_train_on=1,
         max_target_length=2,
     )
@@ -700,7 +620,6 @@ def test_nnx_abstract_state_has_no_intermediates(self):
         enable_checkpointing=False,
         model_name="deepseek3-tiny",
         attention="dot_product",
-        pure_nnx=True,
         use_qwix_quantization=True,
         use_qk_clip=True,  # This sows QK clip intermediates during the forward pass
     )
diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py
index a51b23980f..66d6fa4343 100644
--- a/tests/utils/forward_pass_logit_checker.py
+++ b/tests/utils/forward_pass_logit_checker.py
@@ -71,7 +71,6 @@
 """
 
 import argparse
-import functools
 import os
 from pathlib import Path
 import sys
@@ -83,8 +82,6 @@
 from maxtext.utils.globals import MAXTEXT_TEST_ASSETS_ROOT, HF_IDS
 from maxtext.checkpoint_conversion.utils.hf_utils import convert_jax_weight_to_torch
 from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN
-from maxtext.layers import quantizations
-from maxtext.models import models
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
@@ -356,7 +353,7 @@ def get_data(golden_data_point, config):
 def main(config, test_args):  # pylint: disable=W0621
   """Test the Whole Model of model_name"""
   init_rng = jax.random.PRNGKey(config.init_weights_seed)
-  init_rng, rng1 = jax.random.split(init_rng)
+  init_rng, _ = jax.random.split(init_rng)
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
 
@@ -393,19 +390,13 @@ def main(config, test_args):  # pylint: disable=W0621
   if not test_args.run_hf_model:
     """Comparing maxtext/huggingface model with pre-loaded golden logitis"""
     max_logging.log("Initializing MaxText model")
-    quant = quantizations.configure_quantization(config)
-    if config.pure_nnx_decoder and config.enable_nnx:
-      model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)
-
-      if config.lora.enable_lora:
-        model = lora_utils.apply_lora_to_model(model, mesh, config)
-        if config.lora.lora_restore_path:
-          lora_utils.restore_lora_from_path(model, config)
-      state = None
-    else:
-      model = models.transformer_as_linen(config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, config, False, rng1)
-      state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
+    model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)
+
+    if config.lora.enable_lora:
+      model = lora_utils.apply_lora_to_model(model, mesh, config)
+      if config.lora.lora_restore_path:
+        lora_utils.restore_lora_from_path(model, config)
+    state = None
 
     if test_args.golden_logits_path == "":
       input_golden_data_path = os.path.join(
@@ -641,22 +632,13 @@ def main(config, test_args):  # pylint: disable=W0621
         raise ImportError("peft library is required to load HF LoRA adapter. Run `pip install peft`.") from exc
       hf_model = PeftModel.from_pretrained(hf_model, hf_lora_path)
 
-    quant = quantizations.configure_quantization(config)
-    if config.pure_nnx_decoder and config.enable_nnx:
-      maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)
+    maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN)
 
-      if config.lora.enable_lora:
-        maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config)
-        if config.lora.lora_restore_path:
-          lora_utils.restore_lora_from_path(maxtext_model, config)
-      maxtext_state = None
-    else:
-      maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, rng1)
-      if test_args.ckpt_type == "linen":
-        maxtext_state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
-      else:
-        maxtext_state, _ = model_creation_utils.setup_decode_state_from_nnx(maxtext_model, config, rng1, mesh)
+    if config.lora.enable_lora:
+      maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config)
+      if config.lora.lora_restore_path:
+        lora_utils.restore_lora_from_path(maxtext_model, config)
+    maxtext_state = None
 
     # The long prompt is required to catch position-dependent regressions (e.g. RoPE);
     # the short prompts above cannot detect them. See build_long_prompt().

From 34e33cdc0d5f6ab96d192bd5dee27e4674ac012d Mon Sep 17 00:00:00 2001
From: Lance Wang <lance.wang@cienet.com>
Date: Wed, 19 Aug 2026 15:36:52 +0000
Subject: [PATCH 4/5] [NNX] Delete Linen (2/5): collapse dispatch in trainers
 (pre-train, DiLoCo, GRPO)

---
 src/maxtext/experimental/rl/grpo_trainer.py   | 309 ++------
 src/maxtext/trainers/pre_train/train.py       | 695 ++++++------------
 .../trainers/pre_train/train_compile.py       |  90 +--
 .../golden_logits/golden_dpo_correctness.json |  18 +-
 .../generate_grpo_golden_logits.py            |  93 +--
 tests/integration/diloco_test.py              | 163 ++--
 .../integration/dpo_correctness_base.py       |   3 -
 .../integration/grpo_correctness.py           |  91 +--
 .../grpo_trainer_correctness_test.py          |  54 +-
 .../sft_trainer_correctness_test.py           |  41 +-
 tests/unit/grpo_nnx_test.py                   |  72 +-
 tests/unit/pre_train_loss_mask_test.py        |  61 +-
 12 files changed, 462 insertions(+), 1228 deletions(-)

diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py
index d8661e9814..6d68133822 100644
--- a/src/maxtext/experimental/rl/grpo_trainer.py
+++ b/src/maxtext/experimental/rl/grpo_trainer.py
@@ -38,7 +38,6 @@
 import datetime
 import time
 import os
-import functools
 import threading
 
 from typing import Sequence, Callable, Iterator
@@ -109,43 +108,6 @@
 # -----------------------------------------------------------------------------
 
 
-def _split_grpo_state(state):
-  """Splits the reference parameters from the main training state.
-
-  This is a utility function to separate the reference model's parameters,
-  which are kept frozen, from the policy model's parameters that are actively
-  being trained.
-
-  Args:
-    state: The combined training state, expected to contain a 'reference_params'
-      key within its `params` attribute.
-
-  Returns:
-    A tuple containing:
-      - new_state: The training state with 'reference_params' removed.
-      - reference_params: The extracted reference parameters.
-  """
-  reference_params = state.params["reference_params"]
-  new_state = state.replace(params={k: v for k, v in state.params.items() if k != "reference_params"})
-  return new_state, reference_params
-
-
-def _merge_grpo_state(state, reference_params):
-  """Merges the reference parameters back into the training state.
-
-  This is the inverse operation of `_split_grpo_state`, used to reconstruct
-  the full state object after a training step.
-
-  Args:
-    state: The training state, without 'reference_params'.
-    reference_params: The frozen reference parameters to be added back.
-
-  Returns:
-    A new state object with the 'reference_params' re-integrated.
-  """
-  return state.replace(params=dict(state.params, reference_params=reference_params))
-
-
 @struct.dataclass
 class LossAux:
   """A dataclass to hold auxiliary outputs from the GRPO loss function.
@@ -609,120 +571,22 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
   """Run one GRPO training step.
 
   Computes the GRPO loss and gradients and applies the update to the policy
-  parameters; the reference parameters are held constant. The Linen and NNX
-  paths share this entry point: on the NNX path `model` is an NNX
-  `GraphDef` and `state` is the matching flat `nnx.State` of a
-  `TrainStateNNX`. On the Linen path they are the usual `nn.Module` and
-  `TrainState`.
+  parameters; the reference parameters are held constant. `model` is an NNX
+  `GraphDef` and `state` is the matching flat `nnx.State` of a `TrainStateNNX`.
 
   Args:
-    model: Linen `nn.Module` or NNX `GraphDef`, depending on `config.pure_nnx`.
+    model: NNX `GraphDef` of the `TrainStateNNX`.
     config: Training configuration object.
     state_mesh_shardings: Pytree of shardings matching `state`.
-    params_shardings: Param-only shardings, used for gradient accumulation
-      on the Linen path. Ignored on NNX.
-    state: Linen `TrainState` or NNX `nnx.State` matching `model`.
+    params_shardings: Param-only shardings.
+    state: NNX `nnx.State` matching `model`.
     data: A batch dict produced by the GRPO input pipeline.
-    dropout_rng: PRNG key for dropout (Linen only).
+    dropout_rng: PRNG key for dropout.
 
   Returns:
     A tuple `(new_state, metrics)`.
   """
-  if config.pure_nnx:
-    return _train_step_nnx(model, config, state_mesh_shardings, state, data)
-
-  state, reference_params = _split_grpo_state(state)
-  state_mesh_shardings, reference_params_sharding = _split_grpo_state(state_mesh_shardings)
-  extra_grpo_args = [reference_params]
-  _loss_fn = grpo_loss_fn
-
-  if config.gradient_accumulation_steps > 1:
-
-    def accumulate_gradient(acc_grad_and_loss, data):
-      grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
-      (_, aux), cur_batch_gradient = grad_func(
-          model, config, data, dropout_rng, state.params, *extra_grpo_args, is_train=True
-      )
-      acc_grad_and_loss["loss"] += aux["total_loss"]
-      acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"]
-      acc_grad_and_loss["grad"] = jax.tree_util.tree_map(
-          lambda x, y: x * aux["total_weights"] + y, cur_batch_gradient, acc_grad_and_loss["grad"]
-      )
-      acc_grad_and_loss["total_weights"] += aux["total_weights"]
-      return acc_grad_and_loss, aux
-
-    def reshape_to_microbatch_accumulations(batch_arr):
-      """Reshape global batch to microbatches, assuming batch axis is leading."""
-      microbatches = config.gradient_accumulation_steps
-      microbatch_shape = (microbatches, batch_arr.shape[0] // microbatches) + batch_arr.shape[1:]
-      return jnp.reshape(batch_arr, microbatch_shape)
-
-    data = jax.tree_util.tree_map(reshape_to_microbatch_accumulations, data)
-    init_grad = jax.tree_util.tree_map(jnp.zeros_like, state.params)
-    init_grad_and_loss = {"loss": 0.0, "grad": init_grad, "total_weights": 0, "moe_lb_loss": 0.0}
-
-    grad_and_loss, aux = jax.lax.scan(
-        accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps
-    )
-    loss = (
-        grad_and_loss["loss"] / grad_and_loss["total_weights"]
-        + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps
-    )
-    raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], grad_and_loss["grad"])
-    aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux)
-  else:
-    if config.optimizer_memory_host_offload:
-      cast_params = jax.device_put(state.params, max_utils.with_memory_kind(state_mesh_shardings.params, "device"))
-      cast_params = max_utils.cast_to_bf16(cast_params)
-      state = state.replace(params=cast_params)
-      if config.use_grpo:
-        reference_params = jax.device_put(
-            reference_params, max_utils.with_memory_kind(reference_params_sharding, "device")
-        )
-        reference_params = max_utils.cast_to_bf16(reference_params)
-        extra_grpo_args = [reference_params]
-    grad_func = jax.value_and_grad(_loss_fn, argnums=4, has_aux=True)
-    (loss, aux), raw_grads = grad_func(model, config, data, dropout_rng, state.params, *extra_grpo_args, is_train=True)
-
-  total_weights = aux.total_weights
-  moe_lb_loss = aux.moe_lb_loss
-
-  if config.gradient_clipping_threshold > 0:
-    grads = maxtext_utils.apply_gradient_clipping(raw_grads, state, config.gradient_clipping_threshold)
-  else:
-    grads = raw_grads
-  if config.optimizer_memory_host_offload:
-    state = state.replace(
-        opt_state=jax.device_put(
-            state.opt_state,
-            jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="device"), state_mesh_shardings.opt_state),
-        )
-    )
-  new_state = state.apply_gradients(grads=grads)
-
-  scalar_metrics = {
-      "learning/loss": loss,
-      "learning/avg_reward": aux.avg_reward,
-      "learning/avg_reward_std": aux.avg_reward_std,
-      "learning/avg_advantage": aux.avg_advantage,
-      "learning/avg_kl": aux.avg_kl,
-      "learning/completion_length": aux.completion_length,
-      "learning/moe_lb_loss": moe_lb_loss,
-      "learning/total_weights": total_weights,
-  }
-  if not config.optimizer_memory_host_offload:
-    scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads)
-    scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads)
-    scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(new_state.params)
-    scalar_metrics["learning/avg_reward"] = aux.avg_reward
-  metrics = {
-      "scalar": scalar_metrics,
-      "scalars": {},
-  }
-
-  new_state = _merge_grpo_state(new_state, reference_params)
-
-  return new_state, metrics
+  return _train_step_nnx(model, config, state_mesh_shardings, state, data)
 
 
 def eval_step(model, config, state, data, dropout_rng):
@@ -741,31 +605,7 @@ def eval_step(model, config, state, data, dropout_rng):
   Returns:
     A dictionary of evaluation metrics.
   """
-  if config.pure_nnx:
-    return _eval_step_nnx(model, config, state, data)
-
-  reference_params, extra_grpo_args, _loss_fn = [], [], grpo_loss_fn
-  state, reference_params = _split_grpo_state(state)
-  extra_grpo_args = [reference_params]
-  _loss_fn = grpo_loss_fn
-
-  eval_loss_fn = functools.partial(_loss_fn, model, config, data, dropout_rng, is_train=False)
-  loss, aux = eval_loss_fn(state.params, *extra_grpo_args)
-  total_loss = aux["total_loss"]
-  total_weights = aux["total_weights"]
-  moe_lb_loss = aux["moe_lb_loss"]
-  metrics = {
-      "scalar": {
-          "evaluation/loss": loss,
-          "evaluation/total_loss": total_loss,
-          "evaluation/total_weights": total_weights,
-          "evaluation/moe_lb_loss": moe_lb_loss,
-      },
-  }
-  if config.use_dpo:
-    metrics["scalar"]["evaluation/grpo_reward_accuracy"] = aux["reward_accuracy"]
-
-  return metrics
+  return _eval_step_nnx(model, config, state, data)
 
 
 def setup_train_loop(
@@ -812,54 +652,41 @@ def setup_train_loop(
       - eval_data_iterator: The iterator for the evaluation dataset (or None).
       - state: The initialized training state.
   """
-  if config.pure_nnx != config_inference.pure_nnx:
-    raise ValueError(
-        f"config.pure_nnx ({config.pure_nnx}) and config_inference.pure_nnx " f"({config_inference.pure_nnx}) must agree."
-    )
   with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT):
     max_logging.log("Training mesh used for the workload")
     num_inference_devices = config.inference_devices_per_replica * config.inference_replicas
     training_devices = jax.devices()[num_inference_devices:]
     init_rng = jax.random.PRNGKey(config.init_weights_seed)
 
-    if config.pure_nnx:
-      training_mesh = maxtext_utils.get_mesh_from_config(config, devices=training_devices)
-      training_rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng)
-      model = mt.from_config(config, devices=training_devices, mesh=training_mesh, rngs=training_rngs)
-    else:
-      model = mt.from_config(config, devices=training_devices)
+    training_mesh = maxtext_utils.get_mesh_from_config(config, devices=training_devices)
+    training_rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=init_rng)
+    model = mt.from_config(config, devices=training_devices, mesh=training_mesh, rngs=training_rngs)
     mesh = model.mesh
 
     max_logging.log("Inference mesh used for the workload")
     inference_devices = jax.devices()[:num_inference_devices]
-    if config_inference.pure_nnx:
-      inference_mesh_obj = maxtext_utils.get_mesh_from_config(config_inference, devices=inference_devices)
-      inference_rngs = maxtext_utils_nnx.create_nnx_rngs(config_inference, rng_key=init_rng)
-      inference_model = mt.from_config(
-          config_inference, devices=inference_devices, mesh=inference_mesh_obj, rngs=inference_rngs
-      )
-    else:
-      inference_model = mt.from_config(config_inference, devices=inference_devices)
+    inference_mesh_obj = maxtext_utils.get_mesh_from_config(config_inference, devices=inference_devices)
+    inference_rngs = maxtext_utils_nnx.create_nnx_rngs(config_inference, rng_key=init_rng)
+    inference_model = mt.from_config(
+        config_inference, devices=inference_devices, mesh=inference_mesh_obj, rngs=inference_rngs
+    )
     inference_mesh = inference_model.mesh
 
     learning_rate_schedule, tx = train_utils.create_training_optimizer(config, model)
 
-    if config.pure_nnx:
-      _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh, devices=training_devices)
-
-      def init_state_fn():
-        nnx_model = _create_model_partial()
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-        # Reference uses the same init seed so it starts identical to the policy.
-        reference_model = _create_model_partial()
-        # TrainStateNNX only takes (model, optimizer); reference_model is an NNX
-        # sibling attribute set after construction (nnx.Module is mutable).
-        state = train_state_nnx.TrainStateNNX(nnx_model, optimizer)
-        state.reference_model = reference_model
-        return state
-
-    else:
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng)
+    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh, devices=training_devices)
+
+    def init_state_fn():
+      nnx_model = _create_model_partial()
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+      # Reference uses the same init seed so it starts identical to the policy.
+      reference_model = _create_model_partial()
+      # TrainStateNNX only takes (model, optimizer); reference_model is an NNX
+      # sibling attribute set after construction (nnx.Module is mutable).
+      state = train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+      state.reference_model = reference_model
+      return state
+
     checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn)
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
@@ -868,29 +695,21 @@ def init_state_fn():
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
 
-  if config_inference.pure_nnx:
-    _create_inference_partial, _ = model_creation_utils.create_nnx_abstract_model(
-        config_inference, inference_mesh, devices=inference_devices
-    )
+  _create_inference_partial, _ = model_creation_utils.create_nnx_abstract_model(
+      config_inference, inference_mesh, devices=inference_devices
+  )
 
-    def init_inference_state_fn():
-      inference_nnx_model = _create_inference_partial()
-      return train_state_nnx.TrainStateNNX(inference_nnx_model, None)
+  def init_inference_state_fn():
+    inference_nnx_model = _create_inference_partial()
+    return train_state_nnx.TrainStateNNX(inference_nnx_model, None)
 
-  else:
-    init_inference_state_fn = functools.partial(
-        maxtext_utils.init_initial_state, inference_model, tx, config_inference, False, init_rng
-    )
   inference_state_mesh_shardings = maxtext_utils.get_abstract_state(
       config_inference, inference_mesh, init_inference_state_fn, is_training=False
   )[2]
   if not config.using_pipeline_parallelism:
     # The vocab tensor(s) of shape [vocab, embed] (and transpose) are not sharded by stage
-    if config.pure_nnx:
-      params_for_check = nnx.state(state.model, nnx.Param)
-      sharding.assert_params_sufficiently_sharded(params_for_check, mesh, config.sharding_tolerance)
-    else:
-      sharding.assert_params_sufficiently_sharded(state.params, mesh, config.sharding_tolerance)
+    params_for_check = nnx.state(state.model, nnx.Param)
+    sharding.assert_params_sufficiently_sharded(params_for_check, mesh, config.sharding_tolerance)
 
   return (
       init_rng,
@@ -1003,16 +822,10 @@ def train_loop(config, config_inference, recorder, state=None):
       token=config.hf_access_token,
   )
 
-  if config.pure_nnx:
-    # `reference_model` is a sibling field on TrainStateNNX, populated by
-    # init_state_fn. Nothing to merge here; just verify it is present.
-    if not hasattr(state, "reference_model"):
-      raise RuntimeError("NNX GRPO state is missing reference_model; check setup_train_loop.")
-  else:
-    if "reference_params" not in state.params:
-      reference_params = jax.tree.map(jnp.copy, state.params["params"])
-      state = _merge_grpo_state(state, reference_params)
-    state_mesh_shardings = _merge_grpo_state(state_mesh_shardings, state_mesh_shardings.params["params"])
+  # `reference_model` is a sibling field on TrainStateNNX, populated by
+  # init_state_fn. Nothing to merge here; just verify it is present.
+  if not hasattr(state, "reference_model"):
+    raise RuntimeError("NNX GRPO state is missing reference_model; check setup_train_loop.")
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config, model, mesh, state, state_mesh_shardings, train_step, eval_step, eval_data_iterator
@@ -1037,11 +850,8 @@ def train_loop(config, config_inference, recorder, state=None):
   metric_logger = MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule)
 
   # Write train config params, num model params, and XLA flags to tensorboard
-  if config.pure_nnx:
-    params_for_metrics = nnx.state(state.model, nnx.Param)
-    metric_logger.write_setup_info_to_tensorboard(params_for_metrics)
-  else:
-    metric_logger.write_setup_info_to_tensorboard(state.params["params"])
+  params_for_metrics = nnx.state(state.model, nnx.Param)
+  metric_logger.write_setup_info_to_tensorboard(params_for_metrics)
 
   def generation_worker_fn(
       worker_inference_engine,
@@ -1165,32 +975,21 @@ def generation_worker_fn(
           state, metrics = p_train_step(state, example_batch, train_rng)
       with jax.profiler.StepTraceAnnotation("transfer data", step_num=step):
         if step != 0 and step % config.inference_rollouts == 0:
-          if config.pure_nnx:
-            grpo_utils.pathways_reshard_nnx(
-                config_inference,
-                inference_engine,
-                state.model,
-                state_mesh_shardings.model,
-                inference_state_mesh_shardings.model,
-            )
-          else:
-            grpo_utils.pathways_reshard(
-                config_inference,
-                inference_engine,
-                {"params": state.params["params"]},
-                {"params": state_mesh_shardings.params["params"]},
-                mesh,
-                {"params": inference_state_mesh_shardings.params["params"]},
-            )
+          grpo_utils.pathways_reshard_nnx(
+              config_inference,
+              inference_engine,
+              state.model,
+              state_mesh_shardings.model,
+              inference_state_mesh_shardings.model,
+          )
           with data_buffer_lock:
             data_buffer.clear()
 
       step_time_delta = datetime.datetime.now() - last_step_completion
 
-      # On the Linen path, the reference is embedded in `state.params` and is
-      # stripped before saving. On the NNX path, the reference is a sibling
-      # field on TrainStateNNX, so the whole state can be saved as-is.
-      state_to_save = state if config.pure_nnx else _split_grpo_state(state)[0]
+      # The reference is a sibling field on TrainStateNNX, so the whole state
+      # can be saved as-is.
+      state_to_save = state
       checkpointing.maybe_save_checkpoint(checkpoint_manager, state_to_save, config, data_iterator, step)
 
       if config.dump_hlo and step == start_step:
@@ -1234,7 +1033,7 @@ def generation_worker_fn(
       metric_logger.buffer_and_write_metrics(metrics, step, step_time_delta)
 
       if config.save_checkpoint_on_completion:
-        state_to_save = state if config.pure_nnx else _split_grpo_state(state)[0]
+        state_to_save = state
         checkpointing.maybe_save_checkpoint(checkpoint_manager, state_to_save, config, data_iterator)
       elif checkpoint_manager is not None:
         # in case the last checkpoint_period checkpoint is still in progress
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 274d9acc0a..04d0c42c55 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -25,8 +25,6 @@
 
 from absl import app
 
-import optax
-
 import pathwaysutils  # pylint: disable=unused-import
 
 try:
@@ -40,7 +38,7 @@
 import jax.numpy as jnp
 from jax.sharding import NamedSharding
 
-from flax import linen as nn, nnx, traverse_util
+from flax import nnx, traverse_util
 from flax.linen import partitioning as nn_partitioning
 from flax.nnx import variablelib
 
@@ -77,14 +75,13 @@
 from maxtext.utils import maxtext_utils_nnx
 from maxtext.utils import train_utils
 from maxtext.utils.gradient_accumulation import gradient_accumulation_loss_and_grad
-from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss
+from maxtext.utils.vocabulary_tiling import vocab_tiling_nnx_loss
 
 VertexTensorboardManager, _vertex_tb_is_stub = vertex_tensorboard_modules()
 
 
 def get_first_step(model, state):
-  if isinstance(model, nn.Module):
-    return int(state.step)
+  del model  # NNX-only; kept for call-site signature parity
   if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo)
     step_val = state.step.get_value() if hasattr(state.step, "get_value") else state.step
     return int(step_val)
@@ -100,17 +97,18 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
   """loss_fn for both train and eval.
 
   Args:
-    model: A nn.Module (Linen) or nnx.Module (NNX).
+    model: An NNX model.
     config: Config of parameters
     data: Batch of data to apply to the model
-    dropout_rng: A key to use to generate rng for dropout (Linen); unused for NNX.
-    params: Model params (Linen); unused for NNX (params are part of the model).
+    dropout_rng: Unused for NNX (kept for signature parity).
+    params: Unused for NNX; params are part of the model.
     is_train: True for train_step and False for eval_step
 
   Returns:
     loss: average loss
     aux: a dictionary including intermediate_outputs, xent_sum, and total_weights
   """
+  del dropout_rng, params, sparsity_state  # unused for NNX (kept for signature parity)
   is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion"
   if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion:
     raise ValueError(
@@ -144,175 +142,86 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
   # parameters in the model and checkpoints. Only pass image inputs when present.
   encoder_images = data.get("images") if config.use_multimodal else None
   encoder_image_masks = data.get("image_masks") if config.use_multimodal else None
-  mutable_collections = ["intermediates"]
-  if config.mtp_num_layers > 0 and is_train:
-    # The single model.apply call now triggers the entire chain if MTP is enabled:
-    # Decoder runs -> returns hidden_state -> MTPBlock uses it -> MTPBlock sows losses -> we reap them here.
-    mutable_collections.append("mtp_losses")
-
-  # During evaluation, if the acceptance rate test is enabled, we must
-  # make its specific collection mutable so the MTPBlock can sow into it.
-  if config.mtp_eval_target_module > 0 and not is_train:
-    mutable_collections.append("mtp_acceptance")
-  if config.use_indexer and is_train:
-    mutable_collections.append("indexer_losses")
-
-  sparsity_enabled = is_train and config.weight_sparsity_n and config.weight_sparsity_m
-  if sparsity_enabled:
-    mutable_collections.append("batch_stats")
-  if isinstance(model, nn.Module):
-    # inputs, targets, segments, positions = apply_args
-    if dropout_rng is not None:
-      rng1, aqt_rng = jax.random.split(dropout_rng)
-    else:
-      rng1, aqt_rng = None, None
-
-    # Flax Linen model
-    if sparsity_enabled:
-      model_vars = {"params": params}
-      if sparsity_state:
-        model_vars["batch_stats"] = sparsity_state
-    else:
-      model_vars = params
-    logits, intermediate_outputs = model.apply(
-        model_vars,
-        data["inputs"],
-        data["inputs_position"],
-        decoder_segment_ids=data["inputs_segmentation"],
-        encoder_images=encoder_images,
-        encoder_image_masks=encoder_image_masks,
-        enable_dropout=config.enable_dropout if is_train else False,
-        rngs={"dropout": rng1, "params": aqt_rng},  # pyrefly: ignore[bad-argument-type]
-        mutable=mutable_collections,
-        decoder_target_tokens=data["targets"],
-        decoder_target_mask=data["targets_segmentation"],
-    )
-
-    if (config.use_indexer and not config.indexer_sparse_training) and is_train:
-      # In Dense Warm-up stage, we skip main model loss calculation for efficiency.
-      # The main model parameters are frozen and only the indexer is trained via KL divergence.
-      xent_sum = 0.0
-      total_z_loss = 0.0
-    elif config.num_vocab_tiling > 1:
-      hidden_state_key = ("intermediates", "decoder", "hidden_states")
-      hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
-      xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train)
-    else:
-      if is_block_diffusion:
-        logits = block_diffusion_target_alignment.align_logits_to_targets(
-            logits,
-            config.block_diffusion_logit_alignment,
-            target_positions,
-            data["targets_segmentation"] != 0,
-        )
-      one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size)
-      xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier)
-
-      xent = sharding.maybe_shard_with_logical(
-          xent,
-          ("activation_embed_and_logits_batch", "activation_length"),
-          model.mesh,
-          config.shard_mode,
-          debug_sharding=config.debug_sharding,
-      )
-      z_loss = sharding.maybe_shard_with_logical(
-          z_loss,
-          ("activation_embed_and_logits_batch", "activation_length"),
-          model.mesh,
-          config.shard_mode,
-          debug_sharding=config.debug_sharding,
-      )
-
-      if is_block_diffusion:
-        xent = xent * targets_loss_mask
-        z_loss = z_loss * targets_loss_mask
-      else:
-        xent = xent * (data["targets_segmentation"] != 0)
-        z_loss = z_loss * (data["targets_segmentation"] != 0)
-
-      xent_sum = jnp.sum(xent)
-      total_z_loss = jnp.sum(z_loss)
+  # Flax NNX model: forward pass, then pop Intermediates sown during it.
+  logits = model(
+      decoder_input_tokens=data["inputs"],
+      decoder_positions=data["inputs_position"],
+      decoder_segment_ids=data["inputs_segmentation"],
+      encoder_images=encoder_images,
+      encoder_image_masks=encoder_image_masks,
+      enable_dropout=config.enable_dropout if is_train else False,
+      decoder_target_tokens=data["targets"],
+      decoder_target_mask=data["targets_segmentation"],
+  )
+  # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match
+  # subclasses. Pop them before the generic Intermediate pop below, which would otherwise
+  # take them too and leave the MTP loss silently reading as 0.
+  mtp_losses_state, mtp_acceptance_state = None, None
+  if config.mtp_num_layers > 0:
+    mtp_losses_state = nnx.pop(model, mtp_losses)
+    mtp_acceptance_state = nnx.pop(model, mtp_acceptance)
+
+  indexer_losses_state = None
+  if config.use_indexer:
+    # Pop dedicated indexer_losses to harvest auxiliary KL loss and prevent model state PyTree mismatches.
+    indexer_losses_state = nnx.pop(model, indexer_losses)
+
+  intermediates = nnx.pop(model, nnx.Intermediate)
+  intermediate_outputs = intermediates.to_pure_dict()
+
+  # Store them under the collection name so calculate_mtp_loss and
+  # calculate_mtp_acceptance_rate find them at the path they expect.
+  if mtp_losses_state is not None and mtp_acceptance_state is not None:
+    intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict()
+    intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict()
+
+  if indexer_losses_state is not None:
+    intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict()
+
+  if (config.use_indexer and not config.indexer_sparse_training) and is_train:
+    # In Dense Warm-up stage, we skip main model loss calculation for efficiency.
+    # The main model parameters are frozen and only the indexer is trained via KL divergence.
+    xent_sum = 0.0
+    total_z_loss = 0.0
+  elif config.num_vocab_tiling > 1:
+    hidden_state_key = ("decoder", "hidden_states")
+    hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
+    xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train)
   else:
-    # Flax NNX model: forward pass, then pop Intermediates sown during it.
-    logits = model(
-        decoder_input_tokens=data["inputs"],
-        decoder_positions=data["inputs_position"],
-        decoder_segment_ids=data["inputs_segmentation"],
-        encoder_images=encoder_images,
-        encoder_image_masks=encoder_image_masks,
-        enable_dropout=config.enable_dropout if is_train else False,
-        decoder_target_tokens=data["targets"],
-        decoder_target_mask=data["targets_segmentation"],
-    )
-    # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match
-    # subclasses. Pop them before the generic Intermediate pop below, which would otherwise
-    # take them too and leave the MTP loss silently reading as 0.
-    mtp_losses_state, mtp_acceptance_state = None, None
-    if config.mtp_num_layers > 0:
-      mtp_losses_state = nnx.pop(model, mtp_losses)
-      mtp_acceptance_state = nnx.pop(model, mtp_acceptance)
-
-    indexer_losses_state = None
-    if config.use_indexer:
-      # Pop dedicated indexer_losses to harvest auxiliary KL loss and prevent model state PyTree mismatches.
-      indexer_losses_state = nnx.pop(model, indexer_losses)
-
-    intermediates = nnx.pop(model, nnx.Intermediate)
-    intermediate_outputs = intermediates.to_pure_dict()
-
-    # Store them under the collection name so calculate_mtp_loss and
-    # calculate_mtp_acceptance_rate find them at the same path as the Linen collections.
-    if mtp_losses_state is not None and mtp_acceptance_state is not None:
-      intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict()
-      intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict()
-
-    if indexer_losses_state is not None:
-      intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict()
-
-    if (config.use_indexer and not config.indexer_sparse_training) and is_train:
-      # In Dense Warm-up stage, we skip main model loss calculation for efficiency.
-      # The main model parameters are frozen and only the indexer is trained via KL divergence.
-      xent_sum = 0.0
-      total_z_loss = 0.0
-    elif config.num_vocab_tiling > 1:
-      hidden_state_key = ("decoder", "hidden_states")
-      hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
-      xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train)
-    else:
-      if is_block_diffusion:
-        logits = block_diffusion_target_alignment.align_logits_to_targets(
-            logits,
-            config.block_diffusion_logit_alignment,
-            target_positions,
-            data["targets_segmentation"] != 0,
-        )
-      one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size)
-      xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier)
-
-      xent = sharding.maybe_shard_with_logical(
-          xent,
-          ("activation_embed_and_logits_batch", "activation_length"),
-          model.mesh,
-          config.shard_mode,
-          debug_sharding=config.debug_sharding,
-      )
-      z_loss = sharding.maybe_shard_with_logical(
-          z_loss,
-          ("activation_embed_and_logits_batch", "activation_length"),
-          model.mesh,
-          config.shard_mode,
-          debug_sharding=config.debug_sharding,
+    if is_block_diffusion:
+      logits = block_diffusion_target_alignment.align_logits_to_targets(
+          logits,
+          config.block_diffusion_logit_alignment,
+          target_positions,
+          data["targets_segmentation"] != 0,
       )
+    one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size)
+    xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier)
+
+    xent = sharding.maybe_shard_with_logical(
+        xent,
+        ("activation_embed_and_logits_batch", "activation_length"),
+        model.mesh,
+        config.shard_mode,
+        debug_sharding=config.debug_sharding,
+    )
+    z_loss = sharding.maybe_shard_with_logical(
+        z_loss,
+        ("activation_embed_and_logits_batch", "activation_length"),
+        model.mesh,
+        config.shard_mode,
+        debug_sharding=config.debug_sharding,
+    )
 
-      if is_block_diffusion:
-        xent = xent * targets_loss_mask
-        z_loss = z_loss * targets_loss_mask
-      else:
-        xent = xent * (data["targets_segmentation"] != 0)
-        z_loss = z_loss * (data["targets_segmentation"] != 0)
+    if is_block_diffusion:
+      xent = xent * targets_loss_mask
+      z_loss = z_loss * targets_loss_mask
+    else:
+      xent = xent * (data["targets_segmentation"] != 0)
+      z_loss = z_loss * (data["targets_segmentation"] != 0)
 
-      xent_sum = jnp.sum(xent)
-      total_z_loss = jnp.sum(z_loss)
+    xent_sum = jnp.sum(xent)
+    total_z_loss = jnp.sum(z_loss)
 
   if is_block_diffusion:
     assert targets_loss_mask is not None
@@ -370,24 +279,20 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
   moe_bias_updates = None
   mtp_moe_bias_updates = None
   if config.routed_bias and config.routed_bias_update_rate > 0.0:
-    if isinstance(model, nn.Module):
-      nested_key = ("intermediates", "decoder", "moe_layers", "moe_bias_updates")
-      moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, nested_key, None)
-    else:
-      # NNX intermediates are model-rooted (no "intermediates" prefix),
-      # so match by suffix instead. Unlike collect_intermediates_by_suffix
-      # we must not ravel: the decoder update is a 2-D matrix that's
-      # transposed and MTP update is 1-D matrix.
-      for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs):
-        keys = tuple(k.key for k in path if hasattr(k, "key"))
-        if not keys or keys[-1] != "moe_bias_updates":
-          continue
-        if "decoder" in keys:
-          moe_bias_updates = (val,)
-        elif "mtp_block" in keys:
-          if mtp_moe_bias_updates is None:
-            mtp_moe_bias_updates = []
-          mtp_moe_bias_updates.append(val)
+    # NNX intermediates are model-rooted (no "intermediates" prefix),
+    # so match by suffix instead. Unlike collect_intermediates_by_suffix
+    # we must not ravel: the decoder update is a 2-D matrix that's
+    # transposed and MTP update is 1-D matrix.
+    for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs):
+      keys = tuple(k.key for k in path if hasattr(k, "key"))
+      if not keys or keys[-1] != "moe_bias_updates":
+        continue
+      if "decoder" in keys:
+        moe_bias_updates = (val,)
+      elif "mtp_block" in keys:
+        if mtp_moe_bias_updates is None:
+          mtp_moe_bias_updates = []
+        mtp_moe_bias_updates.append(val)
 
   # Add the model's primary output to the intermediates dict so it can be used
   # by the acceptance rate calculation in eval_step.
@@ -420,29 +325,26 @@ def _find_gate_bias(module: nnx.Module | None) -> nnx.Variable | None:
 
 
 def train_step(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng=None):
-  """Training step for both Linen and NNX models.
+  """Training step for the NNX model.
 
   Args:
-    model: A nn.Module (Linen) or nnx.GraphDef of the TrainStateNNX (NNX).
+    model: An nnx.GraphDef of the TrainStateNNX.
     config: Hyperparameters.
     state_mesh_shardings: PyTree of PartitionSpecs for the train state.
     params_shardings: PyTree of PartitionSpecs for model parameters, used for gradient accumulation.
-    state: Linen TrainState or NNX pure State.
+    state: NNX pure State.
     data: Training data batch.
-    dropout_rng: A key to use to generate rng for dropout (Linen); unused for NNX.
+    dropout_rng: Unused for NNX (kept for jit signature parity).
 
   Returns:
-    new_state: Updated Linen TrainState or NNX pure State.
+    new_state: Updated NNX pure State.
     metrics: Dictionary of model metrics such as loss, training rate, etc.
   """
+  del dropout_rng  # unused for NNX (kept for jit signature parity)
   # pylint: disable=too-many-nested-blocks
   # --- Per-path initialization ---
-  if isinstance(model, nn.Module):
-    params = state.params
-    loss_model, loss_params, loss_rng = model, params, dropout_rng
-  else:
-    state = nnx.merge(model, state)  # reconstruct TrainStateNNX
-    loss_model, loss_params, loss_rng = state.model, None, None
+  state = nnx.merge(model, state)  # reconstruct TrainStateNNX
+  loss_model, loss_params, loss_rng = state.model, None, None
 
   # --- Gradient computation ---
   if config.gradient_accumulation_steps > 1:
@@ -456,80 +358,56 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
         loss_rng,
     )
   else:
-    if isinstance(model, nn.Module):
-      if config.shard_optimizer_over_data:
-        params = jax.tree.map(
-            functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode),
-            params,  # pyrefly: ignore[unbound-name]
-            params_shardings,
-        )
-      sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m
-      pure_params = params["params"] if sparsity_enabled else params  # pyrefly: ignore[unbound-name]
-      batch_stats = params.get("batch_stats", {})
-
-      grad_func = jax.value_and_grad(loss_fn, argnums=4, has_aux=True)
-      (loss, aux), raw_grads = grad_func(
-          model,
-          config,
-          data,
-          dropout_rng,
-          pure_params,
-          sparsity_state=batch_stats,
-          is_train=True,
-      )
-    else:
-      owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True)
-      custom_param_filter = nnx.Any(owg_type)
-      train_param_type = (
-          getattr(nnx, "LoRAParam", nnx.Param)
-          if getattr(getattr(config, "lora", None), "enable_lora", False)
-          else nnx.Param
+    owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True)
+    custom_param_filter = nnx.Any(owg_type)
+    train_param_type = (
+        getattr(nnx, "LoRAParam", nnx.Param)
+        if getattr(getattr(config, "lora", None), "enable_lora", False)
+        else nnx.Param
+    )
+    nnx.pop(state.model, nnx.Intermediate)
+    model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, train_param_type, custom_param_filter, ...)
+    if config.parameter_memory_host_offload:
+      # Params are kept on host (pinned_host) in in_shardings. Move only Param
+      # variables to device before the forward/backward pass so that all dot_general
+      # operands share the same memory space (XLA on GPU requires this).
+      # Using params_shardings (Param-only) avoids Shardy rank mismatches that
+      # occur when applying PartitionSpec() (rank-0 in SDY) to rank-1 RNG key tensors.
+      device_param_shardings = jax.tree_util.tree_map_with_path(
+          maxtext_utils_nnx.move_memory_to_device,
+          params_shardings,
+          is_leaf=lambda x: isinstance(x, NamedSharding),
       )
-      nnx.pop(state.model, nnx.Intermediate)
-      model_graphdef, curr_params, custom_params, rest = nnx.split(
-          state.model, train_param_type, custom_param_filter, ...
+      curr_params = jax.device_put(curr_params, device_param_shardings)
+      nnx.update(state.model, curr_params)  # ensure state.model has device params for optimizer update
+    if config.shard_optimizer_over_data:
+      param_sharding_lookup = {}
+      for p, s in jax.tree_util.tree_leaves_with_path(
+          params_shardings, is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax.sharding.Sharding))
+      ):
+        param_sharding_lookup[p] = s.get_value() if isinstance(s, nnx.Variable) else s
+
+      def _maybe_shard_param(path, var):
+        if path in param_sharding_lookup:
+          return sharding.maybe_shard_with_name(var, param_sharding_lookup[path], shard_mode=config.shard_mode)
+        return var
+
+      curr_params = jax.tree_util.tree_map_with_path(
+          _maybe_shard_param,
+          curr_params,
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
       )
-      if config.parameter_memory_host_offload:
-        # Params are kept on host (pinned_host) in in_shardings. Move only Param
-        # variables to device before the forward/backward pass so that all dot_general
-        # operands share the same memory space (XLA on GPU requires this).
-        # Using params_shardings (Param-only) avoids Shardy rank mismatches that
-        # occur when applying PartitionSpec() (rank-0 in SDY) to rank-1 RNG key tensors.
-        device_param_shardings = jax.tree_util.tree_map_with_path(
-            maxtext_utils_nnx.move_memory_to_device,
-            params_shardings,
-            is_leaf=lambda x: isinstance(x, NamedSharding),
-        )
-        curr_params = jax.device_put(curr_params, device_param_shardings)
-        nnx.update(state.model, curr_params)  # ensure state.model has device params for optimizer update
-      if config.shard_optimizer_over_data:
-        param_sharding_lookup = {}
-        for p, s in jax.tree_util.tree_leaves_with_path(
-            params_shardings, is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax.sharding.Sharding))
-        ):
-          param_sharding_lookup[p] = s.get_value() if isinstance(s, nnx.Variable) else s
-
-        def _maybe_shard_param(path, var):
-          if path in param_sharding_lookup:
-            return sharding.maybe_shard_with_name(var, param_sharding_lookup[path], shard_mode=config.shard_mode)
-          return var
-
-        curr_params = jax.tree_util.tree_map_with_path(
-            _maybe_shard_param,
-            curr_params,
-            is_leaf=lambda x: isinstance(x, nnx.Variable),
-        )
-        nnx.update(state.model, curr_params)
-
-      def diff_wrapper(curr_params, custom_params, rest, config, data):
-        local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
-        loss, aux = loss_fn(local_model, config, data, None, None, is_train=True)
-        non_param_rest = nnx.state(local_model, nnx.Not(nnx.Any(nnx.Param, nnx.Intermediate)))
-        return loss, (aux, non_param_rest)
-
-      grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True)
-      (loss, (aux, non_param_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data)
-      nnx.update(state.model, nnx.State.merge(custom_grads, non_param_rest))
+      nnx.update(state.model, curr_params)
+
+    def diff_wrapper(curr_params, custom_params, rest, config, data):
+      local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True)
+      loss, aux = loss_fn(local_model, config, data, None, None, is_train=True)
+      non_param_rest = nnx.state(local_model, nnx.Not(nnx.Any(nnx.Param, nnx.Intermediate)))
+      return loss, (aux, non_param_rest)
+
+    grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True)
+    (loss, (aux, non_param_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data)
+    nnx.update(state.model, nnx.State.merge(custom_grads, non_param_rest))
 
   raw_grads = jax.tree_util.tree_map(
       lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x,
@@ -554,140 +432,78 @@ def diff_wrapper(curr_params, custom_params, rest, config, data):
   new_opt_state = None
   bias_metrics = {}
 
-  if isinstance(model, nn.Module):
-    if config.gradient_clipping_threshold > 0:
-      grads = maxtext_utils.apply_gradient_clipping(raw_grads, state, config.gradient_clipping_threshold)
-    else:
-      grads = raw_grads
-    if config.optimizer_memory_host_offload:
-      state = state.replace(
-          opt_state=jax.device_put(
-              state.opt_state,
-              jax.tree_util.tree_map(
-                  lambda x: x.with_memory_kind(kind="device"),
-                  state_mesh_shardings.opt_state,
-              ),
-          )
-      )
-    # Move all parameters to device before optimizer update
-    if config.parameter_memory_host_offload:
-      max_logging.log("\nMoving all parameters to device before optimizer update")
-
-      def move(path, value):
-        max_logging.log(f"train.py: Moving f{path} to device")
-        return value.with_memory_kind(kind="device")
-
-      state = state.replace(
-          params=jax.device_put(
-              state.params,
-              jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params),
-          )
-      )
-    # Re-wrap grads to match state.params structure if it's a dict of collections
-    # (when weight_sparsity is enabled, params has both 'params' and 'batch_stats' keys).
-    sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m
-    if sparsity_enabled:
-      full_grads = {"params": grads}
-      if "batch_stats" in state.params:
-        batch_stats_grads = jax.tree_util.tree_map(jnp.zeros_like, state.params.get("batch_stats", {}))
-        full_grads["batch_stats"] = batch_stats_grads
-      full_grads = max_utils.unbox_logicallypartioned(full_grads)
-    else:
-      full_grads = grads
-
-    if getattr(config, "skip_step_on_spikes", False):
-      grad_norm = max_utils.l2norm_pytree(grads)
-      # TrainState.apply_gradients doesn't pass **kwargs to tx.update, so we unpack it manually.
-      updates, new_opt_state = state.tx.update(grads, state.opt_state, state.params, loss=loss, grad_norm=grad_norm)
-      new_params = optax.apply_updates(state.params, updates)
-
-      new_state = state.replace(
-          step=state.step + 1,
-          params=new_params,
-          opt_state=new_opt_state,
-      )
-    else:
-      new_state = state.apply_gradients(grads=full_grads)
-
-    # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family
-    if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None:
-      target_path = ("params", "decoder", "moe_layers", "DeepSeekMoeBlock_0", "MoeBlock_0", "gate", "bias")
-      # Updates the shape to be aligned with state.
-      moe_bias_updates = jnp.array(moe_bias_updates[0]).transpose()
-      new_state = maxtext_utils.update_state_param(new_state, target_path, moe_bias_updates)
+  if config.gradient_clipping_threshold > 0:
+    grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold)
   else:
-    if config.gradient_clipping_threshold > 0:
-      grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold)
-    else:
-      grads = raw_grads
-    if config.optimizer_memory_host_offload:
-      # state.optimizer is an NNX Optimizer module; state_mesh_shardings.optimizer
-      # is an NNX State. Use nnx.state() to get a compatible State for device_put.
-      device_opt_shardings = jax.tree_util.tree_map_with_path(
-          maxtext_utils_nnx.move_memory_to_device,
-          state_mesh_shardings.optimizer,
-          is_leaf=lambda x: isinstance(x, NamedSharding),
-      )
-      opt_state = nnx.state(state.optimizer)
-      new_opt_state = jax.device_put(opt_state, device_opt_shardings)
-      nnx.update(state.optimizer, new_opt_state)
-    if config.skip_step_on_spikes:
-      # The skip-step optimizer is a GradientTransformationExtraArgs that reads
-      # loss/grad_norm to decide whether to zero the update on a spike. nnx
-      # Optimizer.update forwards these kwargs to tx.update.
-      grad_norm = max_utils.l2norm_pytree(grads)
-      state.apply_gradients(grads, loss=loss, grad_norm=grad_norm)
+    grads = raw_grads
+  if config.optimizer_memory_host_offload:
+    # state.optimizer is an NNX Optimizer module; state_mesh_shardings.optimizer
+    # is an NNX State. Use nnx.state() to get a compatible State for device_put.
+    device_opt_shardings = jax.tree_util.tree_map_with_path(
+        maxtext_utils_nnx.move_memory_to_device,
+        state_mesh_shardings.optimizer,
+        is_leaf=lambda x: isinstance(x, NamedSharding),
+    )
+    opt_state = nnx.state(state.optimizer)
+    new_opt_state = jax.device_put(opt_state, device_opt_shardings)
+    nnx.update(state.optimizer, new_opt_state)
+  if config.skip_step_on_spikes:
+    # The skip-step optimizer is a GradientTransformationExtraArgs that reads
+    # loss/grad_norm to decide whether to zero the update on a spike. nnx
+    # Optimizer.update forwards these kwargs to tx.update.
+    grad_norm = max_utils.l2norm_pytree(grads)
+    state.apply_gradients(grads, loss=loss, grad_norm=grad_norm)
+  else:
+    state.apply_gradients(grads)
+  new_state = state
+
+  # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family
+  # pylint: disable=too-many-nested-blocks
+  if config.routed_bias and config.routed_bias_update_rate > 0.0:
+    if getattr(config, "model_name", "").startswith("deepseek4"):
+      max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.")
+      flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {}))
+      for path, update in flat_intermediates.items():
+        if path[-1] != "moe_bias_updates":
+          continue
+        target = new_state.model
+        prefix = path[1:-1] if path[0] == "intermediates" else path[:-1]
+        for key in prefix:
+          if hasattr(target, key):
+            target = getattr(target, key)
+          elif isinstance(target, dict) and key in target:
+            target = target[key]
+          else:
+            target = None
+            break
+        if target is None:
+          continue
+        for _, node in nnx.iter_graph(target):
+          if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None:
+            update_val = update[0] if isinstance(update, (tuple, list)) else update
+            name_prefix = "-".join(map(str, prefix))
+            if getattr(config, "log_moe_bias_norms", False):
+              bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value)
+            node.bias.value = node.bias.value + jnp.array(update_val)
+            if getattr(config, "log_moe_bias_norms", False):
+              bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val))
     else:
-      state.apply_gradients(grads)
-    new_state = state
-
-    # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family
-    # pylint: disable=too-many-nested-blocks
-    if config.routed_bias and config.routed_bias_update_rate > 0.0:
-      if getattr(config, "model_name", "").startswith("deepseek4"):
-        max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.")
-        flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {}))
-        for path, update in flat_intermediates.items():
-          if path[-1] != "moe_bias_updates":
-            continue
-          target = new_state.model
-          prefix = path[1:-1] if path[0] == "intermediates" else path[:-1]
-          for key in prefix:
-            if hasattr(target, key):
-              target = getattr(target, key)
-            elif isinstance(target, dict) and key in target:
-              target = target[key]
-            else:
-              target = None
-              break
-          if target is None:
-            continue
-          for _, node in nnx.iter_graph(target):
-            if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None:
-              update_val = update[0] if isinstance(update, (tuple, list)) else update
-              name_prefix = "-".join(map(str, prefix))
-              if getattr(config, "log_moe_bias_norms", False):
-                bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value)
-              node.bias.value = node.bias.value + jnp.array(update_val)
-              if getattr(config, "log_moe_bias_norms", False):
-                bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val))
-      else:
-        # 1. Update main decoder scanned MoE layers.
-        # The update from the scan is (num_moe_layers, num_experts) and must be transposed.
-        decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder)
-        decoder_bias = _find_gate_bias(decoder_layer)
-        if decoder_bias is not None:
-          decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]).transpose()
-
-        # 2. Update auxiliary MTP MoE layers (if enabled).
-        # Unlike the main decoder, each MTP layer is an individual un-scanned layer
-        # with a 1D bias of shape (num_experts,).
-        if mtp_moe_bias_updates is not None and hasattr(new_state.model, "mtp_block"):
-          for i, update in enumerate(mtp_moe_bias_updates):
-            mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None)
-            mtp_bias = _find_gate_bias(mtp_layer)
-            if mtp_bias is not None:
-              mtp_bias.value = mtp_bias.value + jnp.array(update)
+      # 1. Update main decoder scanned MoE layers.
+      # The update from the scan is (num_moe_layers, num_experts) and must be transposed.
+      decoder_layer = getattr(new_state.model.decoder, "moe_layers", new_state.model.decoder)
+      decoder_bias = _find_gate_bias(decoder_layer)
+      if decoder_bias is not None:
+        decoder_bias.value = decoder_bias.value + jnp.array(moe_bias_updates[0]).transpose()
+
+      # 2. Update auxiliary MTP MoE layers (if enabled).
+      # Unlike the main decoder, each MTP layer is an individual un-scanned layer
+      # with a 1D bias of shape (num_experts,).
+      if mtp_moe_bias_updates is not None and hasattr(new_state.model, "mtp_block"):
+        for i, update in enumerate(mtp_moe_bias_updates):
+          mtp_layer = getattr(new_state.model.mtp_block, f"mtp_layer_{i + 1}", None)
+          mtp_bias = _find_gate_bias(mtp_layer)
+          if mtp_bias is not None:
+            mtp_bias.value = mtp_bias.value + jnp.array(update)
 
   lm_loss = xent_sum / (total_weights + EPS)
   scalar_metrics = {
@@ -702,10 +518,7 @@ def move(path, value):
   }
   scalar_metrics.update(bias_metrics)
   if config.use_qk_clip:
-    if isinstance(model, nn.Module):
-      new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config)
-    else:
-      new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config)
+    new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config)
 
     global_max_logit = qk_clip_utils.calculate_max_logit_metric(intermediate_outputs)
     if global_max_logit is not None:
@@ -714,21 +527,14 @@ def move(path, value):
   if not config.optimizer_memory_host_offload:
     scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads)
     scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads)
-    if isinstance(model, nn.Module):
-      scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(new_state.params)
-    else:
-      model_params = nnx.state(new_state.model, nnx.Param)
-      scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params)
+    model_params = nnx.state(new_state.model, nnx.Param)
+    scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params)
 
   # Surface skip-step rejections as a TB metric. The skip-step optimizer stores
-  # is_skipped in its opt_state: the Linen path gets it from the tx.update return,
-  # the NNX path reads it back off the optimizer it just updated in place.
+  # is_skipped in its opt_state; read it back off the optimizer just updated in place.
   if config.skip_step_on_spikes:
-    if isinstance(model, nn.Module):
-      is_skipped = new_opt_state.get("is_skipped") if isinstance(new_opt_state, dict) else None
-    else:
-      opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get("opt_state", {})
-      is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None
+    opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get("opt_state", {})
+    is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None
     if is_skipped is not None:
       scalar_metrics["optim/step_skipped"] = is_skipped.astype(jnp.float32)
   metrics = {
@@ -738,8 +544,6 @@ def move(path, value):
   if getattr(config, "record_internal_nn_metrics", False):
     record_activation_metrics(metrics, intermediate_outputs, config)
 
-  if isinstance(model, nn.Module):
-    return new_state, metrics
   # Drop Intermediates (e.g. sowed max_logits for QK-Clip) and the MTP sown
   # vars (mtp_losses/mtp_acceptance) before returning. They're absent from
   # state_mesh_shardings and would cause a leaf-count / structure mismatch.
@@ -748,16 +552,9 @@ def move(path, value):
 
 def eval_step(model, config, state, data, dropout_rng=None):
   """eval_step no backprop and new state compared with train_step."""
-  if isinstance(model, nn.Module):
-    sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m
-    pure_params = state.params["params"] if sparsity_enabled else state.params
-    batch_stats = state.params.get("batch_stats", {})
-
-    eval_loss_fn = functools.partial(loss_fn, model, config, data, dropout_rng, is_train=False)
-    loss, aux = eval_loss_fn(pure_params, sparsity_state=batch_stats)
-  else:
-    state = nnx.merge(model, state)  # reconstruct TrainStateNNX
-    loss, aux = loss_fn(state.model, config, data, None, None, is_train=False)
+  del dropout_rng  # unused for NNX (kept for jit signature parity)
+  state = nnx.merge(model, state)  # reconstruct TrainStateNNX
+  loss, aux = loss_fn(state.model, config, data, None, None, is_train=False)
 
   mtp_acceptance_rate = 0.0
   if config.mtp_eval_target_module > 0:
@@ -796,10 +593,8 @@ def training_loop_iteration(
   state = jax_device_state["state"]
   init_rng = jax_device_state["init_rng"]
   mesh = jax_device_state["mesh"]
-  state_mesh_shardings = jax_device_state["state_mesh_shardings"]
   p_train_step = jax_device_state["p_train_step"]
   p_eval_step = jax_device_state["p_eval_step"]
-  model = jax_device_state["model"]
 
   # Unpack python_vars
   step = python_vars["step"]
@@ -817,8 +612,6 @@ def training_loop_iteration(
   config = immutable_data["config"]  # for helpers
   logical_axis_rules_for_train = immutable_data["logical_axis_rules_for_train"]
   logical_axis_rules_for_eval = immutable_data["logical_axis_rules_for_eval"]
-  shard_optimizer_over_data = immutable_data["shard_optimizer_over_data"]
-  shard_mode = immutable_data["shard_mode"]
   eval_interval = immutable_data["eval_interval"]
   eval_steps = immutable_data["eval_steps"]
   start_step = immutable_data["start_step"]
@@ -839,16 +632,14 @@ def training_loop_iteration(
 
   with jax.profiler.StepTraceAnnotation("train", step_num=step):
     example_batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
-    # DiLoCo's inner step takes the rng like the Linen step does.
-    if isinstance(model, nn.Module) or config.enable_diloco:
+    # DiLoCo's inner step takes the rng like the inner NNX step.
+    if config.enable_diloco:
       # pylint: disable=not-callable
       step_rng_args = (jax.jit(jax.random.fold_in)(init_rng, step),)
     else:
       step_rng_args = ()
     with maybe_record_goodput(recorder, GoodputEvent.STEP, step):
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_train):
-        if shard_optimizer_over_data and isinstance(model, nn.Module):
-          state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode)
         state, metrics = p_train_step(state, example_batch, *step_rng_args)
 
   step_time_delta = datetime.datetime.now() - last_step_completion
@@ -932,15 +723,13 @@ def train_loop(config, recorder, state=None):
   start_step = get_first_step(model, state)  # this is the start_step for training
   train_utils.validate_completed_steps(start_step, config.steps)
 
-  if isinstance(model, nn.Module):
-    jit_model = model
-  elif config.enable_diloco:
+  if config.enable_diloco:
     # state is the DiLoCoTrainState; `model` is already the TrainStateNNX graphdef the inner step needs.
     jit_model = model
   else:
     jit_model, state = nnx.split(state)
 
-  if config.pure_nnx and config.enable_diloco:
+  if config.enable_diloco:
     # DiLoCoTrainState.params already holds the param shardings the inner step needs;
     # the Zero-1 opt overlay doesn't apply through the diloco wrapper.
     params_shardings = state_mesh_shardings.params
@@ -962,13 +751,11 @@ def train_loop(config, recorder, state=None):
   with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
     data_sharding = sharding.get_input_data_sharding(config, mesh)
     shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
-    if config.shard_optimizer_over_data and isinstance(model, nn.Module):
-      state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)
-    elif config.shard_optimizer_over_data:
+    if config.shard_optimizer_over_data:
       # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout)
       state = jax.device_put(state, state_mesh_shardings)
-    if isinstance(model, nn.Module) or config.enable_diloco:
-      # The DiLoCo train step takes (state, batch, rng), like the Linen step.
+    if config.enable_diloco:
+      # The DiLoCo train step takes (state, batch, rng), like the inner NNX step.
       lower_args = (state, shaped_batch, init_rng)
     else:
       lower_args = (state, shaped_batch)
@@ -982,9 +769,7 @@ def train_loop(config, recorder, state=None):
   metric_logger_instance = metric_logger.MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule)
 
   # Write train config params, num model params, and XLA flags to tensorboard
-  if isinstance(model, nn.Module):
-    setup_params = state.params
-  elif config.enable_diloco:
+  if config.enable_diloco:
     setup_params = state.params  # DiLoCoTrainState.params: the outer (global) params
   else:
     _, setup_params, _ = nnx.split(state.model, nnx.Param, ...)
diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py
index b705bef562..d87ab1a991 100644
--- a/src/maxtext/trainers/pre_train/train_compile.py
+++ b/src/maxtext/trainers/pre_train/train_compile.py
@@ -38,9 +38,8 @@
 import jax.numpy as jnp
 from jax.sharding import AxisType, Mesh
 from maxtext.common import train_state_nnx
-from maxtext.common.common_types import MODEL_MODE_TRAIN, ShardMode
+from maxtext.common.common_types import ShardMode
 from maxtext.configs import pyconfig
-from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.optimizers import optimizers
 from maxtext.trainers.diloco import diloco
@@ -130,68 +129,49 @@ def _nnx_forward(decoder_input_tokens, decoder_positions, decoder_segment_ids):
 def get_shaped_inputs(topology_mesh, config):
   """Get shaped abstractions of inputs to train_step: state, batch and rng"""
   # Construct the model and optimizer to get shaped versions of the state
-  quant = quantizations.configure_quantization(config)
-  if config.pure_nnx:
-    _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, topology_mesh)
-  else:
-    model = Transformer(config, topology_mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
+  _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, topology_mesh)
   # The learning_rate_schedule is baked into the compiled object.
   learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(config)
   # pass in model for muon
   tx = optimizers.get_optimizer(config, learning_rate_schedule, model)
 
-  # Shaped RNG keys
-  _, example_rng = jax.random.split(jax.random.PRNGKey(0), 2)
-  shaped_rng = jax.ShapeDtypeStruct(example_rng.shape, example_rng.dtype)
-
-  if config.pure_nnx:
-
-    def create_train_state_fn():
-      nnx_model = _create_model_partial()
-      wrt = (
-          getattr(nnx, "LoRAParam", nnx.Param)
-          if getattr(getattr(config, "lora", None), "enable_lora", False)
-          else nnx.Param
-      )
-      optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
-      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+  def create_train_state_fn():
+    nnx_model = _create_model_partial()
+    wrt = (
+        getattr(nnx, "LoRAParam", nnx.Param)
+        if getattr(getattr(config, "lora", None), "enable_lora", False)
+        else nnx.Param
+    )
+    optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
+    return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-    init_state_fn = create_train_state_fn
-  else:
-    init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, example_rng)
+  init_state_fn = create_train_state_fn
 
   # Shaped state
   abstract_state, _, state_mesh_shardings = maxtext_utils.get_abstract_state(config, topology_mesh, init_state_fn, True)
 
-  if config.pure_nnx:
-    # NNX doesn't use Linen logical annotations; derive PartitionSpecs from the physical shardings.
-    logical_annotations = maxtext_utils_nnx.get_partition_spec_nnx(state_mesh_shardings)
-    # For NNX, get_functional_train_with_signature expects the graphdef (static structure),
-    # not the raw model — mirroring how the training loop does nnx.split(train_state).
-    with nn_partitioning.axis_rules(config.logical_axis_rules):
-      abs_train_state = nnx.eval_shape(init_state_fn)
-      graphdef, _ = nnx.split(abs_train_state)
-    model = graphdef
-  else:
-    # unsharded logical annotations
-    logical_annotations = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn)
+  # NNX doesn't use Linen logical annotations; derive PartitionSpecs from the physical shardings.
+  logical_annotations = maxtext_utils_nnx.get_partition_spec_nnx(state_mesh_shardings)
+  # For NNX, get_functional_train_with_signature expects the graphdef (static structure),
+  # not the raw model — mirroring how the training loop does nnx.split(train_state).
+  with nn_partitioning.axis_rules(config.logical_axis_rules):
+    abs_train_state = nnx.eval_shape(init_state_fn)
+    graphdef, _ = nnx.split(abs_train_state)
+  model = graphdef
 
   # Shaped batch
   data_sharding = sharding.get_input_data_sharding(config, topology_mesh)
   shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
 
-  if config.pure_nnx:
-    shaped_train_args = (
-        abstract_state,
-        shaped_batch,
-    )  # NNX doesn't use dropout_rng
-  else:
-    shaped_train_args = (abstract_state, shaped_batch, shaped_rng)
+  shaped_train_args = (
+      abstract_state,
+      shaped_batch,
+  )  # NNX doesn't use dropout_rng
   shaped_train_kwargs = {}
 
   # Collect NNX activation shardings via an abstract forward pass (must run
   # after get_abstract_state, which only traces __init__).
-  if config.debug_sharding and config.pure_nnx:
+  if config.debug_sharding:
     _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh)  # pyrefly: ignore[unbound-name]
 
   return (
@@ -392,20 +372,12 @@ def main(argv: Sequence[str]) -> None:
   # print weights sharding info under debug sharding mode
   if config.debug_sharding:
     max_utils.print_non_trivial_mesh_axis(topology_mesh)
-    if config.pure_nnx:
-      maxtext_utils.print_shardings_params(
-          shaped_train_args[0],
-          state_mesh_shardings,
-          topology_mesh,
-          logical_annotations,
-      )
-    else:
-      maxtext_utils.print_shardings_params(
-          shaped_train_args[0].params,
-          state_mesh_shardings.params,
-          topology_mesh,
-          logical_annotations.params,
-      )
+    maxtext_utils.print_shardings_params(
+        shaped_train_args[0],
+        state_mesh_shardings,
+        topology_mesh,
+        logical_annotations,
+    )
 
   # Compile
   print("Jitting and compiling train step...", flush=True)
diff --git a/tests/assets/golden_logits/golden_dpo_correctness.json b/tests/assets/golden_logits/golden_dpo_correctness.json
index 5552ebde26..7b3dced630 100644
--- a/tests/assets/golden_logits/golden_dpo_correctness.json
+++ b/tests/assets/golden_logits/golden_dpo_correctness.json
@@ -2,17 +2,17 @@
   "explicit_prompt_len_3_column": {
     "loss_step_1": 0.6931471824645996,
     "margin_step_1": 0.0,
-    "loss": 0.6573728919029236,
-    "margin": 0.0728759765625,
-    "chosen_logps": -854.4208984375,
-    "rejected_logps": -530.0099487304688
+    "loss": 0.690365,
+    "margin": 0.005573,
+    "chosen_logps": -882.7361450195312,
+    "rejected_logps": -587.572693
   },
   "default_prompt_len_2_column": {
     "loss_step_1": 0.6931471824645996,
     "margin_step_1": 0.0,
-    "loss": 0.6954630613327026,
-    "margin": -0.004626465495675802,
-    "chosen_logps": -875.7848510742188,
-    "rejected_logps": -502.052978515625
+    "loss": 0.620875,
+    "margin": 0.150177,
+    "chosen_logps": -907.599609375,
+    "rejected_logps": -578.226013
   }
-}
\ No newline at end of file
+}
diff --git a/tests/assets/logits_generation/generate_grpo_golden_logits.py b/tests/assets/logits_generation/generate_grpo_golden_logits.py
index a7d87a1c3f..193030b8a6 100644
--- a/tests/assets/logits_generation/generate_grpo_golden_logits.py
+++ b/tests/assets/logits_generation/generate_grpo_golden_logits.py
@@ -25,7 +25,6 @@
 import unittest
 
 from datasets import load_dataset
-from flax import linen as nn
 from flax import nnx
 import jax
 import jax.numpy as jnp
@@ -33,11 +32,10 @@
 import jsonlines
 from maxtext.configs import pyconfig, types
 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, grpo_loss_fn_nnx
-from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
+from maxtext.common.common_types import Array
+from maxtext.experimental.rl.grpo_trainer import generate_completions, grpo_loss_fn_nnx
+from maxtext.experimental.rl.grpo_utils import 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
@@ -49,40 +47,28 @@
 
 
 def _setup_model(config, mesh, rng):
-  """Builds the model, and for NNX a frozen reference clone, dispatching on pure_nnx.
+  """Builds the model and a frozen reference clone.
 
-  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.
+  Returns (model, reference_model). The model carries its own params (from_pretrained
+  loads the checkpoint or inits) and the reference is a clone of the policy.
   """
-  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)
+  model = model_creation_utils.from_pretrained(config, mesh=mesh, rng_key=rng)
+  return model, nnx.clone(model)
 
 
-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 _logps(config, model, ids, pos, seg, comp_seg):
+  """Policy per-token log-probs."""
+  return compute_log_probs_nnx(model, 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 _reference_logps(config, reference_model, ids, pos, seg, comp_seg):
+  """Reference per-token log-probs, using the cloned reference model."""
+  return compute_log_probs_nnx(reference_model, 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)
+def _grpo_loss(config, model, reference_model, data, rng):
+  """GRPO loss, using the cloned reference model."""
+  return grpo_loss_fn_nnx(model, config, data, rng, None, reference_model)
 
 
 class GRPOTest(unittest.TestCase):
@@ -116,19 +102,12 @@ def setUp(self):
     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:
-      self.state = None
-      self.state_mesh_shardings = None  # NNX param shardings are derived in the generation step.
-    else:
-      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.model, self.reference_model = _setup_model(self.cfg, mesh, self.rng)
     self.data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
     # Without checkpoint
-    self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading, linen_state_no_ckpt = _setup_model(
+    self.model_no_ckpt_loading, self.reference_model_no_ckpt_loading = _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",
@@ -217,29 +196,16 @@ def test_w_trl_and_write_golden_data(self):
         self.cfg.prompt, self.tokenizer_model
     )
     maxtext_per_token_logps, _ = _logps(
-        self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation
+        self.cfg, self.model, input_ids, input_position, input_segmentation, completion_segmentation
     )
 
-    # 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,
         "prompt_completions_position": input_position,
         "prompt_completions_segmentation": input_segmentation,
         "ar_completions_segmentation": completion_segmentation,
     }
-    maxtext_loss, aux = _grpo_loss(
-        self.cfg, self.model, self.reference_model, self.state, reference_params, data, self.rng
-    )
+    maxtext_loss, aux = _grpo_loss(self.cfg, self.model, self.reference_model, 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())
@@ -247,13 +213,11 @@ def test_w_trl_and_write_golden_data(self):
     self.assertEqual(aux.avg_advantage.tolist(), 0.0)
     # since we are at step 0
     maxtext_per_token_logps, _ = _logps(
-        self.cfg, self.model, self.state, input_ids, input_position, input_segmentation, completion_segmentation
+        self.cfg, self.model, input_ids, input_position, input_segmentation, completion_segmentation
     )
     maxtext_per_token_logps_ref, _ = _reference_logps(
         self.cfg,
-        self.model,
         self.reference_model,
-        reference_params,
         input_ids,
         input_position,
         input_segmentation,
@@ -274,7 +238,6 @@ def test_w_trl_and_write_golden_data(self):
     maxtext_per_token_logps_no_ckpt_loading, _ = _logps(
         self.cfg_no_ckpt_loading,
         self.model_no_ckpt_loading,
-        self.state_no_ckpt_loading,
         input_ids,
         input_position,
         input_segmentation,
@@ -285,8 +248,6 @@ def test_w_trl_and_write_golden_data(self):
         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,
     )
@@ -301,13 +262,9 @@ 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
+    # Params live on the model; the inference engine reads them directly.
+    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)
     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, gen_param_shardings, None),
diff --git a/tests/integration/diloco_test.py b/tests/integration/diloco_test.py
index 3b5650dbf2..6ec3b6618a 100644
--- a/tests/integration/diloco_test.py
+++ b/tests/integration/diloco_test.py
@@ -23,7 +23,6 @@
 
 import chex
 from flax.experimental import nnx
-from flax.training import train_state
 import jax
 import jax.numpy as jnp
 import jax.sharding
@@ -88,72 +87,36 @@ def test_diloco_training_simulation_with_mesh(self):
       tx = optax.sgd(learning_rate=0.1)
       rngs = nnx.Rngs(params=jax.random.key(seed=42))
       model = SimpleNNXModel(rngs=rngs)
-      graphdef, params = nnx.split(model)
 
-      if test_config.pure_nnx:
-        optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
-        # diloco_test_state expects a TrainStateNNX instance when pure_nnx is True.
-        initial_test_state = TrainStateNNX(model, optimizer)
+      optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
+      # diloco_test_state expects a TrainStateNNX instance.
+      initial_test_state = TrainStateNNX(model, optimizer)
 
-        # For NNX, train_step needs to take the TrainStateNNX and mutate it
+      # train_step takes the TrainStateNNX and mutates it.
 
-        def _test_train_step(state, batch, prng_key: diloco.PRNGKey):
-          del prng_key
+      def _test_train_step(state, batch, prng_key: diloco.PRNGKey):
+        del prng_key
 
-          def loss_fn(model, batch):
-            inputs, labels = batch
-            logits = jax.vmap(model)(inputs)
-            residual = logits - labels
-            return jnp.mean(jnp.square(residual))
+        def loss_fn(model, batch):
+          inputs, labels = batch
+          logits = jax.vmap(model)(inputs)
+          residual = logits - labels
+          return jnp.mean(jnp.square(residual))
 
-          loss, grads = nnx.value_and_grad(loss_fn)(state.model, batch)
-          state.optimizer.update(state.model, grads)
-          return state, loss
-
-      else:
-
-        def nnx_apply_fn(params, inputs):
-          model_replica = nnx.merge(graphdef, params)
-          return model_replica(inputs)
-
-        # 2. Vmap this new wrapper function
-        vmapped_apply = jax.vmap(nnx_apply_fn, in_axes=(None, 0))
-
-        def _test_train_step(state: train_state.TrainState, batch, prng_key: diloco.PRNGKey):
-          """A simple MSE loss train step to enable numerics testing."""
-          del prng_key
-
-          def loss_fn(params, batch):
-            inputs, labels = batch
-            logits = vmapped_apply(params, inputs)
-            residual = logits - labels
-            sq_residual = jnp.square(residual)
-            msq_residual = jnp.mean(sq_residual)
-            return msq_residual
-
-          loss, grad = jax.value_and_grad(loss_fn)(state.params, batch)
-          return state.apply_gradients(grads=grad), loss
-
-        initial_test_state = train_state.TrainState.create(
-            apply_fn=vmapped_apply,
-            params=params,
-            tx=tx,
-        )
+        loss, grads = nnx.value_and_grad(loss_fn)(state.model, batch)
+        state.optimizer.update(state.model, grads)
+        return state, loss
 
       diloco_test_state, _ = diloco.build_diloco_state(test_config, lambda: initial_test_state)
       chex.assert_equal(diloco_test_state.step, 0)
-      if test_config.pure_nnx:
-        _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
-
-        # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
-        # We need to unwrap them if they do.
-        diloco_params_pure = jax.tree_util.tree_map(
-            lambda x: x.value if hasattr(x, "value") else x,
-            diloco_test_state.params,
-        )
-        chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
-      else:
-        chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params)
+      _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
+
+      # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
+      # We need to unwrap them if they do.
+      diloco_params_pure = jax.tree_util.tree_map(
+          lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params
+      )
+      chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
 
       diloco_train_step = diloco.build_diloco_train_step(test_config, _test_train_step)
       inputs = jnp.array(
@@ -201,18 +164,14 @@ def loss_fn(params, batch):
       chex.assert_equal(diloco_test_state.step, 1.0)
       chex.assert_equal(loss, 1.0)
       # Assert no updates to the global model yet (no synchronization)
-      if test_config.pure_nnx:
-        _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
-
-        # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
-        # We need to unwrap them if they do.
-        diloco_params_pure = jax.tree_util.tree_map(
-            lambda x: x.value if hasattr(x, "value") else x,
-            diloco_test_state.params,
-        )
-        chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
-      else:
-        chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params)
+      _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
+
+      # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
+      # We need to unwrap them if they do.
+      diloco_params_pure = jax.tree_util.tree_map(
+          lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params
+      )
+      chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
 
       # Run the second step (no synchronization).
       # Replica 0:
@@ -242,18 +201,14 @@ def loss_fn(params, batch):
       chex.assert_equal(diloco_test_state.step, 2.0)
       chex.assert_trees_all_close(loss, 0.49, rtol=1e-2, atol=1e-2)
       # Assert no updates to the global model yet (no synchronization)
-      if test_config.pure_nnx:
-        _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
-
-        # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
-        # We need to unwrap them if they do.
-        diloco_params_pure = jax.tree_util.tree_map(
-            lambda x: x.value if hasattr(x, "value") else x,
-            diloco_test_state.params,
-        )
-        chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
-      else:
-        chex.assert_trees_all_equal(diloco_test_state.params, initial_test_state.params)
+      _, params_pure, _ = nnx.split(initial_test_state.model, nnx.Param, ...)
+
+      # diloco_test_state.params might contain nnx.Variables instead of pure arrays.
+      # We need to unwrap them if they do.
+      diloco_params_pure = jax.tree_util.tree_map(
+          lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params
+      )
+      chex.assert_trees_all_equal(diloco_params_pure, params_pure.to_pure_dict())
 
       # Run the third step, which synchronizes afterwards.
       # Replica 0:
@@ -288,33 +243,21 @@ def loss_fn(params, batch):
       chex.assert_trees_all_close(loss, 0.2401, rtol=1e-2, atol=1e-2)
       # Assert that inner and outer parameters are all equal now that
       # synchronization has happened.
-      if test_config.pure_nnx:
-        _, inner_params, _ = nnx.split(diloco_test_state.inner_state.model, nnx.Param, ...)
-        inner_params_pure = jax.tree_util.tree_map(
-            lambda x: x.value if hasattr(x, "value") else x,
-            inner_params.to_pure_dict(),
-        )
-        diloco_params_pure_3 = jax.tree_util.tree_map(
-            lambda x: x.value if hasattr(x, "value") else x,
-            diloco_test_state.params,
-        )
-        chex.assert_trees_all_equal(
-            diloco_params_pure_3,
-            jax.tree.map(lambda arr: arr[0, ...], inner_params_pure),
-        )
-        chex.assert_trees_all_equal(
-            diloco_params_pure_3,
-            jax.tree.map(lambda arr: arr[1, ...], inner_params_pure),
-        )
-      else:
-        chex.assert_trees_all_equal(
-            diloco_test_state.params,
-            jax.tree.map(lambda arr: arr[0, ...], diloco_test_state.inner_state.params),
-        )
-        chex.assert_trees_all_equal(
-            diloco_test_state.params,
-            jax.tree.map(lambda arr: arr[1, ...], diloco_test_state.inner_state.params),
-        )
+      _, inner_params, _ = nnx.split(diloco_test_state.inner_state.model, nnx.Param, ...)
+      inner_params_pure = jax.tree_util.tree_map(
+          lambda x: x.value if hasattr(x, "value") else x, inner_params.to_pure_dict()
+      )
+      diloco_params_pure_3 = jax.tree_util.tree_map(
+          lambda x: x.value if hasattr(x, "value") else x, diloco_test_state.params
+      )
+      chex.assert_trees_all_equal(
+          diloco_params_pure_3,
+          jax.tree.map(lambda arr: arr[0, ...], inner_params_pure),
+      )
+      chex.assert_trees_all_equal(
+          diloco_params_pure_3,
+          jax.tree.map(lambda arr: arr[1, ...], inner_params_pure),
+      )
 
       # Run the fourth step (no synchronization).
       # Replica 0:
diff --git a/tests/post_training/integration/dpo_correctness_base.py b/tests/post_training/integration/dpo_correctness_base.py
index eab78cae1e..11346f065b 100644
--- a/tests/post_training/integration/dpo_correctness_base.py
+++ b/tests/post_training/integration/dpo_correctness_base.py
@@ -202,9 +202,6 @@ def build_tiny_qwen2_jax_config(
         "per_device_batch_size=1",
         f"max_target_length={max_target_length}",
         "skip_jax_distributed_system=True",
-        "enable_nnx=True",
-        "pure_nnx=True",
-        "pure_nnx_decoder=False",
         "remat_policy=full",
         "log_config=0",
         # Tiny architecture specifications.
diff --git a/tests/post_training/integration/grpo_correctness.py b/tests/post_training/integration/grpo_correctness.py
index d2783ea372..98a568d08c 100644
--- a/tests/post_training/integration/grpo_correctness.py
+++ b/tests/post_training/integration/grpo_correctness.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """GRPO correctness tests"""
 
-import functools
 import os
 import unittest
 
@@ -23,11 +22,9 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from maxtext.configs import pyconfig, types
-from maxtext.common.common_types import MODEL_MODE_TRAIN
-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.experimental.rl.grpo_trainer import grpo_loss_fn_nnx
+from maxtext.experimental.rl.grpo_utils import 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
@@ -68,17 +65,11 @@ def setUp(self):
     self.rng = jax.random.PRNGKey(42)
     devices_array = maxtext_utils.create_device_mesh(self.cfg)
     mesh = Mesh(devices_array, self.cfg.mesh_axes)
-    if self.cfg.pure_nnx:
-      # 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.reference_model = None
-      self.state, _ = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn)
+    # 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
     tokenizer_path = ensure_tokenizer_downloaded("llama3.1-tokenizer", skip_test_on_failure=True)
     self.tokenizer_model = transformers.AutoTokenizer.from_pretrained(
         tokenizer_path,
@@ -152,57 +143,25 @@ def _prepare_trl_inputs(self):
     return input_ids, attention_mask, logits_to_keep
 
   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,
-        inputs_position,
+    """NNX forward pass logits."""
+    return self.model(
+        decoder_input_tokens=inputs,
+        decoder_positions=inputs_position,
         decoder_segment_ids=inputs_segmentation,
         enable_dropout=False,
-        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,
+    """Policy per-token log-probs."""
+    return compute_log_probs_nnx(
+        self.model, 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},
+    """Reference per-token log-probs from the cloned reference model."""
+    del reference_params  # NNX reads the cloned reference model directly
+    return compute_log_probs_nnx(
+        self.reference_model,
         input_ids,
         input_position,
         input_segmentation,
@@ -212,10 +171,9 @@ def _reference_logps(self, input_ids, input_position, input_segmentation, comple
     )
 
   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)
+    """GRPO loss from the NNX policy and cloned reference model."""
+    del reference_params  # NNX reads the cloned reference model directly
+    return grpo_loss_fn_nnx(self.model, self.cfg, data, self.rng, None, self.reference_model)
 
   def test_logits(self):
     def _prepare_inputs():
@@ -311,12 +269,9 @@ def test_loss_kl_div(self):
         input_ids, input_position, input_segmentation, completion_segmentation
     )
 
-    # 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.
+    # The reference is a frozen copy of the step-0 policy, held as a cloned model
+    # built in setUp; the NNX loss path reads it directly.
     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,
diff --git a/tests/post_training/integration/grpo_trainer_correctness_test.py b/tests/post_training/integration/grpo_trainer_correctness_test.py
index 95bfd28190..725e19a455 100644
--- a/tests/post_training/integration/grpo_trainer_correctness_test.py
+++ b/tests/post_training/integration/grpo_trainer_correctness_test.py
@@ -25,11 +25,9 @@
   pytest tests/post_training/integration/grpo_trainer_correctness_test.py
 """
 
-import functools
 import os
 import sys
 import unittest
-from flax import linen as nn
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
@@ -37,17 +35,14 @@
 import maxtext as mt
 from maxtext.configs import pyconfig, types
 from maxtext.utils.globals import 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, grpo_loss_fn_nnx, setup_train_loop
-from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
+from maxtext.experimental.rl.grpo_trainer import grpo_loss_fn_nnx, setup_train_loop
+from maxtext.experimental.rl.grpo_utils import compute_log_probs_nnx
 from tests.utils.test_helpers import ensure_tokenizer_downloaded
 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
@@ -73,49 +68,28 @@ def get_golden_data(config):
 def setup_maxtext_model(config, mesh):
   """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.
+  Returns (model, state, reference, init_rng, state_mesh_shardings, data_sharding). The
+  model carries its own params (state is None) and reference is a cloned frozen model.
   """
   init_rng = jax.random.PRNGKey(config.init_weights_seed)
   data_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec(None))
 
-  if config.pure_nnx:
-    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)
-  reference_params = jax.tree.map(jnp.copy, state.params["params"])
-  state = _merge_grpo_state(state, reference_params)
-  return (
-      maxtext_model,
-      state,
-      reference_params,
-      init_rng,
-      state_mesh_shardings,
-      data_sharding,
-  )
+  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
 
 
 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)
+  """Per-token log-probs from the NNX model."""
+  del state, rngs  # unused on the NNX path
+  return compute_log_probs_nnx(model, ids, pos, seg, comp_seg, config, is_train=False)
 
 
 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)
+  """GRPO loss with the NNX policy and frozen reference model."""
+  del state  # unused on the NNX path
+  return grpo_loss_fn_nnx(model, config, data, rng, None, reference)
 
 
 def prepare_maxtext_inputs(input_str, tokenizer_model):
diff --git a/tests/post_training/integration/sft_trainer_correctness_test.py b/tests/post_training/integration/sft_trainer_correctness_test.py
index 253ba7ade9..edcfba2fe5 100644
--- a/tests/post_training/integration/sft_trainer_correctness_test.py
+++ b/tests/post_training/integration/sft_trainer_correctness_test.py
@@ -24,7 +24,6 @@
   pytest tests/post_training/integration/sft_trainer_correctness_test.py
 """
 
-import functools
 import os.path
 import sys
 import unittest
@@ -34,11 +33,8 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 import jsonlines
-from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.configs import pyconfig
 from maxtext.input_pipeline import input_pipeline_utils
-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
@@ -118,40 +114,23 @@ def setup_maxtext_model(config):
   init_rng = jax.random.PRNGKey(config.init_weights_seed)
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = Mesh(devices_array, config.mesh_axes)
-  if config.pure_nnx:
-    # 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
+  # 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
 
 
 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"],
-      maxtext_data["inputs_position"],
+  maxtext_model, _, _ = setup_maxtext_model(config)
+  # 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,
-      rngs=rng,
-      mutable="intermediates",
   )
-  return maxtext_logits
 
 
 def get_token_log_probs(logits, inputs):
diff --git a/tests/unit/grpo_nnx_test.py b/tests/unit/grpo_nnx_test.py
index ddc2de5da0..0ccf30f7a0 100644
--- a/tests/unit/grpo_nnx_test.py
+++ b/tests/unit/grpo_nnx_test.py
@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for `grpo_loss_fn_nnx`, `compute_log_probs_nnx`, plus a small
-Linen-path regression block (the repo's existing Linen GRPO integration test
-is TPU-only)."""
+"""Unit tests for `grpo_loss_fn_nnx` and `compute_log_probs_nnx`."""
 
 import types
 import unittest
 
-import flax.linen as nn
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -95,7 +92,7 @@ def setUp(self):
     self.data = _make_grpo_batch()
 
   def test_aux_structure_matches_linen(self):
-    """`grpo_loss_fn_nnx` returns the same `LossAux` dataclass shape as `grpo_loss_fn`."""
+    """`grpo_loss_fn_nnx` returns a `LossAux` dataclass with the expected fields."""
     loss, aux = grpo_trainer.grpo_loss_fn_nnx(
         self.policy, self.config, self.data, None, None, self.reference, is_train=True
     )
@@ -295,70 +292,5 @@ def test_host_offload_matches_no_offload(self):
       np.testing.assert_allclose(np.asarray(a), np.asarray(b), rtol=1e-6, atol=1e-6)
 
 
-# ---------------------------------------------------------------------------
-# Linen-path regression smoke tests
-# ---------------------------------------------------------------------------
-
-
-class _MockLinenTransformer(nn.Module):
-  """Tiny Linen module with the `model.apply(...)` signature that Linen `compute_log_probs` expects."""
-
-  vocab_size: int
-  embed_dim: int
-
-  @nn.compact
-  def __call__(self, inputs, positions, decoder_segment_ids=None, enable_dropout=False):
-    del positions, decoder_segment_ids, enable_dropout
-    embed = nn.Embed(num_embeddings=self.vocab_size, features=self.embed_dim, name="embed")(inputs)
-    return nn.Dense(features=self.vocab_size, name="proj")(embed)
-
-
-class TestLinenGrpoRegression(unittest.TestCase):
-  """Smoke tests that the Linen `grpo_loss_fn` and `compute_log_probs` still run on Linen-shaped inputs."""
-
-  def setUp(self):
-    self.config = _make_grpo_config()
-    self.config.pure_nnx = False  # Force the Linen dispatch branch.
-    self.config.gradient_accumulation_steps = 1
-    self.data = _make_grpo_batch()
-    self.model = _MockLinenTransformer(vocab_size=8, embed_dim=4)
-    rng = jax.random.key(0)
-    inputs = self.data["prompt_completions"]
-    self.params = self.model.init(rng, inputs, inputs, decoder_segment_ids=jnp.ones_like(inputs), enable_dropout=False)
-    self.reference_params = jax.tree_util.tree_map(jnp.copy, self.params)
-
-  def test_linen_grpo_loss_fn_still_runs(self):
-    """Linen `grpo_loss_fn` returns a finite loss + a `LossAux`."""
-    loss, aux = grpo_trainer.grpo_loss_fn(
-        self.model,
-        self.config,
-        self.data,
-        jax.random.key(1),
-        self.params,
-        self.reference_params["params"],  # On Linen, reference_params is the inner subtree.
-        is_train=True,
-    )
-    self.assertTrue(jnp.isfinite(loss))
-    self.assertTrue(hasattr(aux, "total_loss"))
-    self.assertTrue(hasattr(aux, "moe_lb_loss"))
-    self.assertTrue(hasattr(aux, "total_weights"))
-
-  def test_linen_compute_log_probs_still_runs(self):
-    """Linen `compute_log_probs` produces shape `[B, S-1]`."""
-    log_probs, _ = grpo_utils.compute_log_probs(
-        self.model,
-        self.params,
-        self.data["prompt_completions"],
-        self.data["prompt_completions_position"],
-        self.data["prompt_completions_segmentation"],
-        self.data["ar_completions_segmentation"],
-        self.config,
-        is_train=False,
-        rngs={"dropout": jax.random.key(2), "params": jax.random.key(3)},
-    )
-    S = self.data["prompt_completions"].shape[1]
-    self.assertEqual(log_probs.shape, (self.data["prompt_completions"].shape[0], S - 1))
-
-
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py
index c58349ac98..977aaef641 100644
--- a/tests/unit/pre_train_loss_mask_test.py
+++ b/tests/unit/pre_train_loss_mask_test.py
@@ -20,7 +20,6 @@
 import unittest
 from unittest import mock
 
-from flax import linen as nn
 from flax import nnx
 import jax
 import jax.numpy as jnp
@@ -84,29 +83,6 @@ def __call__(
     return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32)
 
 
-class _UniformLinenDecoder(nn.Module):
-  """Returns uniform logits through the Linen call contract."""
-
-  vocab_size: int
-  mesh: object
-
-  @nn.compact
-  def __call__(
-      self,
-      decoder_input_tokens,
-      decoder_positions,
-      decoder_segment_ids=None,
-      encoder_images=None,
-      encoder_image_masks=None,
-      enable_dropout=False,
-      decoder_target_tokens=None,
-      decoder_target_mask=None,
-  ):
-    del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks
-    del enable_dropout, decoder_target_tokens, decoder_target_mask
-    return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32)
-
-
 def _make_data(include_loss_mask=True):
   """Builds a batch whose explicit loss mask differs from segmentation."""
   data = {
@@ -123,7 +99,7 @@ def _make_data(include_loss_mask=True):
 
 
 class PreTrainLossMaskTest(unittest.TestCase):
-  """Checks explicit masks in both Linen and NNX loss branches."""
+  """Checks explicit target-loss masking in the pre-training loss."""
 
   def setUp(self):
     super().setUp()
@@ -142,20 +118,6 @@ def _use_block_diffusion(self):
     self.config.attention_type = "block_diffusion"
     self.config.training_objective = "block_diffusion"
 
-  def _linen_model_and_variables(self, data):
-    """Initializes the test Linen decoder for the supplied batch."""
-    mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context"))
-    model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh)
-    variables = model.init(
-        jax.random.key(0),
-        data["inputs"],
-        data["inputs_position"],
-        decoder_segment_ids=data["inputs_segmentation"],
-        decoder_target_tokens=data["targets"],
-        decoder_target_mask=data["targets_segmentation"],
-    )
-    return model, variables
-
   def _assert_explicit_mask_result(self, loss, aux):
     expected_mask = _make_data()["targets_loss_mask"] != 0
     expected_xent = jnp.sum(self.per_token_xent * expected_mask)
@@ -173,27 +135,6 @@ def test_nnx_loss_uses_targets_loss_mask(self):
 
     self._assert_explicit_mask_result(loss, aux)
 
-  def test_linen_loss_uses_targets_loss_mask(self):
-    self._use_block_diffusion()
-    data = _make_data()
-    model, variables = self._linen_model_and_variables(data)
-    with self._cross_entropy_patch():
-      loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True)
-
-    self._assert_explicit_mask_result(loss, aux)
-
-  def test_linen_causal_loss_skips_diffusion_alignment(self):
-    data = _make_data(include_loss_mask=False)
-    model, variables = self._linen_model_and_variables(data)
-    with self._cross_entropy_patch():
-      loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True)
-
-    expected_mask = data["targets_segmentation"] != 0
-    expected_xent = jnp.sum(self.per_token_xent * expected_mask)
-    self.assertEqual(int(aux["total_weights"]), 7)
-    self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent))
-    self.assertAlmostEqual(float(loss), float(expected_xent / 7.0))
-
   def test_causal_fallback_uses_targets_segmentation(self):
     data = _make_data(include_loss_mask=False)
     model = _UniformNnxDecoder(self.config.vocab_size)

From 5890a0549e6c2820000b5de1d0672eb3de9e548f Mon Sep 17 00:00:00 2001
From: Lance Wang <lance.wang@cienet.com>
Date: Thu, 2 Jul 2026 17:00:00 +0000
Subject: [PATCH 5/5] [NNX] Delete Linen (3/5): collapse dispatch in inference
 (maxengine, kvcache, vLLM, LoRA)

---
 src/maxtext/inference/maxengine/maxengine.py | 588 +++++--------------
 src/maxtext/inference/vllm_decode.py         |   4 -
 src/maxtext/utils/lora_utils.py              |  48 +-
 tests/integration/maxengine_test.py          |  10 +-
 tests/post_training/unit/lora_utils_test.py  |   2 -
 tests/unit/aqt_serve_roundtrip_nnx_test.py   |   3 -
 tests/unit/maxengine_nnx_test.py             |   3 -
 7 files changed, 166 insertions(+), 492 deletions(-)

diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py
index 62c5fe0345..4785b8bd92 100644
--- a/src/maxtext/inference/maxengine/maxengine.py
+++ b/src/maxtext/inference/maxengine/maxengine.py
@@ -31,18 +31,14 @@
 else:
   from jax.experimental.layout import DeviceLocalLayout as DLL  # type: ignore # pylint: disable=no-name-in-module
 
-from flax import linen as nn
 from flax import nnx
 from flax import struct
 from flax.linen import partitioning as nn_partitioning
-import flax
 
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
-from maxtext.models import models
 from maxtext.layers import quantizations
 from maxtext.inference import inference_utils
-from maxtext.multimodal import processor as mm_processor
 from maxtext.utils import lora_utils
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
@@ -127,45 +123,39 @@ def __init__(self, config: Any, devices: Any | None = None):
 
     # Model and Optimizer definition.
     quant = quantizations.configure_quantization(config)
-    if config.pure_nnx:
-      # `serve` only when the on-disk checkpoint already carries `qrhs.frozen`
-      # (no full-precision kernel). For `checkpoint_is_quantized=False` with
-      # quant enabled we stay in `train` mode and let AQT quantize per-forward
-      # against the full-precision kernel — same numerical result as `serve`
-      # for absmax calibration, just slower.
-      nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train"
-      # We need both PREFILL and AR abstract models because the cache vars inherit
-      # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and
-      # bulk_insert searches for the substring "cache_batch" in the AR-mode names.
-      # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids
-      # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm".
-      _create_model = model_creation_utils.get_nnx_create_model_fn(
-          config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str
-      )
-      _create_model_ar = model_creation_utils.get_nnx_create_model_fn(
-          config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str
-      )
-      self._nnx_quant_mode_str = nnx_quant_mode_str
-      with nn_partitioning.axis_rules(config.logical_axis_rules):
-        abstract_model = nnx.eval_shape(_create_model)
-        abstract_model_ar = nnx.eval_shape(_create_model_ar)
-      self.model = abstract_model
-      self.model_ar = abstract_model_ar
-      # 3-way split so JIT bodies can pass (params, cache, rest) separately to
-      # nnx.merge. `rest` (RNG state etc.) is materialized in load_params.
-      graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...)
-      self.graphdef = graphdef
-      # Layers may bake their construction-time model_mode into static attributes,
-      # so a call must be merged with the graphdef built for that same mode.
-      graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...)
-      self.graphdef_ar = graphdef_ar
-      self._create_model_fn = _create_model
-      self._nnx_rest_state = None
-    else:
-      self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL)
-      self.graphdef = None
-      self.graphdef_ar = None
-      self._create_model_fn = None
+    # `serve` only when the on-disk checkpoint already carries `qrhs.frozen`
+    # (no full-precision kernel). For `checkpoint_is_quantized=False` with
+    # quant enabled we stay in `train` mode and let AQT quantize per-forward
+    # against the full-precision kernel — same numerical result as `serve`
+    # for absmax calibration, just slower.
+    nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train"
+    # We need both PREFILL and AR abstract models because the cache vars inherit
+    # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and
+    # bulk_insert searches for the substring "cache_batch" in the AR-mode names.
+    # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids
+    # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm".
+    _create_model = model_creation_utils.get_nnx_create_model_fn(
+        config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str
+    )
+    _create_model_ar = model_creation_utils.get_nnx_create_model_fn(
+        config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str
+    )
+    self._nnx_quant_mode_str = nnx_quant_mode_str
+    with nn_partitioning.axis_rules(config.logical_axis_rules):
+      abstract_model = nnx.eval_shape(_create_model)
+      abstract_model_ar = nnx.eval_shape(_create_model_ar)
+    self.model = abstract_model
+    self.model_ar = abstract_model_ar
+    # 3-way split so JIT bodies can pass (params, cache, rest) separately to
+    # nnx.merge. `rest` (RNG state etc.) is materialized in load_params.
+    graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...)
+    self.graphdef = graphdef
+    # Layers may bake their construction-time model_mode into static attributes,
+    # so a call must be merged with the graphdef built for that same mode.
+    graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...)
+    self.graphdef_ar = graphdef_ar
+    self._create_model_fn = _create_model
+    self._nnx_rest_state = None
     self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None))
 
     self.abstract_params = None
@@ -351,65 +341,7 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar
     if rng is None:
       rng = jax.random.PRNGKey(0)
 
-    if self.config.pure_nnx:
-      return self._load_params_nnx(params=params, rng=rng)
-
-    if self.model.quant and self.config.checkpoint_is_quantized:
-      print("Loading from the quantized checkpoint...")
-      self.model.quant.quant_mode = quantizations.get_quant_mode("serve")
-
-    rng1, rng2, rng3 = jax.random.split(rng, 3)
-    if params:
-      print("Resharding given params")
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng)
-      _, self.state_mesh_annotations, state_mesh_shardings = maxtext_utils.get_abstract_state(
-          self.config, self._mesh, init_state_fn, False
-      )
-      # reshard given params based on shardings from config in MaxEngine
-      params = jax.device_put(params, state_mesh_shardings.params)
-      state = maxtext_utils.init_decode_state(None, params)
-      state = max_utils.unbox_logicallypartioned(state)
-    else:
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng1)
-      state, self.state_mesh_annotations = maxtext_utils.setup_decode_state(self.config, self._mesh, None, init_state_fn)
-    # pylint: disable=isinstance-second-argument-not-valid-type
-    self.abstract_params = jax.tree_util.tree_map(
-        lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding)
-        if isinstance(x, jax.Array)
-        else None,
-        state.params,
-    )
-
-    self.prefill_kv_cache_annotations = maxtext_utils.get_prefill_kv_cache_annotations(
-        self.model, self.config, rng2, self._mesh
-    )
-    self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
-        lambda x: jax.sharding.NamedSharding(self._mesh, x),
-        self.prefill_kv_cache_annotations,
-    )
-
-    if self.config.stack_prefill_result_cache:
-      # Add extra axis for the axis generated by the stack.
-      self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
-          lambda x: jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(None, *x.spec)),
-          self.prefill_kv_cache_shardings,
-      )
-      self.prefill_kv_cache_shardings = self.prefill_kv_cache_shardings["decoder"]["layers_0"]
-
-    self.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(self.model, self.config, rng2, self._mesh)
-    self.kv_cache_shardings = jax.tree_util.tree_map(
-        lambda x: jax.sharding.NamedSharding(self._mesh, x),
-        self.kv_cache_annotations,
-    )
-
-    if self.model.quant and not self.config.checkpoint_is_quantized:
-      params = self.quantize_params(state, rng3)
-    else:
-      params = state.params
-
-    self.print_stats("After load_params")
-
-    return params
+    return self._load_params_nnx(params=params, rng=rng)
 
   def _load_params_nnx(self, params, rng):
     """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings.
@@ -558,142 +490,77 @@ def apply_adapter(self, base_params, adapter_config, adapter_params):
 
     lora_rank = int(adapter_config["r"])
     lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank
-    if self.config.pure_nnx:
-      lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor)
-    else:
-      lora_utils.apply_lora_on_base_params(base_params, adapter_params, lora_scale_factor)
+    lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor)
 
   def unapply_adapter(self, base_params, adapter_config, adapter_params):
     """Unapply the adapter params from the merged params to get back the base params."""
 
     lora_rank = int(adapter_config["r"])
     lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank
-    if self.config.pure_nnx:
-      lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor)
-    else:
-      lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor)
+    lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor)
 
   def quantize_params(self, state, rng: PRNGKeyType | None = None):
     """Forward pass to quantize decode params."""
     if rng is None:
       rng = jax.random.PRNGKey(0)
-    if self.config.pure_nnx:
-      # NNX takes a different code path: convert-on-load lives in `_load_params_nnx`
-      # via `_convert_and_quantize_nnx`, which runs the dummy forward against a
-      # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model.
-      # The standalone `quantize_params(state, rng)` API expects a Linen-shape
-      # `state.params` dict and isn't reachable on the NNX pathway in maxengine
-      # (load_params already dispatched to _load_params_nnx).
-      raise NotImplementedError(
-          "Use load_params() on NNX — the convert step runs inside _load_params_nnx via "
-          "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API."
-      )
-
-    self.model.quant.quant_mode = quantizations.get_quant_mode("convert")
-
-    @jax.jit
-    def model_apply(_p, _rng):
-      image_shape = mm_processor.get_dummy_image_shape_for_init(
-          model_name=self.config.model_name,
-          batch_size=self.config.micro_batch_size_to_train_on,
-      )
-      audio_shape = mm_processor.get_dummy_audio_shape_for_init(self.config)
-      return self.model.apply(
-          _p | {"aqt": {}},
-          jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None,
-          # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input.
-          encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32)
-          if self.config.use_multimodal and "llama4" in self.config.model_name
-          else None,
-          encoder_audios=jnp.ones(audio_shape, dtype=jnp.float32) if self.config.use_audio else None,
-          decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          enable_dropout=False,
-          model_mode=MODEL_MODE_PREFILL,
-          rngs={"params": _rng},
-          mutable=True,
-      )
-
-    _, new_vars = model_apply(state.params, rng)
-    # Remove param values which have corresponding qtensors in aqt to save memory.
-    params = {}
-    params["aqt"] = new_vars["aqt"]
-    params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"])
-    self.abstract_params = jax.tree_util.tree_map(
-        lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding),
-        params,
+    # NNX takes a different code path: convert-on-load lives in `_load_params_nnx`
+    # via `_convert_and_quantize_nnx`, which runs the dummy forward against a
+    # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model.
+    # The standalone `quantize_params(state, rng)` API expects a Linen-shape
+    # `state.params` dict and isn't reachable on the NNX pathway in maxengine
+    # (load_params already dispatched to _load_params_nnx).
+    raise NotImplementedError(
+        "Use load_params() on NNX — the convert step runs inside _load_params_nnx via "
+        "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API."
     )
-    maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params)
-    self.model.quant.quant_mode = quantizations.get_quant_mode("serve")
-    return params
 
   def _maybe_stack_prefill_result_cache(self, cache):
     """Stack the caches across the layers."""
     if not self.config.stack_prefill_result_cache:
       return cache
 
-    if self.config.pure_nnx:
-      if self.config.scan_layers:
-        # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack.
-        return cache
-      # scan_layers=False: stack the per-layer subtrees under decoder into one
-      # subtree with a leading layer axis (matching the scan_layers=True shape).
-      if "dense_layers_0" in cache["decoder"] or "moe_layers_0" in cache["decoder"]:
-        first_dense = self.config.first_num_dense_layers
-        num_moe = self.config.num_decoder_layers - first_dense
-        layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
-      else:
-        layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
-
-      layer_cache = [cache["decoder"][key] for key in layer_keys]
-      stacked = jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
-      return {"decoder": {"layers": stacked}}
-
-    layer_keys = []
-    for i in range(self.config.num_decoder_layers):
-      layer_keys.append(f"layers_{i}")
-
-    layer_cache = [cache["decoder"][layer_key] for layer_key in layer_keys]
+    if self.config.scan_layers:
+      # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack.
+      return cache
+    # scan_layers=False: stack the per-layer subtrees under decoder into one
+    # subtree with a leading layer axis (matching the scan_layers=True shape).
+    if "dense_layers_0" in cache["decoder"] or "moe_layers_0" in cache["decoder"]:
+      first_dense = self.config.first_num_dense_layers
+      num_moe = self.config.num_decoder_layers - first_dense
+      layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
+    else:
+      layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
 
-    return jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
+    layer_cache = [cache["decoder"][key] for key in layer_keys]
+    stacked = jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
+    return {"decoder": {"layers": stacked}}
 
   def _maybe_unstack_prefill_result_cache(self, cache):
     """Unstack the caches across the layers."""
     if not self.config.stack_prefill_result_cache:
       return cache
 
-    if self.config.pure_nnx:
-      if self.config.scan_layers:
-        # Mirror _maybe_stack_prefill_result_cache: the cache already carries the
-        # layer axis, so there is nothing to unstack.
-        return cache
-      # scan_layers=False: split the leading layer axis back into per-layer subtrees.
-      stacked = cache["decoder"]["layers"]
-      res_cache = {"decoder": {}}
-      is_deepseek = (
-          getattr(self.model, "is_deepseek", False)
-          or (hasattr(self.model, "decoder") and getattr(self.model.decoder, "is_deepseek", False))
-          or (hasattr(self.config, "decoder_block") and str(self.config.decoder_block).lower() == "deepseek")
-      )
-      if is_deepseek:
-        first_dense = self.config.first_num_dense_layers
-        num_moe = self.config.num_decoder_layers - first_dense
-        layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
-      else:
-        layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
-
-      for idx, key in enumerate(layer_keys):
-        res_cache["decoder"][key] = jax.tree.map(lambda x, i=idx: x[i], stacked)
-      return res_cache
-
-    flat_cache, treedef = jax.tree.flatten(cache)
-    layer_cache = [jax.tree.unflatten(treedef, flat_cache_vars) for flat_cache_vars in zip(*flat_cache, strict=True)]
+    if self.config.scan_layers:
+      # Mirror _maybe_stack_prefill_result_cache: the cache already carries the
+      # layer axis, so there is nothing to unstack.
+      return cache
+    # scan_layers=False: split the leading layer axis back into per-layer subtrees.
+    stacked = cache["decoder"]["layers"]
     res_cache = {"decoder": {}}
+    is_deepseek = (
+        getattr(self.model, "is_deepseek", False)
+        or (hasattr(self.model, "decoder") and getattr(self.model.decoder, "is_deepseek", False))
+        or (hasattr(self.config, "decoder_block") and str(self.config.decoder_block).lower() == "deepseek")
+    )
+    if is_deepseek:
+      first_dense = self.config.first_num_dense_layers
+      num_moe = self.config.num_decoder_layers - first_dense
+      layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
+    else:
+      layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
 
-    for i in range(self.config.num_decoder_layers):
-      res_cache["decoder"][f"layers_{i}"] = layer_cache[i]
-
+    for idx, key in enumerate(layer_keys):
+      res_cache["decoder"][key] = jax.tree.map(lambda x, i=idx: x[i], stacked)
     return res_cache
 
   def prefill_aot(  # pylint: disable=too-many-positional-arguments
@@ -786,10 +653,7 @@ def _prefill_jit(
     if existing_prefix is not None:
       if not self.use_chunked_prefill:
         raise ValueError("Using chunked prefill is needed for existing_prefix.")
-      # NNX threads existing_prefix.cache via the nnx_cache local below; only
-      # the Linen path merges cache into input_params (params is a dict there).
-      if not self.config.pure_nnx:
-        input_params = params | {"cache": existing_prefix.cache}
+      # NNX threads existing_prefix.cache via the nnx_cache local below.
       start_position = existing_prefix.common_prefix_tokens.shape[0]
       # TODO(yuyanpeng): rename previous_chunk
       previous_chunk = jnp.expand_dims(existing_prefix.common_prefix_tokens, 0)
@@ -821,52 +685,30 @@ def _prefill_jit(
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
-    if self.config.pure_nnx:
-      # Prefill always operates on batch=1 (one padded prompt at a time).
-      nnx_cache = (
-          existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL)
+    # Prefill always operates on batch=1 (one padded prompt at a time).
+    nnx_cache = (
+        existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL)
+    )
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=input_params,
+          cache_dict=nnx_cache,
+          decoder_input_tokens=input_tokens,
+          decoder_positions=positions,
+          decoder_segment_ids=sequence_indicator,
+          encoder_images=images,
+          encoder_image_masks=image_masks,
+          encoder_videos=videos,
+          encoder_video_masks=video_masks,
+          encoder_video_grid_thw=video_grid_thw,
+          encoder_audios=audio_values,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+          previous_chunk=previous_chunk,
+          true_length=true_length,
+          slot=slot,
       )
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=input_params,
-            cache_dict=nnx_cache,
-            decoder_input_tokens=input_tokens,
-            decoder_positions=positions,
-            decoder_segment_ids=sequence_indicator,
-            encoder_images=images,
-            encoder_image_masks=image_masks,
-            encoder_videos=videos,
-            encoder_video_masks=video_masks,
-            encoder_video_grid_thw=video_grid_thw,
-            encoder_audios=audio_values,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            previous_chunk=previous_chunk,
-            true_length=true_length,
-            slot=slot,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            input_params,
-            input_tokens,
-            positions,
-            encoder_images=images,
-            encoder_image_masks=image_masks,
-            encoder_videos=videos,
-            encoder_video_masks=video_masks,
-            encoder_video_grid_thw=video_grid_thw,
-            encoder_audios=audio_values,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-            previous_chunk=previous_chunk,
-            true_length=true_length,
-            slot=slot,
-        )
+    new_vars = {"cache": new_cache_dict}
     if return_prompt_logp:
       prompt_logp = inference_utils.prompt_logprobs_from_prefill(flat_logits, input_tokens, true_length)
     else:
@@ -1089,32 +931,19 @@ def _prefill_multisampling_jit(
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
-    if self.config.pure_nnx:
-      # Prefill is batch=1 (one prompt); multi-sampling only draws several first
-      # tokens from the shared logits below. Mirror the _prefill_jit NNX branch.
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
-            decoder_input_tokens=input_tokens,
-            decoder_positions=positions,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            params,
-            input_tokens,
-            positions,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    # Prefill is batch=1 (one prompt); multi-sampling only draws several first
+    # tokens from the shared logits below. Mirror the _prefill_jit NNX branch.
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
+          decoder_input_tokens=input_tokens,
+          decoder_positions=positions,
+          decoder_segment_ids=sequence_indicator,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+      )
+    new_vars = {"cache": new_cache_dict}
 
     next_pos = jnp.full((1, 1), true_length, dtype=jnp.int32)
     selected_logits = jax.lax.dynamic_slice(
@@ -1225,33 +1054,20 @@ def prefill_concat(
     input_tokens = jnp.expand_dims(padded_tokens, 0)  # [BATCH, SEQUENCE]
     decoder_positions = jnp.expand_dims(decoder_positions, 0)
     decoder_segment_ids = jnp.expand_dims(decoder_segment_ids, 0)
-    rng, new_rng = jax.random.split(rng)
-    if self.config.pure_nnx:
-      # Packed prompts run as a single batch=1 prefill; the packed positions and
-      # segment ids keep the prompts separated. Mirror the _prefill_jit NNX branch.
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
-            decoder_input_tokens=input_tokens,
-            decoder_positions=decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            params,
-            input_tokens,
-            decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    rng, _ = jax.random.split(rng)
+    # Packed prompts run as a single batch=1 prefill; the packed positions and
+    # segment ids keep the prompts separated. Mirror the _prefill_jit NNX branch.
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
+          decoder_input_tokens=input_tokens,
+          decoder_positions=decoder_positions,
+          decoder_segment_ids=decoder_segment_ids,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+      )
+    new_vars = {"cache": new_cache_dict}
     cache = new_vars["cache"]
     cache = self._maybe_stack_prefill_result_cache(cache)
     if return_prompt_logp:
@@ -1399,28 +1215,16 @@ def _generate_jit(
     previous_token = decode_state["tokens"]
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
     # run one step generation
-    if self.config.pure_nnx:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        out_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=decode_state["cache"],
-            decoder_input_tokens=previous_token,
-            decoder_positions=decode_state["next_pos"],
-            enable_dropout=False,
-            model_mode=MODEL_MODE_AUTOREGRESSIVE,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        out_logits, new_vars = self.model.apply(
-            params | {"cache": decode_state["cache"]},
-            previous_token,
-            decode_state["next_pos"],
-            enable_dropout=False,
-            model_mode=MODEL_MODE_AUTOREGRESSIVE,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      out_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=decode_state["cache"],
+          decoder_input_tokens=previous_token,
+          decoder_positions=decode_state["next_pos"],
+          enable_dropout=False,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+      )
+    new_vars = {"cache": new_cache_dict}
     out_logits = jax.lax.with_sharding_constraint(out_logits, self.replicated_sharding)
     new_cache = jax.lax.with_sharding_constraint(new_vars["cache"], self.kv_cache_shardings)
     # sampling tokens
@@ -1937,107 +1741,7 @@ def init_decode_state(
       **kwargs,  # pylint: disable=unused-argument
   ) -> DecodeState:
     """Initialises any state which a generation step transforms."""
-    if rng is None:
-      rng = jax.random.PRNGKey(0)
-
-    if self.config.pure_nnx:
-      return self._init_decode_state_nnx()
-
-    # pylint: disable=unused-argument
-    def init(abstract_params):
-      x = jnp.ones(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      dummy_image = jnp.ones(
-          mm_processor.get_dummy_image_shape_for_init(
-              model_name=self.config.model_name, batch_size=int(self.config.per_device_batch_size * self.mesh.size)
-          ),
-          dtype=jnp.int32,
-      )
-      dummy_audio = jnp.ones(
-          mm_processor.get_dummy_audio_shape_for_init(self.config),
-          dtype=jnp.float32,
-      )
-      _, cache = self.model.apply(
-          abstract_params,
-          x,
-          x,
-          encoder_images=dummy_image if self.config.use_multimodal else None,
-          encoder_audios=dummy_audio if self.config.use_audio else None,
-          enable_dropout=False,
-          model_mode=MODEL_MODE_AUTOREGRESSIVE,
-          rngs={"params": rng},
-          mutable=["cache"],
-          slot=0,
-      )
-
-      next_pos = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      generated_tokens = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      tokens = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      token_logp = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.float32,
-      )
-      return {
-          "logits": jnp.zeros(
-              (
-                  int(self.config.per_device_batch_size * self.mesh.size),
-                  1,
-                  self.config.vocab_size,
-              )
-          ),
-          "cache": cache["cache"],
-          "next_pos": next_pos,
-          "generated_tokens": generated_tokens,
-          "tokens": tokens,
-          "token_logp": token_logp,
-      }
-
-    with nn_partitioning.axis_rules(self.config.logical_axis_rules):
-      abstract_outputs = jax.eval_shape(init, self.abstract_params)
-    logical_annotations = nn.get_partition_spec(abstract_outputs)
-
-    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-      mesh_annotations = nn.logical_to_mesh(logical_annotations)
-
-    shardings = jax.tree_util.tree_map(
-        lambda mesh_annotation: jax.sharding.NamedSharding(self._mesh, mesh_annotation),
-        mesh_annotations,
-    )
-
-    if self._compiled_initialize_fn is None:
-
-      @functools.partial(jax.jit, out_shardings=shardings)
-      def initialize():
-        return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs)
-
-      self._compiled_initialize_fn = initialize
-
-    init_state = self._compiled_initialize_fn()
-    cache = init_state["cache"]
-
-    def is_lp(k):
-      return isinstance(k, flax.linen.spmd.LogicallyPartitioned)
-
-    self.kv_cache_annotations_named = jax.tree_util.tree_map(
-        lambda x: tuple(x.logical_axes)
-        if hasattr(x, "logical_axes")
-        else (tuple(x.names) if hasattr(x, "names") else ()),
-        cache,
-        is_leaf=is_lp,
-    )
-    zeroed = max_utils.unbox_logicallypartioned(init_state)
-    return zeroed
+    return self._init_decode_state_nnx()
 
   def _init_decode_state_nnx(self) -> DecodeState:
     """NNX equivalent of init_decode_state. Returns a decode_state dict with a pure-dict cache."""
@@ -2130,18 +1834,10 @@ def set_engine_vars_from_base_engine(
   """Set internal vars from base_engine, which has already loaded the checkpoint and has sharding,
   mesh, and kv cache related vars set.
   """
-  if not engine.config.pure_nnx and base_engine.model.quant:
-    # NNX bakes the quant mode in at construction (via _nnx_quant_mode_str) rather
-    # than mutating model.quant.quant_mode, so there's nothing to copy on that path.
-    engine.model.quant.quant_mode = base_engine.model.quant.quant_mode
   engine.state_mesh_annotations = base_engine.state_mesh_annotations
   engine.abstract_params = base_engine.abstract_params
-  if engine.config.pure_nnx:
-    # Linen's get_kv_cache_annotations calls model.init(); NNX modules have no
-    # .init, so use the abstract-model variant (mirrors _load_params_nnx).
-    engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh)
-  else:
-    engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(engine.model, engine.config, rng, engine.mesh)  # pylint: disable=protected-access
+  # NNX modules have no .init, so use the abstract-model variant (mirrors _load_params_nnx).
+  engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh)
   engine.kv_cache_shardings = jax.tree_util.tree_map(
       lambda x: jax.sharding.NamedSharding(engine.mesh, x),
       engine.kv_cache_annotations,  # pylint: disable=protected-access
diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py
index 03514ce9cf..ea6d386bdc 100644
--- a/src/maxtext/inference/vllm_decode.py
+++ b/src/maxtext/inference/vllm_decode.py
@@ -99,8 +99,6 @@ def decode_with_vllm(config: Config) -> None:
               "debug_sharding": config.debug_sharding,
               "prefuse_moe_weights": config.prefuse_moe_weights,
               "scan_layers": config.scan_layers,
-              "enable_nnx": config.enable_nnx,
-              "pure_nnx_decoder": config.pure_nnx_decoder,
           },
           "sharding": {
               "sharding_strategy": {
@@ -248,8 +246,6 @@ def decode_with_tunix(
           "debug_sharding": config.debug_sharding,
           "prefuse_moe_weights": config.prefuse_moe_weights,
           "scan_layers": config.scan_layers,
-          "enable_nnx": config.enable_nnx,
-          "pure_nnx_decoder": config.pure_nnx_decoder,
       }
   }
 
diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 08de8ed8e0..f304873122 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -15,7 +15,6 @@
 """Common LoRA utils needed to support LoRA adapters."""
 
 from collections.abc import Mapping
-from functools import partial
 import json
 import os
 import re
@@ -171,10 +170,7 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte
     if not gcs_utils.gcs_path_exists(f"{adapter_weights_path}/commit_success.txt"):
       raise FileNotFoundError(f"Failed to read lora_weights from {adapter_weights_path}.")
 
-    if config.pure_nnx:
-      lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
-    else:
-      lora_state, _ = get_lora_abstract_state(base_abstract_state_params, lora_config)
+    lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
 
     with nn_partitioning.axis_rules(config.logical_axis_rules):
       lora_params = checkpointing.load_params_from_path(
@@ -218,37 +214,31 @@ def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkp
 
   if lora_adapter_path:
     max_logging.log(f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}")
-    if config.pure_nnx:
-      # pylint: disable=import-outside-toplevel
-      from maxtext.common import train_state_nnx
-      from maxtext.utils import model_creation_utils
-
-      _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
-
-      def create_train_state_fn():
-        nnx_model = _create_model_partial()
-        wrt = (
-            getattr(nnx, "LoRAParam", nnx.Param)
-            if getattr(getattr(config, "lora", None), "enable_lora", False)
-            else nnx.Param
-        )
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    # pylint: disable=import-outside-toplevel
+    from maxtext.common import train_state_nnx
+    from maxtext.utils import model_creation_utils
+
+    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
+
+    def create_train_state_fn():
+      nnx_model = _create_model_partial()
+      wrt = (
+          getattr(nnx, "LoRAParam", nnx.Param)
+          if getattr(getattr(config, "lora", None), "enable_lora", False)
+          else nnx.Param
+      )
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
+    init_state_fn = create_train_state_fn
     unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True)
 
     lora_config_path = lora_adapter_path + "adapter_config.json"
 
     lora_config = gcs_utils.read_json_from_gcs(lora_config_path)
 
-    if config.pure_nnx:
-      base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
-      lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
-    else:
-      lora_state, lora_state_annotations = get_lora_abstract_state(unboxed_abstract_state.params, lora_config)
+    base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
+    lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
 
     lora_weights_path = f"{lora_adapter_path}/0/items"
 
diff --git a/tests/integration/maxengine_test.py b/tests/integration/maxengine_test.py
index f17a523d52..e49bc8b7a9 100644
--- a/tests/integration/maxengine_test.py
+++ b/tests/integration/maxengine_test.py
@@ -96,8 +96,8 @@ def test_stack_and_unstack_prefill_cache_nnx(self):
   # default. test_basic_prefill_nnx / test_basic_decode_nnx below cover the NNX path.
 
   def _init_nnx_pyconfig(self, **kwargs):
-    """Same as init_pyconfig but with the NNX flags turned on."""
-    return self.init_pyconfig(pure_nnx=True, enable_nnx=True, pure_nnx_decoder=True, **kwargs)
+    """Same as init_pyconfig (NNX is the only path now)."""
+    return self.init_pyconfig(**kwargs)
 
   def _build_nnx_params(self, cfg, mesh):
     """Materialize an NNX Transformer and return its nnx.Param state."""
@@ -173,7 +173,7 @@ def test_basic_decode_nnx(self):
       )
 
   def test_quantize_passes_gate_for_nnx(self):
-    """pure_nnx + quantization (convert-on-load) reaches the actual machinery in train mode."""
+    """NNX + quantization (convert-on-load) reaches the actual machinery in train mode."""
     # checkpoint_is_quantized defaults to False — full-precision on disk, AQT
     # quantizes per-forward against the loaded kernel (train mode).
     cfg = self._init_nnx_pyconfig(quantization="int8")
@@ -187,7 +187,7 @@ def test_quantize_passes_gate_for_nnx(self):
       pass  # any other failure (e.g. checkpoint not found) is fine for this test
 
   def test_load_pre_quantized_nnx_passes_quant_gate(self):
-    """pure_nnx + quantization + checkpoint_is_quantized=True clears the load gate."""
+    """NNX + quantization + checkpoint_is_quantized=True clears the load gate."""
     cfg = self._init_nnx_pyconfig(quantization="int8", checkpoint_is_quantized=True)
     engine = maxengine.MaxEngine(cfg, jax.devices())
     self.assertEqual(engine._nnx_quant_mode_str, "serve")  # pylint: disable=protected-access
@@ -219,7 +219,7 @@ def test_quantized_prefill_nnx_train_mode(self):
     self.assertTrue(jnp.all(jnp.isfinite(prefill_result["logits"])))
 
   def test_lora_load_single_adapter_reaches_loader_on_nnx(self):
-    """pure_nnx + LoRA: load_single_adapter dispatches to the NNX loader.
+    """NNX + LoRA: load_single_adapter dispatches to the NNX loader.
 
     A nonexistent adapter path should raise FileNotFoundError from the
     loader itself. A NotImplementedError here would mean the dispatch
diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py
index f1c29c4f16..ffc5458212 100644
--- a/tests/post_training/unit/lora_utils_test.py
+++ b/tests/post_training/unit/lora_utils_test.py
@@ -55,8 +55,6 @@
     "base_mlp_dim": 256,
     "max_prefill_predict_length": 4,
     "model_name": "llama2-7b",
-    "enable_nnx": True,
-    "pure_nnx_decoder": True,
     "override_model_config": True,
     "weight_dtype": "bfloat16",
 }
diff --git a/tests/unit/aqt_serve_roundtrip_nnx_test.py b/tests/unit/aqt_serve_roundtrip_nnx_test.py
index 00d6825de7..006a4580d2 100644
--- a/tests/unit/aqt_serve_roundtrip_nnx_test.py
+++ b/tests/unit/aqt_serve_roundtrip_nnx_test.py
@@ -74,9 +74,6 @@ def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized):
         sys.argv[0],
         base_yml,
         "model_name=gpt3-52k",
-        "pure_nnx=true",
-        "enable_nnx=true",
-        "pure_nnx_decoder=true",
         "max_target_length=64",
         "max_prefill_predict_length=16",
         "per_device_batch_size=1",
diff --git a/tests/unit/maxengine_nnx_test.py b/tests/unit/maxengine_nnx_test.py
index 08303f817a..074db10c22 100644
--- a/tests/unit/maxengine_nnx_test.py
+++ b/tests/unit/maxengine_nnx_test.py
@@ -41,9 +41,6 @@ def _nnx_config(self, **kwargs):
         "max_target_length": 8,
         "per_device_batch_size": 1,
         "enable_checkpointing": False,
-        "pure_nnx": True,
-        "enable_nnx": True,
-        "pure_nnx_decoder": True,
     } | kwargs
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs)