[NNX] Delete Linen (12.1b): collapse dispatch in sharding and muon

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
 
@@ -142,26 +140,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)
@@ -193,7 +182,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
@@ -217,30 +206,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__":
@@ -249,4 +224,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)

src/maxtext/utils/sharding.py

@@ -28,7 +28,6 @@
 from maxtext.configs import pyconfig
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
-import optax
 
 _LOGGED_ACTIVATION_SHARDINGS = set()
 _ACTIVATION_SHARDINGS_DUMP = []
@@ -712,26 +711,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(
@@ -865,8 +845,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(

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
@@ -220,37 +219,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."""
 

tests/unit/optimizers_test.py

@@ -574,12 +574,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):

tests/unit/sharding_nnx_test.py

@@ -30,7 +30,6 @@
 
 @dataclass
 class _Cfg:
-  pure_nnx: bool = True
   shard_optimizer_over_data: bool = False
 
 
@@ -95,9 +94,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)

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}")