_apply_sharding disallow mixed Explicit/Auto mesh

PiperOrigin-RevId: 859299941

Author: Cristian Garcia

flax/core/spmd.py

@@ -32,22 +32,34 @@ def get_pspec(sharding_names, sharding_rules = None) -> PartitionSpec:
     return PartitionSpec(*from_sharding_rules(sharding_names, rules))
   return PartitionSpec(*sharding_names)
 
-def _apply_sharding(value, sharding):
-  with jax.disable_jit(False):
-    return jax.jit(lambda x: x, out_shardings=sharding)(value)
+def _apply_sharding(value, sharding, mesh):
+  if mesh.are_all_axes_explicit:
+    return jax.sharding.reshard(value, sharding)
+  elif mesh.are_all_axes_auto:
+    return jax.lax.with_sharding_constraint(value, sharding)
+  else:
+    raise ValueError(
+        'Mesh must have all axes as Explicit or all axes as Auto. '
+        f'Got mixed axis types: {mesh.axis_types}')
+
 
-def shard_value(value, sharding_names, sharding_rules, mesh):
+def shard_value(
+  value, sharding_names, sharding_rules,
+  mesh: jax.sharding.AbstractMesh | jax.sharding.Mesh | None
+):
   if not sharding_names:
     return value
-  if not mesh and not meta.global_mesh_defined():
+
+  if mesh is None:
+    mesh = meta.get_global_mesh()
+
+  if mesh is None:
     raise ValueError(
       'An auto mesh context or metadata is required if creating a variable'
       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 mesh is not None:
-    return _apply_sharding(value, NamedSharding(mesh, pspec))
-  return _apply_sharding(value, pspec)
+  return _apply_sharding(value, NamedSharding(mesh, pspec), mesh)
 
 
 

tests/nnx/spmd_test.py

@@ -301,57 +301,35 @@ def __init__(self, rngs):
     assert jax.tree.leaves(abs_state)[0].sharding.is_equivalent_to(
       NamedSharding(mesh, P(None, 'model')), ndim=2)
 
-  def test_explicit_sharding(self):
-    mesh = jax.make_mesh(
-      (2, 2),
-      ('row', 'col'),
-      axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Explicit),
-    )
-    v = nnx.Variable(
-      jnp.ones((4, 4)),
-      sharding_names=('row', 'col'),
-      mesh=mesh,
-    )
-    self.assertEqual(v.sharding.mesh, mesh)
-    self.assertEqual(
-      v.sharding.spec,
-      P('row', 'col'),
-    )
+  @parameterized.parameters('auto', 'explicit', 'mixed')
+  def test_sharding_axis_types(self, mode):
+    if mode == 'auto':
+      axis_types = (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto)
+    elif mode == 'explicit':
+      axis_types = (jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit)
+    else:
+      axis_types = (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Explicit)
 
-  def test_explicit_sharding_disable_jit(self):
     mesh = jax.make_mesh(
       (2, 2),
       ('row', 'col'),
-      axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Explicit),
+      axis_types=axis_types,
     )
-    with jax.disable_jit(True):
+    if mode == 'mixed':
+      with self.assertRaises(ValueError):
+        nnx.Variable(
+          jnp.ones((4, 4)),
+          sharding_names=('row', 'col'),
+          mesh=mesh,
+        )
+    else:
       v = nnx.Variable(
         jnp.ones((4, 4)),
         sharding_names=('row', 'col'),
         mesh=mesh,
       )
-    self.assertEqual(v.sharding.mesh, mesh)
-    self.assertEqual(
-      v.sharding.spec,
-      P('row', 'col'),
-    )
-
-  def test_explicit_sharding_mesh_context(self):
-    mesh = jax.make_mesh(
-      (2, 2),
-      ('row', 'col'),
-      axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Explicit),
-    )
-    with jax.set_mesh(mesh):
-      v = nnx.Variable(
-        jnp.ones((4, 4)),
-        sharding_names=('row', 'col'),
-      )
-    self.assertEqual(v.sharding.mesh, mesh)
-    self.assertEqual(
-      v.sharding.spec,
-      P('row', 'col'),
-    )
+      self.assertEqual(v.sharding.mesh, mesh)
+      self.assertEqual(v.sharding.spec, P('row', 'col'))
 
 def has_sharding_spec(array):
     sharding = array.sharding