Merge pull request #4637 from IvyZX:custom-meta

Flax Authors · Flax Authors · commit 1e4838007194 · 2025-03-18T17:01:55.000-07:00
PiperOrigin-RevId: 738191979
diff --git a/flax/nnx/spmd.py b/flax/nnx/spmd.py
@@ -38,17 +38,24 @@ class HasSharding(tp.Protocol):
 def _has_sharding(x: tp.Any) -> tp.TypeGuard[HasSharding]:
   return hasattr(x, 'sharding') and x.sharding is not None
 
-def add_axis(tree: A, index: int, params: tp.Mapping) -> A:
-  axis_name = _get_partition_name(params)
+def add_axis(tree: A, index: int, transform_metadata: tp.Mapping) -> A:
+  axis_name, other_meta = _get_partition_name_and_metadata(transform_metadata)
+
+  def insert_field(fields, index, value):
+    iterable = list(fields)
+    while len(iterable) < index:
+      iterable.append(None)
+    iterable.insert(index, value)
+    return tuple(iterable)
 
   def _add_axis(x: tp.Any):
     if isinstance(x, variablelib.VariableState):
       if _has_sharding(x) and x.sharding is not None:
-        sharding: list[str | None] = list(x.sharding)
-        while len(sharding) < index:
-          sharding.append(None)
-        sharding.insert(index, axis_name)
-        x.sharding = tuple(sharding)  # type: ignore
+        x.sharding = insert_field(x.sharding, index, axis_name)
+
+      for k, v in other_meta.items():
+        if hasattr(x, k) and (t := getattr(x, k)) and isinstance(t, tuple):
+          setattr(x, k, insert_field(t, index, v))
 
       assert isinstance(x, variablelib.VariableState)
       x.add_axis(index, axis_name)
@@ -59,15 +66,23 @@ def _add_axis(x: tp.Any):
   )
 
 
-def remove_axis(tree: A, index: int, params: tp.Mapping[tp.Any, tp.Any]) -> A:
-  axis_name = _get_partition_name(params)
+def remove_axis(tree: A, index: int, transform_metadata: tp.Mapping[tp.Any, tp.Any]) -> A:
+  axis_name, other_meta = _get_partition_name_and_metadata(transform_metadata)
+
+  def remove_field(fields, index, value):
+    iterable = list(fields)
+    assert iterable.pop(index) == value
+    return tuple(iterable)
 
   def _remove_axis(x: tp.Any):
     if isinstance(x, variablelib.VariableState):
       if hasattr(x, 'sharding') and x.sharding is not None:
-        sharding = list(x.sharding)
-        assert sharding.pop(index) == axis_name
-        x.sharding = tuple(sharding)
+        x.sharding = remove_field(x.sharding, index, axis_name)
+
+      for k, v in other_meta.items():
+        if hasattr(x, k) and (t := getattr(x, k)) and isinstance(t, tuple):
+          setattr(x, k, remove_field(t, index, v))
+
       x.remove_axis(index, axis_name)
     return x
 
@@ -78,13 +93,17 @@ def _remove_axis(x: tp.Any):
   )
 
 
-def _get_partition_name(params: tp.Mapping[tp.Any, tp.Any]) -> str:
-  if PARTITION_NAME not in params:
+def _get_partition_name_and_metadata(
+  transform_metadata: tp.Mapping[tp.Any, tp.Any]
+) -> tuple[str, tp.Mapping[tp.Any, tp.Any]]:
+  if PARTITION_NAME not in transform_metadata:
     raise ValueError(
       'Trying to transform a Partitioned variable but "partition_name" '
-      f'is not specified in scan_metadata: {params}'
+      f'is not specified in transform_metadata: {transform_metadata}'
     )
-  return params[PARTITION_NAME]
+  other_meta = dict(transform_metadata)  # shallow copy
+  other_meta.pop(PARTITION_NAME)
+  return transform_metadata[PARTITION_NAME], other_meta
 
 
 def get_partition_spec(tree: A) -> A:
diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py
@@ -108,7 +108,7 @@ class MLP(nnx.Module):
       @nnx.split_rngs(splits=5)
       @nnx.vmap(
           in_axes=(0, 0),
-          transform_metadata={nnx.PARTITION_NAME: 'layers'},
+          transform_metadata={nnx.PARTITION_NAME: 'layers', 'nickname': 'nick'},
       )
       def __init__(self, rngs: nnx.Rngs):
         self.linear = nnx.Linear(
@@ -117,6 +117,7 @@ def __init__(self, rngs: nnx.Rngs):
           kernel_init=nnx.with_metadata(
             nnx.initializers.lecun_normal(),
             sharding=('din', 'dout'),
+            nickname=('in', 'out'),
             on_add_axis=lambda _, idx, name: kadds.append((idx, name)),
             on_remove_axis=lambda _, idx, name: kremoves.append((idx, name)),
           ),
@@ -145,6 +146,7 @@ def __call__(self, x: jax.Array):
     m = MLP(rngs=nnx.Rngs(0))
     self.assertEqual(m.linear.kernel.shape, (5, 3, 3))
     self.assertEqual(m.linear.kernel.sharding, ('layers', 'din', 'dout'))
+    self.assertEqual(m.linear.kernel.nickname, ('nick', 'in', 'out'))
     self.assertEqual(m.linear.bias.shape, (5, 3))
     # One add_axis called to add the `nnx.vmap` dimension
     self.assertEqual(kadds, [(0, 'layers')])
