Fix linting

Author: hsuan-lun-chiang

src/maxtext/configs/base.yml

@@ -1086,8 +1086,8 @@ position_id_per_seconds: 25
 subslice_shape: ""
 
 # NNX
-enable_nnx: false
-pure_nnx_decoder: false
+enable_nnx: True
+pure_nnx_decoder: True
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/maxtext/layers/nnx_decoders.py

@@ -311,7 +311,7 @@ def __init__(
 
         num_moe = config.num_decoder_layers - config.first_num_dense_layers
 
-        self.moe_layer = self._create_scanned_layers(moe_cls, length=num_moe, rngs=rngs)
+        self.moe_layers = self._create_scanned_layers(moe_cls, length=num_moe, rngs=rngs)
       elif self.is_gemma3:
         attention_pattern_length = len(gemma3.GEMMA3_ATTENTION_PATTERN)
         scan_length = config.num_decoder_layers // attention_pattern_length
@@ -437,7 +437,7 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, **kwargs)
     prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
     graphdef, params, state = nnx.split(
         layers, nnx.Param, ...
-    )
+    )  # state: the mutable state we carry (KV cache, RNGs, etc.)
 
     scan_axis = self.config.param_scan_axis
     if scan_axis != 0:
@@ -458,10 +458,9 @@ def layer_fn(carry, scanned_vars):
         current_params = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), current_params)
 
       layer = nnx.merge(graphdef, current_params, current_state)
-
       layer_out = layer(carry, *args, **valid_kwargs)
-
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+      new_current_state = nnx.state(layer)
 
       return new_carry, new_current_state
 
@@ -823,43 +822,41 @@ def _apply_single_engram_layer(self, y, current_idx, layer_stack, *args, **kwarg
     graphdef, state = nnx.split(layer_stack)
     params, rest = state.split(nnx.Param, ...)
     scan_axis = self.config.param_scan_axis
-    
+
     # Helper to generate N-dimensional basic slices (e.g., x[:, idx, :])
     def _extract_slice(x, idx, axis):
       slices = tuple(idx if i == axis else slice(None) for i in range(x.ndim))
       return x[slices]
-    
+
     # Slice using native indexing instead of jnp.take
     sliced_params = jax.tree.map(lambda x: _extract_slice(x, current_idx, scan_axis), params)
     sliced_rest = jax.tree.map(lambda x: _extract_slice(x, current_idx, 0), rest)
-    
+
     single_layer = nnx.merge(graphdef, sliced_params, sliced_rest)
-    
+
     # Run the single layer
     out = single_layer(
-        y, *args, 
-        decoder_input_tokens=kwargs.get("decoder_input_tokens"), 
-        **kwargs.get("layer_kwargs", {})
+        y, *args, decoder_input_tokens=kwargs.get("decoder_input_tokens"), **kwargs.get("layer_kwargs", {})
     )
     y = out[0] if isinstance(out, tuple) else out
-    
+
     # Re-merge the updated state back into the specific slice of the stack
     new_state = nnx.state(single_layer)
     new_params, new_rest = new_state.split(nnx.Param, ...)
-    
+
     updated_params = jax.tree.map(
         lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(
             s, jnp.expand_dims(new_s, axis=scan_axis), current_idx, axis=scan_axis
-        ), 
-        params, new_params
+        ),
+        params,
+        new_params,
     )
     updated_rest = jax.tree.map(
-        lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(
-            s, jnp.expand_dims(new_s, axis=0), current_idx, axis=0
-        ), 
-        rest, new_rest
+        lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(s, jnp.expand_dims(new_s, axis=0), current_idx, axis=0),
+        rest,
+        new_rest,
     )
-    
+
     nnx.update(layer_stack, updated_params, updated_rest)
     return y
 
@@ -870,38 +867,32 @@ def _apply_scanned_chunk(self, y, current_idx, next_boundary, layer_stack, *args
       graphdef, state = nnx.split(layer_stack)
       params, rest = state.split(nnx.Param, ...)
       scan_axis = self.config.param_scan_axis
-      
+
       # Slice the chunk state along the correct axes
       chunk_params = jax.tree.map(
-          lambda x: jax.lax.dynamic_slice_in_dim(x, current_idx, scan_length, axis=scan_axis), 
-          params
-      )
-      chunk_rest = jax.tree.map(
-          lambda x: jax.lax.dynamic_slice_in_dim(x, current_idx, scan_length, axis=0), 
-          rest
+          lambda x: jax.lax.dynamic_slice_in_dim(x, current_idx, scan_length, axis=scan_axis), params
       )
+      chunk_rest = jax.tree.map(lambda x: jax.lax.dynamic_slice_in_dim(x, current_idx, scan_length, axis=0), rest)
       chunk_stack = nnx.merge(graphdef, chunk_params, chunk_rest)
-      
+
       # Apply sequentially
       y, chunk_stack = self._apply_layers_sequentially(
           chunk_stack, y, *args, length=scan_length, **kwargs.get("layer_kwargs", {})
       )
-      
+
       # Update the original stack state
       new_state = nnx.state(chunk_stack)
       new_params, new_rest = new_state.split(nnx.Param, ...)
-      
+
       updated_params = jax.tree.map(
-          lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(s, new_s, current_idx, axis=scan_axis),
-          params, new_params
+          lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(s, new_s, current_idx, axis=scan_axis), params, new_params
       )
       updated_rest = jax.tree.map(
-          lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(s, new_s, current_idx, axis=0),
-          rest, new_rest
+          lambda s, new_s: jax.lax.dynamic_update_slice_in_dim(s, new_s, current_idx, axis=0), rest, new_rest
       )
-      
+
       nnx.update(layer_stack, updated_params, updated_rest)
-        
+
     return y
 
   def _apply_interleaved_scanned_layers(self, y, layer_stack, start_idx, end_idx, engram_indices, *args, **kwargs):
@@ -990,7 +981,7 @@ def __call__(
 
           y = self._apply_interleaved_scanned_layers(
               y,
-              self.moe_layer,
+              self.moe_layers,
               0,
               (cfg.num_decoder_layers - cfg.first_num_dense_layers),
               [e - cfg.first_num_dense_layers for e in cfg.engram_layers],
@@ -1007,7 +998,7 @@ def __call__(
           if cfg.use_batch_split_schedule:
             policy = self.get_remat_policy()
 
-            mock_params = self._build_linen_params(self.moe_layer)
+            mock_params = self._build_linen_params(self.moe_layers)
 
             y = deepseek_batchsplit.scan_batch_split_layers(
                 y,
@@ -1021,8 +1012,8 @@ def __call__(
                 policy=policy,
             )
           else:
-            y, self.moe_layer = self._apply_layers_sequentially(
-                self.moe_layer, y, *layer_args, length=num_moe, **layer_kwargs
+            y, self.moe_layers = self._apply_layers_sequentially(
+                self.moe_layers, y, *layer_args, length=num_moe, **layer_kwargs
             )
       elif self.is_gemma3:
         y = self._apply_gemma3_scanned_blocks(

tests/unit/nnx_decoder_test.py

@@ -48,69 +48,59 @@
 # from maxtext.layers.nnx_decoders import decoder_as_linen
 # from maxtext.common.common_types import MODEL_MODE_TRAIN
 
+
 class TestNNXDecoderLayerLogicalAxesUnmocked(unittest.TestCase):
-    """
-    Executes pure, unmocked forward passes through NNXDecoderLayer to 
-    guarantee coverage of the logical_axis_names assignment block.
-    """
-
-    def setUp(self):
-        super().setUp()
-        self.rngs = nnx.Rngs(params=0, dropout=1)
-        self.base_cfg = _make_config()
-        self.mesh = _make_mesh(self.base_cfg)
-
-    def _make_dummy_inputs(self, cfg):
-        batch = cfg.global_batch_size_to_train_on
-        seq_len = cfg.max_target_length
-        emb_dim = cfg.emb_dim
-        
-        # Use jnp.ones to ensure stable, non-stochastic arrays for the forward pass
-        inputs = jnp.ones((batch, seq_len, emb_dim), dtype=cfg.dtype)
-        segment_ids = jnp.ones((batch, seq_len), dtype=jnp.int32)
-        positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len))
-        
-        return inputs, segment_ids, positions
-
-    def test_forward_pass_prefill_mode(self):
-        """Forces execution of: if self.model_mode == MODEL_MODE_PREFILL"""
-        cfg = _make_config()
-        layer = NNXDecoderLayer(
-            config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_PREFILL, rngs=self.rngs
-        )
-        inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
-
-        # A real forward pass ensures all sharding and normalization lines are executed
-        out, _ = layer(
-            inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_PREFILL
-        )
-        self.assertEqual(out.shape, inputs.shape)
-
-    def test_forward_pass_ep_as_context(self):
-        """Forces execution of: elif self.config.expert_shard_attention_option == EP_AS_CONTEXT..."""
-        cfg = _make_config(expert_shard_attention_option=EP_AS_CONTEXT)
-        layer = NNXDecoderLayer(
-            config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_TRAIN, rngs=self.rngs
-        )
-        inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
-
-        out, _ = layer(
-            inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_TRAIN
-        )
-        self.assertEqual(out.shape, inputs.shape)
-
-    def test_forward_pass_default_axes(self):
-        """Forces execution of the default 'else' fallback."""
-        cfg = _make_config(expert_shard_attention_option="none")
-        layer = NNXDecoderLayer(
-            config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_TRAIN, rngs=self.rngs
-        )
-        inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
-
-        out, _ = layer(
-            inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_TRAIN
-        )
-        self.assertEqual(out.shape, inputs.shape)
+  """
+  Executes pure, unmocked forward passes through NNXDecoderLayer to
+  guarantee coverage of the logical_axis_names assignment block.
+  """
+
+  def setUp(self):
+    super().setUp()
+    self.rngs = nnx.Rngs(params=0, dropout=1)
+    self.base_cfg = _make_config()
+    self.mesh = _make_mesh(self.base_cfg)
+
+  def _make_dummy_inputs(self, cfg):
+    batch = cfg.global_batch_size_to_train_on
+    seq_len = cfg.max_target_length
+    emb_dim = cfg.emb_dim
+
+    # Use jnp.ones to ensure stable, non-stochastic arrays for the forward pass
+    inputs = jnp.ones((batch, seq_len, emb_dim), dtype=cfg.dtype)
+    segment_ids = jnp.ones((batch, seq_len), dtype=jnp.int32)
+    positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len))
+
+    return inputs, segment_ids, positions
+
+  def test_forward_pass_prefill_mode(self):
+    """Forces execution of: if self.model_mode == MODEL_MODE_PREFILL"""
+    cfg = _make_config()
+    layer = NNXDecoderLayer(config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_PREFILL, rngs=self.rngs)
+    inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
+
+    # A real forward pass ensures all sharding and normalization lines are executed
+    out, _ = layer(inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_PREFILL)
+    self.assertEqual(out.shape, inputs.shape)
+
+  def test_forward_pass_ep_as_context(self):
+    """Forces execution of: elif self.config.expert_shard_attention_option == EP_AS_CONTEXT..."""
+    cfg = _make_config(expert_shard_attention_option=EP_AS_CONTEXT)
+    layer = NNXDecoderLayer(config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_TRAIN, rngs=self.rngs)
+    inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
+
+    out, _ = layer(inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_TRAIN)
+    self.assertEqual(out.shape, inputs.shape)
+
+  def test_forward_pass_default_axes(self):
+    """Forces execution of the default 'else' fallback."""
+    cfg = _make_config(expert_shard_attention_option="none")
+    layer = NNXDecoderLayer(config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_TRAIN, rngs=self.rngs)
+    inputs, segment_ids, positions = self._make_dummy_inputs(cfg)
+
+    out, _ = layer(inputs, segment_ids, positions, deterministic=True, model_mode=MODEL_MODE_TRAIN)
+    self.assertEqual(out.shape, inputs.shape)
+
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()