Update

Author: hsuan-lun-chiang

src/MaxText/layers/nnx_decoders.py

@@ -194,10 +194,10 @@ def __call__(
     layer_output = _maybe_shard_with_logical(layer_output, logical_axis_names)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )
@@ -380,10 +380,11 @@ def create_layer_fn(rng):
     except:  # pylint: disable=bare-except
       pass
 
+    scan_axis = self.config.param_scan_axis
     layers_vmapped = nnx.vmap(
         create_layer_fn,
         in_axes=0,
-        out_axes=0,
+        out_axes=scan_axis,
         axis_name="layers",
         transform_metadata={nnx.PARTITION_NAME: "layers"},
     )(forked_rngs)