Rename sharding_names to out_sharding in NNX Variable metadata

This CL renames the sharding_names attribute to out_sharding for better
consistency with the sharding API. The new name more clearly indicates
the purpose of this metadata field.
## Changes
- Bump Flax version to 0.12.4
- Core changes in variablelib.py:
  - Add sharding_names to out_sharding metadata remapping for backward compatibility
  - Add deprecated sharding_names property that returns out_sharding with a warning
- Update nnx/spmd.py, core/spmd.py, core/meta.py, linen/spmd.py to use out_sharding
- Update all NNX tests to use the new attribute name
- Update qwix flax_util.py to check for out_sharding first, with fallback to sharding_names
- Update maxtext initializers.py to check for out_sharding first
- Update documentation and examples to use out_sharding
## Backward Compatibility
Existing code using sharding_names will continue to work via:
- Metadata remapping during Variable creation
- Deprecated Variable.sharding_names property

PiperOrigin-RevId: 859745972

Author: Cristian Garcia

README.md

@@ -158,7 +158,7 @@ To cite this repository:
   author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
   title = {{F}lax: A neural network library and ecosystem for {JAX}},
   url = {http://github.com/google/flax},
-  version = {0.12.3},
+  version = {0.12.4},
   year = {2024},
 }
 ```

docs_nnx/flip/4844-var-eager-sharding.md

@@ -56,12 +56,12 @@ with jax.set_mesh(mesh):
   ...
 ```
 
-For JAX explicit mode, remove the `sharding_names=` annotation on the `nnx.Variable`.
+For JAX explicit mode, remove the `out_sharding=` annotation on the `nnx.Variable`.
 
 
 # Implementation
 [implementation]: #implementation
 
-When an `nnx.Variable` is created, check for the metadata `sharding_names`, and if present, check if under a valid global mesh context of was supplied with a valid mesh. If no, throw error; if yes, call `jax.lax.with_sharding_constraint` to apply sharding constraint on the value.
+When an `nnx.Variable` is created, check for the metadata `out_sharding`, and if present, check if under a valid global mesh context of was supplied with a valid mesh. If no, throw error; if yes, call `jax.lax.with_sharding_constraint` to apply sharding constraint on the value.
 
 Note that this only works in auto sharding mode. User should use JAX-level APIs to annotate shardings for explicit mode.

docs_nnx/guides/flax_gspmd.ipynb

@@ -123,7 +123,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "nnx.Param(jnp.ones(4,4), sharding_names=(None, 'model'), eager_sharding=True, mesh=auto_mesh)"
+    "nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh=auto_mesh)"
    ]
   },
   {
@@ -134,7 +134,7 @@
     "\n",
     "Let's begin by sharding the simplest component possible - a Flax variable.\n",
     "\n",
-    "When you define a Flax variable, you can pass in a metadata field called `sharding_names`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.\n",
+    "When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.\n",
     "\n",
     "**You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook."
    ]
@@ -191,7 +191,7 @@
     "with jax.set_mesh(auto_mesh):\n",
     "  w = nnx.Param(\n",
     "    rngs.lecun_normal()((4, 8)),\n",
-    "    sharding_names=(None, 'model')\n",
+    "    out_sharding=(None, 'model')\n",
     "  )\n",
     "  print(w.sharding.spec)\n",
     "  jax.debug.visualize_array_sharding(w)  # already sharded!"

docs_nnx/guides/flax_gspmd.md

@@ -64,14 +64,14 @@ with nnx.use_eager_sharding(False):
 You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way.
 
 ```{code-cell} ipython3
-nnx.Param(jnp.ones(4,4), sharding_names=(None, 'model'), eager_sharding=True, mesh=auto_mesh)
+nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh=auto_mesh)
 ```
 
 ## Shard a single-array model
 
 Let's begin by sharding the simplest component possible - a Flax variable.
 
-When you define a Flax variable, you can pass in a metadata field called `sharding_names`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.
+When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.
 
 **You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook.
 
@@ -81,7 +81,7 @@ rngs = nnx.Rngs(0)
 with jax.set_mesh(auto_mesh):
   w = nnx.Param(
     rngs.lecun_normal()((4, 8)),
-    sharding_names=(None, 'model')
+    out_sharding=(None, 'model')
   )
   print(w.sharding.spec)
   jax.debug.visualize_array_sharding(w)  # already sharded!

docs_nnx/guides/transforms.ipynb

@@ -815,30 +815,30 @@
      "output_type": "stream",
      "text": [
       "Inner m.param.shape = (3, 5)\n",
-      "Inner m.param.sharding_names = ('a', None)\n",
+      "Inner m.param.out_sharding = ('a', None)\n",
       "Outter m.param.shape = (3, 4, 5)\n",
-      "Outter m.param.sharding_names = ('a', 'b', None)\n"
+      "Outter m.param.out_sharding = ('a', 'b', None)\n"
      ]
     }
    ],
    "source": [
     "mesh = jax.make_mesh((1, 1), ('a', 'b'))\n",
     "\n",
     "class Weights(nnx.Module):\n",
-    "  def __init__(self, array: jax.Array, sharding_names: tuple[str | None, ...]):\n",
-    "    self.param = nnx.Param(array, sharding_names=sharding_names)\n",
+    "  def __init__(self, array: jax.Array, out_sharding: tuple[str | None, ...]):\n",
+    "    self.param = nnx.Param(array, out_sharding=out_sharding)\n",
     "\n",
     "@nnx.vmap(in_axes=1, transform_metadata={nnx.PARTITION_NAME: 'b'})\n",
     "def f(m: Weights):\n",
     "  print(f'Inner {m.param.shape = }')\n",
-    "  print(f'Inner {m.param.sharding_names = }')\n",
+    "  print(f'Inner {m.param.out_sharding = }')\n",
     "\n",
     "with jax.set_mesh(mesh):\n",
-    "  m = Weights(jnp.ones((3, 4, 5)), sharding_names=('a', 'b', None))\n",
+    "  m = Weights(jnp.ones((3, 4, 5)), out_sharding=('a', 'b', None))\n",
     "  f(m)\n",
     "\n",
     "print(f'Outter {m.param.shape = }')\n",
-    "print(f'Outter {m.param.sharding_names = }')"
+    "print(f'Outter {m.param.out_sharding = }')"
    ]
   },
   {
@@ -862,19 +862,19 @@
      "output_type": "stream",
      "text": [
       "Outter m.param.shape = (3, 4, 5)\n",
-      "Outter m.param.sharding_names = ('a', 'b', None)\n"
+      "Outter m.param.out_sharding = ('a', 'b', None)\n"
      ]
     }
    ],
    "source": [
     "@nnx.vmap(out_axes=1, axis_size=4, transform_metadata={nnx.PARTITION_NAME: 'b'})\n",
     "def init_vmap():\n",
-    "  return Weights(jnp.ones((3, 5)), sharding_names=('a', None))\n",
+    "  return Weights(jnp.ones((3, 5)), out_sharding=('a', None))\n",
     "\n",
     "with jax.set_mesh(mesh):\n",
     "  m = init_vmap()\n",
     "print(f'Outter {m.param.shape = }')\n",
-    "print(f'Outter {m.param.sharding_names = }')"
+    "print(f'Outter {m.param.out_sharding = }')"
    ]
   }
  ],

docs_nnx/guides/transforms.md

@@ -391,20 +391,20 @@ Let's see an example of this in action:
 mesh = jax.make_mesh((1, 1), ('a', 'b'))
 
 class Weights(nnx.Module):
-  def __init__(self, array: jax.Array, sharding_names: tuple[str | None, ...]):
-    self.param = nnx.Param(array, sharding_names=sharding_names)
+  def __init__(self, array: jax.Array, out_sharding: tuple[str | None, ...]):
+    self.param = nnx.Param(array, out_sharding=out_sharding)
 
 @nnx.vmap(in_axes=1, transform_metadata={nnx.PARTITION_NAME: 'b'})
 def f(m: Weights):
   print(f'Inner {m.param.shape = }')
-  print(f'Inner {m.param.sharding_names = }')
+  print(f'Inner {m.param.out_sharding = }')
 
 with jax.set_mesh(mesh):
-  m = Weights(jnp.ones((3, 4, 5)), sharding_names=('a', 'b', None))
+  m = Weights(jnp.ones((3, 4, 5)), out_sharding=('a', 'b', None))
   f(m)
 
 print(f'Outter {m.param.shape = }')
-print(f'Outter {m.param.sharding_names = }')
+print(f'Outter {m.param.out_sharding = }')
 ```
 
 Here, you added a `sharding` metadata to the `nnx.Param` variables, and used `transform_metadata` to update the `sharding` metadata to reflect the axis changes. Specifically, you can see that the first axis `b` was removed from the `sharding` metadata when inside of `nnx.vmap`, and then added back when outside of `nnx.vmap`.
@@ -414,10 +414,10 @@ You can verify that this also works when `nnx.Module`s are created inside the tr
 ```{code-cell} ipython3
 @nnx.vmap(out_axes=1, axis_size=4, transform_metadata={nnx.PARTITION_NAME: 'b'})
 def init_vmap():
-  return Weights(jnp.ones((3, 5)), sharding_names=('a', None))
+  return Weights(jnp.ones((3, 5)), out_sharding=('a', None))
 
 with jax.set_mesh(mesh):
   m = init_vmap()
 print(f'Outter {m.param.shape = }')
-print(f'Outter {m.param.sharding_names = }')
+print(f'Outter {m.param.out_sharding = }')
 ```

examples/nnx_toy_examples/10_fsdp_and_optimizer.py

@@ -56,15 +56,15 @@ class MLP(nnx.Module):
   def __init__(self, din, dmid, dout, rngs: nnx.Rngs):
     self.w1 = nnx.Param(
       nnx.initializers.lecun_normal()(rngs.params(), (din, dmid)),
-      sharding_names=mesh_rules('embed', 'mlp'),
+      out_sharding=mesh_rules('embed', 'mlp'),
     )
     self.b1 = nnx.Param(
       jnp.zeros((dmid,)),
-      sharding_names=mesh_rules('mlp'),
+      out_sharding=mesh_rules('mlp'),
     )
     self.w2 = nnx.Param(
       nnx.initializers.lecun_normal()(rngs.params(), (dmid, dout)),
-      sharding_names=mesh_rules('embed', 'mlp'),
+      out_sharding=mesh_rules('embed', 'mlp'),
     )
 
   def __call__(self, x: jax.Array):

flax/core/meta.py

@@ -297,13 +297,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['out_sharding'] = 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('out_sharding')
     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

@@ -24,13 +24,13 @@
     Sharding,
 )
 
-def get_pspec(sharding_names, sharding_rules = None) -> PartitionSpec:
+def get_pspec(sharding, sharding_rules = None) -> PartitionSpec:
   """Given an `nnx.Variable`, return its `PartitionSpec`."""
   if get_logical_axis_rules() or sharding_rules:
     context_rules = get_logical_axis_rules()
     rules = composite_rules(context_rules, sharding_rules)
-    return PartitionSpec(*from_sharding_rules(sharding_names, rules))
-  return PartitionSpec(*sharding_names)
+    return PartitionSpec(*from_sharding_rules(sharding, rules))
+  return PartitionSpec(*sharding)
 
 def _apply_sharding(value, sharding, mesh):
   if mesh.are_all_axes_explicit:
@@ -44,10 +44,9 @@ def _apply_sharding(value, sharding, mesh):
 
 
 def shard_value(
-  value, sharding_names, sharding_rules,
-  mesh: jax.sharding.AbstractMesh | jax.sharding.Mesh | None
+  value, sharding, sharding_rules, mesh: jax.sharding.AbstractMesh | jax.sharding.Mesh | None
 ):
-  if not sharding_names:
+  if not sharding:
     return value
 
   if mesh is None:
@@ -56,9 +55,9 @@ def shard_value(
   if mesh is None:
     raise ValueError(
       'An auto mesh context or metadata is required if creating a variable'
-      f' with annotation {sharding_names=}. '
+      f' with annotation {sharding=}. '
       'For more guidance, see https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html.')
-  pspec = get_pspec(sharding_names, sharding_rules)
+  pspec = get_pspec(sharding, sharding_rules)
   return _apply_sharding(value, NamedSharding(mesh, pspec), mesh)
 
 

flax/linen/spmd.py

@@ -303,15 +303,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['out_sharding'] = 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('out_sharding')
     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})