fix(nnx): keep only newly created params in _apply_layers_sequentially

src/maxtext/layers/nnx_decoders.py

@@ -1081,6 +1081,12 @@ def _extract_matching_state(template, full):
     dynamic_graph_init = bool(getattr(self, "disable_quant_stats_update", False))
     updated_graphdef = [graphdef]
 
+    # Parameters fed in as scan inputs come back out unchanged, so they must not be
+    # re-emitted as scan outputs (see layer_fn). Anything else the body produces --
+    # including parameters materialized while tracing, such as Qwix LoRA adapters,
+    # which are ``nnx.Param`` subclasses -- still has to leave the scan.
+    carried_param_paths = {path for path, _ in nnx.to_flat_state(params)}
+
     use_kv = kv_caches_stacked is not None
     use_forced_routing = forced_routed_experts_scanned is not None
 
@@ -1128,7 +1134,13 @@ def layer_fn(carry, scanned_vars):
       if dynamic_graph_init:
         new_graphdef, updated_params, updated_state = nnx.split(layer, nnx.Param, ...)
         updated_graphdef[0] = new_graphdef
-        returned_params = updated_params
+        # Drop the parameters that were carried in: jax.lax.scan stacks every
+        # output, so returning them would materialize a second copy of the
+        # stacked layer weights. Parameters created inside the body are still
+        # returned.
+        returned_params = nnx.from_flat_state(
+            [(path, value) for path, value in nnx.to_flat_state(updated_params) if path not in carried_param_paths]
+        )
         new_current_state = nnx.State.merge(returned_params, updated_state)
       else:
         # Avoid returning and stacking read-only parameters inside the scan body.
@@ -1192,9 +1204,11 @@ def layer_fn(carry, scanned_vars):
 
     if dynamic_graph_init:
       # If graph changed, we need to merge with the new graphdef.
-      # Note: scanned_state here is the full state (Params + rest).
+      # Note: scanned_state holds only the params created inside the body plus
+      # the rest; the carried-in params are read back off `layers`, which keeps
+      # their array identity so a second adapter shares one base.
       new_params, new_rest = scanned_state.split(nnx.Param, ...)
-      out_layers = nnx.merge(updated_graphdef[0], new_params, new_rest)
+      out_layers = nnx.merge(updated_graphdef[0], nnx.state(layers, nnx.Param), new_params, new_rest)
     else:
       clean_state = nnx.filter_state(scanned_state, nnx.Not(nnx.RngState))
       nnx.update(layers, clean_state)

tests/unit/nnx_decoders_test.py

@@ -1574,5 +1574,42 @@ def __call__(self, x, **kwargs):
       maxtext_utils_nnx.nnx_add_and_sync_scan_axis = original_add_scan_axis
 
 
+class TestApplyLayersSequentiallyDynamicGraphInit(unittest.TestCase):
+  """Params created inside the scan body must not drag the base stack out with them."""
+
+  class _AdapterLayer(nnx.Module):
+    """A layer that materializes a new param while tracing, as Qwix LoRA does."""
+
+    def __init__(self):
+      self.p = nnx.Param(jax.numpy.zeros((2,)))
+
+    def __call__(self, x, **kwargs):
+      self.adapter = nnx.LoRAParam(jax.numpy.ones((2,)))
+      return x + self.p.value + self.adapter.value, None
+
+  def _run(self, param_scan_axis):
+    cfg = _make_config(param_scan_axis=param_scan_axis)
+    decoder = NNXDecoder(config=cfg, mesh=_make_mesh(cfg), model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0))
+    layers = nnx.vmap(self._AdapterLayer, in_axes=(), out_axes=param_scan_axis, axis_size=2)()
+    # Qwix sets this on every module for the duration of its init pass.
+    decoder.disable_quant_stats_update = True
+    base_before = layers.p.value
+    # pylint: disable=protected-access
+    _, out_layers, _ = decoder._apply_layers_sequentially(layers=layers, x_in=jax.numpy.zeros((2,)), length=2)
+    return base_before, out_layers
+
+  def test_created_param_escapes_the_scan(self):
+    for axis in (0, 1):
+      with self.subTest(param_scan_axis=axis):
+        _, out_layers = self._run(axis)
+        self.assertTrue(hasattr(out_layers, "adapter"))
+
+  def test_base_params_are_not_restacked(self):
+    for axis in (0, 1):
+      with self.subTest(param_scan_axis=axis):
+        base_before, out_layers = self._run(axis)
+        self.assertIs(out_layers.p.value, base_before)
+
+
 if __name__ == "__main__":
   unittest.main()