fix: replace nnx.stateaxes

Author: mesakhcienet

src/maxtext/layers/nnx_pipeline.py

@@ -84,11 +84,16 @@ def build_batched_rngs(shape):
 
       return nnx.Rngs(**kwargs)
 
+    # Define a wrapper that explicitly splits the instantiated layer
+    def create_layer_fn(rngs):
+      layer = stage_factory(rngs)
+      return nnx.split(layer, nnx.Param, ...)
+
     # Vmap over stages natively adds 'layers' metadata to the logical partition spec!
     vmap_stages = nnx.vmap(
-        stage_factory,
+        create_layer_fn,
         in_axes=0,
-        out_axes=nnx.StateAxes({nnx.Param: 0, ...: 0}),
+        out_axes=(None, 0, 0),  # None for static graphdef, 0 for params, 0 for rest
         axis_name=self.spmd_axis_name,
         transform_metadata={nnx.PARTITION_NAME: "layers"},
     )
@@ -97,14 +102,16 @@ def build_batched_rngs(shape):
       vmap_repeats = nnx.vmap(
           vmap_stages,
           in_axes=0,
-          out_axes=nnx.StateAxes({nnx.Param: 0, ...: 0}),
+          out_axes=(None, 0, 0), # Graphdef remains unbatched (None) through the second vmap
           transform_metadata={nnx.PARTITION_NAME: "circular_repeats"},
       )
       batched_rngs = build_batched_rngs((self.config.num_pipeline_repeats, self.num_stages))
-      self.layers = vmap_repeats(batched_rngs)
+      graphdef, params, rest = vmap_repeats(batched_rngs)
+      self.layers = nnx.merge(graphdef, params, rest)
     else:
       batched_rngs = build_batched_rngs((self.num_stages,))
-      self.layers = vmap_stages(batched_rngs)
+      graphdef, params, rest = vmap_stages(batched_rngs)
+      self.layers = nnx.merge(graphdef, params, rest)
 
     # -------------------------------------------------------------------------
     # Sharding Configs

src/maxtext/layers/pipeline.py

@@ -1017,11 +1017,14 @@ def build_batched_rngs(shape):
 
       return nnx.Rngs(**kwargs)
 
-    # Vmap over stages natively adds 'layers' metadata to the logical partition spec!
+    def create_layer_fn(rngs):
+      layer = stage_factory(rngs)
+      return nnx.split(layer, nnx.Param, ...)
+
     vmap_stages = nnx.vmap(
-        stage_factory,
+        create_layer_fn,
         in_axes=0,
-        out_axes=nnx.StateAxes({nnx.Param: 0, ...: 0}),
+        out_axes=(None, 0, 0),
         axis_name=self.spmd_axis_name,
         transform_metadata={nnx.PARTITION_NAME: "layers"},
     )
@@ -1030,14 +1033,16 @@ def build_batched_rngs(shape):
       vmap_repeats = nnx.vmap(
           vmap_stages,
           in_axes=0,
-          out_axes=nnx.StateAxes({nnx.Param: 0, ...: 0}),
+          out_axes=(None, 0, 0),
           transform_metadata={nnx.PARTITION_NAME: "circular_repeats"},
       )
       batched_rngs = build_batched_rngs((self.config.num_pipeline_repeats, self.num_stages))
-      self.layers = vmap_repeats(batched_rngs)
+      graphdef, params, rest = vmap_repeats(batched_rngs)
+      self.layers = nnx.merge(graphdef, params, rest)
     else:
       batched_rngs = build_batched_rngs((self.num_stages,))
-      self.layers = vmap_stages(batched_rngs)
+      graphdef, params, rest = vmap_stages(batched_rngs)
+      self.layers = nnx.merge(graphdef, params, rest)
 
   def get_current_repeat_from_stages(self, weights, loop_iteration, physical_partition_spec=None):
     """get current repeat from stages"""