[NNX] NNX migration (12/N): delete Linen code paths, classes, and compatibility flags

NNX is now the only model path (PR11 flipped pure_nnx/enable_nnx/pure_nnx_decoder
to True), so these flags are no longer dispatch points. Delete the Linen code:

- Collapse all flag and isinstance(model, nn.Module) dispatch to the NNX branch
  across ~22 src files (train.py, maxtext_utils, train_utils, sharding, diloco,
  maxengine, layerwise_quantization, grpo_trainer, lora_utils, checkpointing,
  convert_gpt3_ckpt_from_paxml, ...). Zero executable flag reads remain in src.
- Delete TransformerLinenPure; the Linen decoder stack Decoder / DecoderLayer /
  SequentialBlockDecoderLayers (decoders.py 1525->47, only deepstack_process kept);
  and 28 dead *_as_linen ToLinen wrappers. The wrapped NNX classes are unchanged.
- Remove the pure_nnx / enable_nnx / pure_nnx_decoder flags from configs/types.py,
  base.yml, inference/vllm.yml, and pyconfig.
- Delete 21 obsolete Linen-only tests; drop redundant flag args elsewhere.

Kept for focused follow-ups: the transformer_as_linen / init_initial_state
NNX->Linen bridge (checkpoint-conversion tools), the Linen GRPO reference
grpo_loss_fn (torch-gated correctness tests), and the Linen pipeline.py (NNX
pipeline parallelism pending PR11.5; test_pipeline_subset skipped).

Verified on CPU: NNX unit suite 213 passed / 28 skipped (3 non-regression fails),
5 train_compile AOT cases, nnx_decoders_test 40 passed, lint 10/10, base.yml
config-load smoke. Net -5,346 lines across 59 files.

Author: ecnal-cienet

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -35,7 +35,6 @@
 """
 
 import argparse
-import functools
 import gc
 import os
 import sys
@@ -47,11 +46,7 @@
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.common import checkpointing
-from maxtext.common.common_types import MODEL_MODE_TRAIN
-from maxtext.layers import quantizations
 from maxtext.common import train_state_nnx
-from maxtext.models.models import transformer_as_linen
-from maxtext.optimizers import optimizers
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
@@ -92,23 +87,15 @@ def convert(paxml_ckpt_path, maxtext_model_name, base_output_directory, run_name
   devices_array = maxtext_utils.create_device_mesh(cfg)
   mesh = Mesh(devices_array, cfg.mesh_axes)
 
-  if cfg.pure_nnx:
-    rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
-    model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
-    _, tx = train_utils.create_training_optimizer(cfg, model)
-    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)
+  rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
+  model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
+  _, tx = train_utils.create_training_optimizer(cfg, model)
+  _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)
 
-    def init_state_fn():
-      nnx_model = _create_model_partial()
-      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
-
-  else:
-    quant = quantizations.configure_quantization(cfg)
-    model = transformer_as_linen(cfg, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-    learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(cfg)
-    tx = optimizers.get_optimizer(cfg, learning_rate_schedule)
-    init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, cfg, True, init_rng)
+  def init_state_fn():
+    nnx_model = _create_model_partial()
+    optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+    return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
   checkpoint_manager = checkpointing.create_orbax_checkpoint_manager(
       cfg.checkpoint_dir,
@@ -201,21 +188,15 @@ def init_state_fn():
       "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None),
   }
 
-  if cfg.pure_nnx:
-    # NNX state-tree paths after `nnx.split(TrainStateNNX)`:
-    #   model params     -> ['model']<rest>.value
-    #   adam mu / nu     -> ['optimizer']['opt_state']['mu' | 'nu']<rest>.value
-    #   step             -> ['optimizer']['step'].value
-    #   opt count        -> ['optimizer']['opt_state']['count'].value
-    state_map = {
-        ".optimizer.step.value": ("step", None),
-        ".optimizer.opt_state.count.value": ("opt_states_0.no_prefix_0.count", None),
-    }
-  else:
-    state_map = {
-        ".step": ("step", None),
-        ".opt_state.count": ("opt_states_0.no_prefix_0.count", None),
-    }
+  # NNX state-tree paths after `nnx.split(TrainStateNNX)`:
+  #   model params     -> ['model']<rest>.value
+  #   adam mu / nu     -> ['optimizer']['opt_state']['mu' | 'nu']<rest>.value
+  #   step             -> ['optimizer']['step'].value
+  #   opt count        -> ['optimizer']['opt_state']['count'].value
+  state_map = {
+      ".optimizer.step.value": ("step", None),
+      ".optimizer.opt_state.count.value": ("opt_states_0.no_prefix_0.count", None),
+  }
 
   def get_layer_prefix(keystr_pax):
     # different path format between decoder_layer variable
@@ -228,26 +209,15 @@ def get_layer_prefix(keystr_pax):
 
   for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items():
     prefix_pax_opt_state = get_layer_prefix(keystr_pax)
-    if cfg.pure_nnx:
-      state_map[f".model{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f".optimizer.opt_state.mu{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f".optimizer.opt_state.nu{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
-    else:
-      state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f".opt_state.nu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
+    state_map[f".model{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn)
+    state_map[f".optimizer.opt_state.mu{keystr_maxtext}.value"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f".optimizer.opt_state.nu{keystr_maxtext}.value"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+        transform_fn,
+    )
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
@@ -299,7 +269,7 @@ def map_fn(key_path, value):
   max_logging.log("converted state finished")
   max_utils.print_mem_stats("converted state finished")
 
-  step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step
+  step_value = int(converted_state.optimizer.step.value)
   if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
     max_logging.log(f"saved a checkpoint at step {step_value}")
   # Upon preemption, exit when and only when all ongoing saves are complete.

src/maxtext/common/checkpointing.py

@@ -949,19 +949,14 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
   if step is not None:
     actual_step = int(step)
   else:
-    if config.pure_nnx:
-      actual_step = int(state.optimizer.step) - 1
-    else:
-      # Linen TrainState has .step attribute
-      actual_step = int(state.step) - 1
+    actual_step = int(state.optimizer.step) - 1
 
   if checkpoint_manager.latest_step() == actual_step:
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
-  if config.pure_nnx:
-    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
-    state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict())
+  # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
+  state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict())
 
   # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic.
   # This occurs if this function was called:

src/maxtext/configs/base.yml

@@ -1168,11 +1168,6 @@ position_id_per_seconds: 25
 # Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium.
 subslice_shape: ""
 
-# NNX
-enable_nnx: true
-pure_nnx_decoder: true
-pure_nnx: true
-
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net
 gdn_conv_kernel_dim: 4

src/maxtext/configs/inference/vllm.yml

@@ -16,8 +16,6 @@ base_config: "base.yml"
 attention: "vllm_rpa"
 model_call_mode: "inference"
 
-# NNX required for vLLM integration
-enable_nnx: true
 # Avoid re-initializing JAX distributed system when using vLLM
 skip_jax_distributed_system: true
 # Scanned layers are not supported with vLLM integration

src/maxtext/configs/pyconfig_deprecated.py

@@ -193,8 +193,7 @@ def validate_expert_shard_attention_option(expert_shard_attention_option: str) -
     )
 
 
-def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int, enable_nnx: bool):
-  del enable_nnx  # NNX vocab tiling supported via vocab_tiling_nnx_loss in vocabulary_tiling.py
+def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int):
   if (per_device_batch_size * max_target_length) % num_vocab_tiling != 0:
     raise ValueError("Per device batch size times sequence length should be divisible by the number of vocab tiles.")
 
@@ -238,9 +237,7 @@ def validate_keys(keys):
   validate_model_call_mode(keys["model_call_mode"])
   validate_prefill_and_target_lengths(keys["max_prefill_predict_length"], keys["max_target_length"])
   validate_rope_type(keys["rope_type"])
-  validate_vocab_tiling(
-      keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"], keys["enable_nnx"]
-  )
+  validate_vocab_tiling(keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"])
   if keys["enable_rampup_batch_size"]:
     validate_rampup_batch_size(
         keys["per_device_batch_size_start"],

src/maxtext/configs/types.py

@@ -895,11 +895,8 @@ class HardwareAndMesh(BaseModel):
       CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity."
   )
   allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.")
-  enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.")
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
-  pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.")
-  pure_nnx: bool = Field(False, description="Whether to enable pure NNX mode.")
   remove_size_one_mesh_axis_from_type: bool = Field(
       True, description="Whether to remove size one mesh axis from type through jax.config."
   )
@@ -2498,8 +2495,6 @@ def validate_and_set_hlo_dump_defaults():
     if self.distill_beta > 0.0:
       if not self.scan_layers:
         raise ValueError("a value of self.distill_beta > 0.0 requires self.scan_layers = True")
-      if not self.enable_nnx:
-        raise ValueError("a value of self.distill_beta > 0.0 requires self.enable_nnx = True")
 
     # Validate distillation schedule parameters
     if self.distill_alpha_end is not None and not 0.0 <= self.distill_alpha_end <= 1.0: