fix: update

Author: mesakhcienet

src/maxtext/layers/nnx_decoders.py

@@ -413,7 +413,10 @@ def _create_scanned_layers(self, decoder_layer_class, length: int, metadata_axis
 
     # Create a reference layer to capture the module graph structure (graphdef).
     # This layer's params are discarded — only the structure is kept.
-    ref_rngs = nnx.Rngs(0)
+    # Must use the first slice of the forked rngs (not a dummy Rngs(0)) so the
+    # graphdef has the same number of RNG state leaves as the scan-created layers.
+    first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
+    ref_rngs = nnx.merge(rngs_graphdef, first_rng_state)
     ref_layer = decoder_layer_class(
         config=self.config, mesh=self.mesh, quant=self.quant,
         model_mode=self.model_mode, rngs=ref_rngs, **layer_kwargs

src/maxtext/layers/nnx_wrappers.py

@@ -435,9 +435,6 @@ def _module_kwargs():
       _fix_for_qwix_quantization(module)
       method_fn = _get_module_method(module, nnx_method)
       out = method_fn(module, *args, **kwargs)
-      # Free the NNX module eagerly to avoid holding both NNX params and
-      # Linen variable copies in memory simultaneously during init tracing.
-      del module
       return out
 
     # create the nnx module