Rename sharding_names to sharding_metadata

Author: samanklesaria

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/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)