Repository: jax-ml/jax
Pull request #30083 (2025-07-09 03:19:31)

Allow `P(('a', 'b'))` in aval.sharding where mesh is `('a', 2, Explicit), ('b', 2, Auto)`.
Allow `P(('a', 'b'))` in aval.sharding where mesh is `('a', 2, Explicit), ('b', 2, Auto)`.

This will lower to `P('a', None)` for ndim=2. But `P(('a', 'b'))` where `a=Explicit, b=Manual` is not allowed.

This is required to make shard_map/smap partial manual **eager mode** work correctly.

Co-authored-by: Matthew Johnson <mattjj@google.com>



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 344c8b2b3..570008e65 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1980,17 +1980,16 @@ def _make_lengths_same(sharding, ndim):
     return sharding.update(spec=P(*pspec[:ndim], unreduced=pspec.unreduced))
   assert False, "unreachable"
 
-# TODO(yashkatariya): Only works with User/Auto. Generalize it to work with
-# Collective too.
 def modify_spec_for_auto_manual(spec, mesh) -> P:
-  new_spec = []
+  new_spec = []  # type: ignore
   for s in spec:
-    if not s:
-      new_spec.append(s)
+    if s is None:
+      new_spec.append(s)  # type: ignore
+    elif isinstance(s, tuple):
+      new_spec.append(tuple(
+          p for p in s if mesh._name_to_type[p] == AxisType.Explicit))
     else:
-      temp_s = s[0] if isinstance(s, tuple) else s
-      new_spec.append(s if mesh._name_to_type[temp_s] == AxisType.Explicit
-                      else None)
+      new_spec.append(s if mesh._name_to_type[s] == AxisType.Explicit else None)  # type: ignore
   new_unreduced = {u for u in spec.unreduced
                    if mesh._name_to_type[u] == AxisType.Explicit}
   new_reduced = {u for u in spec.reduced
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index df54f1ab9..ae543ce3d 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1137,15 +1137,11 @@ def all_unconstrained(s, aval):
 class UnconstrainedVariants(NamedTuple):
   contains_unconstrained: bool
   all_unconstrained: bool
-  unconstrained_dims: set[int] | None
 
 def _get_unconstrained_variants(s, aval) -> UnconstrainedVariants:
   us = contains_unconstrained(s)
-  unconstrained_dims = ({i for i, p in enumerate(s.spec)  # pytype: disable=attribute-error
-                         if p is PartitionSpec.UNCONSTRAINED} if us else None)
   return UnconstrainedVariants(
-      contains_unconstrained=us, all_unconstrained=all_unconstrained(s, aval),
-      unconstrained_dims=unconstrained_dims)
+      contains_unconstrained=us, all_unconstrained=all_unconstrained(s, aval))
 
 
 def check_jaxpr_constants(closed_jaxpr: core.ClosedJaxpr):
@@ -1842,9 +1838,14 @@ def lower_jaxpr_to_fun(
             not uv.all_unconstrained):
           if config.use_shardy_partitioner.value:
             s = modify_sdy_sharding_wrt_axis_types(s, o_aval.sharding.mesh)
+            unconstrained_dims = None  # delete this after shardy is default
+          else:
+            unconstrained_dims = (
+                set(range(o_aval.ndim)) if o_aval.sharding.mesh._any_axis_auto
+                else None)
           temp_flat_outputs.append(wrap_with_sharding_op(
               entry_lowering_ctx, o, o_aval, s,
-              unspecified_dims=uv.unconstrained_dims))
+              unspecified_dims=unconstrained_dims))
         else:
           temp_flat_outputs.append(o)
       flat_outputs = temp_flat_outputs
@@ -2898,9 +2899,7 @@ def lower_with_sharding_in_types(ctx, op, aval, sharding_proto=None):
             if sharding_proto is None else sharding_proto)
     unspecified_dims = None
     if aval.sharding.mesh._any_axis_auto:
-      # TODO(yashkatariya): Maybe if any mesh axis is auto, mark all axes
-      # as unspecified?
-      unspecified_dims = {i for i, s in enumerate(aval.sharding.spec) if s is None}
+      unspecified_dims = set(range(aval.ndim))
     return wrap_with_sharding_op(ctx, op, aval, proto, unspecified_dims)
 
 
diff --git a/jax/_src/named_sharding.py b/jax/_src/named_sharding.py
index 1194c1795..5c60afe56 100644
--- a/jax/_src/named_sharding.py
+++ b/jax/_src/named_sharding.py
@@ -283,12 +283,11 @@ def get_array_mapping(
 class SdyDim:
   axes: Sequence[str]
   is_open: bool
-  priority: int | None = None
 
   def build(self) -> sdy.DimensionShardingAttr:
     return sdy.DimensionShardingAttr.get(
         [sdy.AxisRefAttr.get(axis) for axis in self.axes],
-        is_closed=not self.is_open, priority=self.priority)
+        is_closed=not self.is_open)
 
   def __repr__(self):
     return f'SdyDim({self._custom_repr()})'
@@ -298,8 +297,7 @@ class SdyDim:
     open_repr = ''
     if self.is_open:
       open_repr = ', ?' if self.axes else '?'
-    priority_repr = '' if self.priority is None else f'p{self.priority}'
-    return f'{{{axes_repr}{open_repr}}}{priority_repr}'
+    return f'{{{axes_repr}{open_repr}}}'
 
 def _get_axes(axes, mesh_shape):
   if not axes:
@@ -351,9 +349,7 @@ def modify_sdy_sharding_wrt_axis_types(sdy_sharding: SdyArray, mesh):
   if mesh._any_axis_auto:
     dim_shardings, used_axes = [], []  # type: ignore
     for d in sdy_sharding.dim_shardings:
-      # TODO(yashkatariya): Maybe if any mesh axis is auto, mark all axes as open?
-      dim_shardings.append(SdyDim(axes=[], is_open=True)
-                           if not d.axes and not d.is_open else d)
+      dim_shardings.append(SdyDim(axes=d.axes, is_open=True))
       used_axes.extend(d.axes)
     remaining_axes = set(mesh.axis_names) - set(used_axes)
     replicated_axes = tuple(r for r in remaining_axes
@@ -498,6 +494,19 @@ def _check_unique_resources(pspec: PartitionSpec, arg_name: str, mesh=None
             f' for {mesh_lib.show_axes(multiple_uses)}'),
         mesh=mesh, pspec=pspec)
 
+def check_pspec_mix_axis_type(mesh, pspec):
+  for spec in pspec:
+    if isinstance(spec, tuple):
+      if all(mesh._name_to_type[spec[0]] == mesh._name_to_type[p]
+             for p in spec):
+        continue
+      if any(mesh._name_to_type[p] == AxisType.Manual for p in spec):
+        raise ValueError(
+            'Tuple subset of `PartitionSpec` cannot contain `Manual` mixed'
+            f' with `Auto` or `Explicit`. Got pspec {pspec} and subset'
+            f' {spec} with axis types:'
+            f' ({", ".join(str(mesh._name_to_type[p]) for p in spec)})')
+
 def _check_mesh_resource_axis(mesh, pspec):
   for p in pspec:
     if p is PartitionSpec.UNCONSTRAINED or p is None:
@@ -508,11 +517,7 @@ def _check_mesh_resource_axis(mesh, pspec):
         raise ValueError(
             f"Resource axis: {r} of {pspec} "
             f"is not found in mesh: {tuple(mesh.shape.keys())}.")
-    if not all(mesh._name_to_type[p[0]] == mesh._name_to_type[r] for r in p):
-      raise ValueError(
-          'AxisTypes should be the same in a tuple subset of PartitionSpec:'
-          f' {pspec}. Got subset {p} with axis'
-          f' types: ({", ".join(str(mesh._name_to_type[r]) for r in p)})')
+  check_pspec_mix_axis_type(mesh, pspec)
   if (AxisType.Auto not in mesh._axis_types_dict and
       PartitionSpec.UNCONSTRAINED in pspec):
     raise ValueError(


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index b98f25abc..7bdd3e5fb 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1336,6 +1336,28 @@ class ShardingTest(jtu.JaxTestCase):
         ValueError, 'Got invalid memory kind'):
       NamedSharding(abstract_mesh, P(), memory_kind='weird_device')
 
+  def test_pspec_mix_axis_types(self):
+    mesh = AbstractMesh(
+        (2, 2, 2, 2), ('a', 'b', 'c', 'd'),
+        axis_types=(AxisType.Explicit, AxisType.Explicit, AxisType.Auto,
+                    AxisType.Manual))
+    aval = jax.core.ShapedArray((16, 8, 4, 2), np.float32)
+
+    out = aval.update(sharding=NamedSharding(mesh, P(('a', 'b', 'c'), 'd')))
+    self.assertEqual(out.sharding.spec, P(('a', 'b'), None, None, None))
+
+    out = aval.update(sharding=NamedSharding(mesh, P(('a', 'c'), 'b', 'd')))
+    self.assertEqual(out.sharding.spec, P('a', 'b', None, None))
+
+    out = aval.update(sharding=NamedSharding(mesh, P(('a', 'b'), 'c', 'd')))
+    self.assertEqual(out.sharding.spec, P(('a', 'b'), None, None, None))
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'Tuple subset of `PartitionSpec` cannot contain `Manual` mixed with'
+        ' `Auto` or `Explicit`'):
+      aval.update(sharding=NamedSharding(mesh, P(('a', 'd'), 'b', 'c')))
+
   def test_pspec_unreduced(self):
     pspec = P('a', 'b', None, unreduced={'c'}, reduced={'d'})
     self.assertEqual(
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 404113f1a..eaef1dc5b 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4868,9 +4868,9 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(out.sharding, s)
     lowered_text = g.lower(arr).as_text()
     if config.use_shardy_partitioner.value:
-      self.assertIn('<@mesh, [{?}, {"y"}]>', lowered_text)
+      self.assertIn('<@mesh, [{?}, {"y", ?}]>', lowered_text)
     else:
-      self.assertIn("unspecified_dims=[0]", lowered_text)
+      self.assertIn("unspecified_dims=[0,1]", lowered_text)
 
   def test_prng_key_wsc(self):
     mesh = jtu.create_mesh((2,), 'x')
@@ -6175,11 +6175,6 @@ class ShardingInTypesTest(jtu.JaxTestCase):
       else:
         self.assertTrue(lowered_text.count("unspecified_dims") == 4)
 
-    with self.assertRaisesRegex(
-        ValueError,
-        "AxisTypes should be the same in a tuple subset of PartitionSpec"):
-      NamedSharding(mesh2, P(('x', 'y')))
-
   @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
   def test_where_with_scalar(self, mesh):
     np_inp = np.arange(16.).reshape(8, 2)
@@ -6417,9 +6412,9 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     lowered_text = f.lower(arr).as_text()
     if config.use_shardy_partitioner.value:
       self.assertTrue(lowered_text.count(
-          '[{"x"}, {?}, {?}], replicated={"y"}') == 3)
+          '[{"x", ?}, {?}, {?}], replicated={"y"}') == 3)
     else:
-      self.assertTrue(lowered_text.count("unspecified_dims=[1,2]") == 3)
+      self.assertTrue(lowered_text.count("unspecified_dims=[0,1,2]") == 3)
 
   @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
   def test_auto_mode_mix(self, mesh):
@@ -8218,6 +8213,29 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = jax.jit(f)(jnp.ones((2, 3)))
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
 
+  @jtu.with_explicit_mesh((2, 2), ('x', 'y'),
+                          axis_types=(AxisType.Explicit, AxisType.Auto))
+  def test_mix_axis_type_in_pspec(self, mesh):
+    np_inp = np.arange(16).reshape(8, 2)
+    arr = jax.device_put(np_inp, P(('x', 'y'), None))
+
+    @jax.jit
+    def f(x):
+      self.assertEqual(x.aval.sharding.spec, P('x', None))
+      out = x * 2
+      self.assertEqual(out.aval.sharding.spec, P('x', None))
+      return out
+
+    out = f(arr)
+    self.assertArraysEqual(out, np_inp * 2)
+
+    lowered_text = f.lower(arr).as_text()
+    if config.use_shardy_partitioner.value:
+      self.assertEqual(lowered_text.count("{?}"), 3)
+      self.assertEqual(lowered_text.count('{"x", ?}'), 3)
+    else:
+      self.assertEqual(lowered_text.count('unspecified_dims=[0,1]'), 3)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):
@@ -8989,8 +9007,8 @@ class ShardyTest(jtu.JaxTestCase):
         mesh_shape=(('data', 4), ('model', 8), ('expert', 2)),
         dim_shardings=[
             sharding_impls.SdyDim(axes=['data', 'expert'], is_open=False),
-            sharding_impls.SdyDim(axes=['model'], is_open=True, priority=2)])
-    self.assertEqual(repr(sharding), "SdyArray([{'data', 'expert'}, {'model', ?}p2])")
+            sharding_impls.SdyDim(axes=['model'], is_open=True)])
+    self.assertEqual(repr(sharding), "SdyArray([{'data', 'expert'}, {'model', ?}])")
 
   def test_array_sharding_repr_with_logical_ids(self):
     abstract_mesh = jax.sharding.AbstractMesh((4, 8, 2), ('x', 'y', 'z'))
@@ -9002,9 +9020,9 @@ class ShardyTest(jtu.JaxTestCase):
 
   def test_dimension_sharding_repr(self):
     dim_sharding = sharding_impls.SdyDim(
-        axes=['data', 'model'], is_open=True, priority=2)
+        axes=['data', 'model'], is_open=True)
     self.assertEqual(repr(dim_sharding),
-                     "SdyDim({'data', 'model', ?}p2)")
+                     "SdyDim({'data', 'model', ?})")
 
   def test_tensor_dialect(self):
     # While this doesn't emit any `mlir::TensorDialect` ops, some pass in the