fix(nnx): optimize scan memory by excluding read-only parameters from scan outputs

When using `jax.lax.scan` in Flax NNX decoder layers (e.g. DeepSeek-V3 / Gemma 2), returning `nnx.state(layer)` inside `layer_fn` caused `jax.lax.scan` to stack read-only model parameters (`nnx.Param`) across all 58 MoE layers into Python scan outputs. This allocated ~12.7 GiB of extra HBM memory and triggered `RESOURCE_EXHAUSTED` OOM errors during large model execution.

This change uses `_, _, updated_state = nnx.split(layer, nnx.Param, ...)` inside `layer_fn` to exclude read-only parameters from scan carry/output states, returning only mutable layer states to `jax.lax.scan`.

Verification:
- Local NNX decoder unit tests pass 100% (44/44 in `tests/unit/nnx_decoders_test.py`).
- Pre-commit code quality checks pass 100% (`codespell`, `pylint`, `pyink`, `yamllint`).
- End-to-end logit check on `mlperf-v5p` cluster (`v5p-128`) for DeepSeek-V3 671B ran forward pass to completion without OOM.

TAG=agy
CONV=0c977f36-644a-472a-981c-1856f8cac8a6

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_decoders.py

@@ -356,13 +356,17 @@ def layer_fn(carry, scanned_vars):
           **kwargs,
       )
       new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
-      return new_carry, nnx.state(layer)
+      # Avoid returning and stacking read-only parameters inside the scan body.
+      # This prevents huge unnecessary memory allocation.
+      _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+      return new_carry, updated_state
 
     final_carry, scanned_state = jax.lax.scan(layer_fn, inputs, (params, state))
 
     if scan_axis != 0:
       scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
-      scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
+      if scanned_params:
+        scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
       scanned_state = nnx.State.merge(scanned_params, scanned_other)
 
     nnx.update(self.scanned_layers, scanned_state)
@@ -957,7 +961,10 @@ def layer_fn(carry, scanned_vars):
         returned_params = updated_params
         new_current_state = nnx.State.merge(returned_params, updated_state)
       else:
-        new_current_state = nnx.state(layer)
+        # Avoid returning and stacking read-only parameters inside the scan body.
+        # This prevents huge unnecessary memory allocation.
+        _, _, updated_state = nnx.split(layer, nnx.Param, ...)
+        new_current_state = updated_state
 
       if use_kv:
         return new_carry, (new_current_state, updated_kv)