Rename sharding_names to sharding_metadata

Author: samanklesaria

flax/core/meta.py

@@ -300,13 +300,13 @@ def get_sharding(self, mesh: jax.sharding.Mesh) -> jax.sharding.Sharding:
   def to_nnx_metadata(self) -> dict[str, Any]:
     """Return a dict of metadata that can translate into an `nnx.Variable`."""
     metadata = dict(vars(self))
-    metadata['sharding_names'] = metadata.pop('names')
+    metadata['sharding_metadata'] = metadata.pop('names')
     return metadata
 
   @classmethod
   def from_nnx_metadata(cls, metadata: dict[str, Any]):
     """Given a dict of `nnx.Variable` format metadata, create a `nn.Partitioned`."""
-    metadata['names'] = metadata.pop('sharding_names')
+    metadata['names'] = metadata.pop('sharding_metadata')
     fields = {x.name for x in dataclasses.fields(cls)}
     return cls(**{k: v for k, v in metadata.items() if k in fields})
 

flax/core/spmd.py

@@ -45,6 +45,8 @@ def shard_value(value, sharding_names, sharding_rules, mesh):
       f' with annotation {sharding_names=}. '
       'For more guidance, see https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html.')
   pspec = get_pspec(sharding_names, sharding_rules)
+  if isinstance(sharding_names, NamedSharding) and mesh is not None:
+    assert sharding_names.mesh == mesh
   if mesh is not None:
     return _apply_sharding(value, NamedSharding(mesh, pspec))
   return _apply_sharding(value, pspec)
@@ -107,8 +109,10 @@ def composite_rules(rule1, rule2):
 
 
 def from_sharding_rules(
-    sharding: Sharding, sharding_rules: LogicalRules
+    sharding, sharding_rules: LogicalRules
 ) -> Sharding:
+  if isinstance(sharding, NamedSharding):
+    sharding = sharding.spec
   rules = {alias: on_mesh for (alias, on_mesh) in sharding_rules}
   return tuple(
       rules[str(s)] if (s and str(s) in rules) else s for s in sharding

flax/linen/spmd.py

@@ -290,15 +290,15 @@ def to_nnx_metadata(self) -> dict[str, Any]:
     """Return a dict of metadata that can translate into an `nnx.Variable`."""
     metadata = vars(self)
     if 'names' in metadata:
-      metadata['sharding_names'] = metadata.pop('names')
+      metadata['sharding_metadata'] = metadata.pop('names')
     if 'rules' in metadata:
       metadata['sharding_rules'] = metadata.pop('rules')
     return metadata
 
   @classmethod
   def from_nnx_metadata(cls, metadata: dict[str, Any]):
     """Given a dict of `nnx.Variable` format metadata, create a `nn.LogicallyPartitioned`."""
-    metadata['names'] = metadata.pop('sharding_names')
+    metadata['names'] = metadata.pop('sharding_metadata')
     metadata['rules'] = metadata.pop('sharding_rules')
     fields = {x.name for x in dataclasses.fields(cls)}
     return cls(**{k: v for k, v in metadata.items() if k in fields})

flax/nnx/spmd.py

@@ -45,8 +45,8 @@ def insert_field(fields, index, value):
   def _add_axis(x: tp.Any):
     if isinstance(x, variablelib.Variable):
       metadata = x.get_metadata()
-      if 'sharding_names' in metadata and metadata['sharding_names']:
-        sharding = metadata['sharding_names']
+      if 'sharding_metadata' in metadata and metadata['sharding_metadata']:
+        sharding = metadata['sharding_metadata']
         x.set_metadata(sharding_names=insert_field(sharding, index, axis_name))
 
       for k, v in other_meta.items():
@@ -74,7 +74,7 @@ def remove_field(fields, index, value):
 
   def _remove_axis(x: tp.Any):
     if isinstance(x, variablelib.Variable):
-      if hasattr(x, 'sharding_names') and x.sharding_names is not None:
+      if hasattr(x, 'sharding_metadata') and x.sharding_names is not None:
         x.set_metadata(
           sharding_names=remove_field(x.sharding_names, index, axis_name)
         )
@@ -119,7 +119,7 @@ def with_partitioning(
   """A wrapper over any initializer to add sharding annotation data to a `Variable`."""
   return variablelib.with_metadata(
     initializer,
-    sharding_names=sharding,
+    sharding_metadata=sharding,
     mesh=mesh,
     **metadata,
   )
@@ -128,8 +128,8 @@ def with_partitioning(
 def get_var_pspec(v: variablelib.Variable) -> PartitionSpec | None:
   """Given an `nnx.Variable`, return its `PartitionSpec`."""
   metadata = v.get_metadata()
-  if 'sharding_names' in metadata and metadata['sharding_names']:
-    sharding = metadata['sharding_names']
+  if 'sharding_metadata' in metadata and metadata['sharding_metadata']:
+    sharding = metadata['sharding_metadata']
     if core_spmd.get_logical_axis_rules() or 'sharding_rules' in metadata:
       context_rules = core_spmd.get_logical_axis_rules()
       local_rules = metadata.get('sharding_rules', ())
@@ -174,4 +174,4 @@ def get_abstract_model(init_fn, mesh):
       lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s),
       abs_state, get_named_sharding(abs_state, mesh)
     )
-  return gdef, abs_state
+  return gdef, abs_state

flax/nnx/variablelib.py

@@ -21,6 +21,7 @@
 import threading
 import typing as tp
 from typing import Any
+import warnings
 from flax import config
 
 import jax
@@ -375,16 +376,20 @@ def __init__(
       metadata['on_remove_axis'] = var_t.on_remove_axis
 
     if 'sharding' in metadata:
-      metadata['sharding_names'] = metadata.pop('sharding')
+      metadata['sharding_metadata'] = metadata.pop('sharding')
+
+    if 'sharding_names' in metadata: # for bw compat
+      warnings.warn("'sharding_names' is deprecated. Use 'sharding_metadata' instead.", DeprecationWarning)
+      metadata['sharding_metadata'] = metadata.pop('sharding_names')
 
     object.__setattr__(self, '_var_metadata', metadata)
     # run create_value hooks
     value = self.create_value(self.raw_value)
 
     # shard the value if applicable
-    if metadata.get('eager_sharding', using_eager_sharding()) and 'sharding_names' in metadata:
+    if metadata.get('eager_sharding', using_eager_sharding()) and 'sharding_metadata' in metadata:
       value = core_spmd.shard_value(
-        value, metadata['sharding_names'], metadata.get('sharding_rules', None),
+        value, metadata['sharding_metadata'], metadata.get('sharding_rules', None),
         metadata.get('mesh', None))
 
     # Create the ref out of the array value
@@ -394,6 +399,9 @@ def __init__(
     object.__setattr__(self, 'raw_value', value)
 
   def __getattr__(self, name: str) -> tp.Any:
+    if name == 'sharding_names': # for backward compatibility
+      warnings.warn("'sharding_names' is deprecated. Use 'sharding_metadata' instead.", DeprecationWarning)
+      return self.sharding_metadata
     if name in object.__getattribute__(self, '_var_metadata'):
       return self._var_metadata[name]
     return getattr(self.raw_value, name)

tests/nnx/bridge/wrappers_test.py

@@ -410,7 +410,7 @@ def test_nnx_to_linen_metadata(self):
       pspec_tree = nn.get_partition_spec(variables)
     assert y.shape == (1, 64)
     self.assertIsInstance(variables['params']['kernel'], nnx.bridge.NNXMeta)
-    assert variables['params']['kernel'].metadata['sharding_names'] == ('in', 'out')
+    assert variables['params']['kernel'].metadata['sharding_metadata'] == ('in', 'out')
     self.assertEqual(pspec_tree['params']['kernel'],
                      jax.sharding.PartitionSpec('in', 'out'))
     np.testing.assert_allclose(y, x @ variables['params']['kernel'].value)
@@ -519,8 +519,8 @@ def __call__(self, x):
       w, b = model.inner.dot['w'], model.inner.b
       np.testing.assert_allclose(model(x), x @ w + b)
     self.assertIsInstance(w, nnx.Param)
-    assert hasattr(w, 'sharding_names') and w.sharding_names == ('in', 'out')
-    assert hasattr(b, 'sharding_names') and b.sharding_names == ('out-alias', )
+    assert hasattr(w, 'sharding_metadata') and w.sharding_metadata == ('in', 'out')
+    assert hasattr(b, 'sharding_metadata') and b.sharding_metadata == ('out-alias', )
 
   def test_linen_nnx_linen(self):
     # TODO: add when we can safely `lazy_init` the NNX module inside `ToLinen` without