Repository: jax-ml/jax
Pull request #25764 (2025-01-07 23:48:47)

[sharding_in_types] Aval sharding under full auto mode should contain None and not UNCONSTRAINED because axis_types + pspec give the full picture.
[sharding_in_types] Aval sharding under full auto mode should contain None and not UNCONSTRAINED because axis_types + pspec give the full picture.



Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 0868990e1..5017f7f5f 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -39,7 +39,7 @@ from jax._src import config
 from jax._src import effects
 from jax._src import compute_on
 from jax._src import mesh as mesh_lib
-from jax._src.partition_spec import PartitionSpec as P, UnconstrainedSingleton
+from jax._src.partition_spec import PartitionSpec as P
 from jax._src.errors import (
     ConcretizationTypeError, TracerArrayConversionError, TracerBoolConversionError,
     TracerIntegerConversionError, UnexpectedTracerError)
@@ -1659,12 +1659,12 @@ def _maybe_modify_sharding(sharding):
 
   new_spec = []
   for s in sharding.spec:
-    if s is None or isinstance(s, UnconstrainedSingleton):
+    if s is None:
       new_spec.append(s)
     else:
       temp_s = s[0] if isinstance(s, tuple) else s
       new_spec.append(
-          P.UNCONSTRAINED
+          None
           if sharding.mesh._name_to_type[temp_s] == mesh_lib.AxisTypes.Auto else s)
   return sharding.with_spec(new_spec)
 
@@ -1762,8 +1762,6 @@ def _get_shape_sharding_str(shape, spec):
   for s1, s2 in zip(shape, spec):
     if s2 is None:
       out.append(f"{s1}")
-    elif isinstance(s2, UnconstrainedSingleton):
-      out.append(f"{s1}")
     elif isinstance(s2, tuple):
       ss = ','.join(s for s in s2)
       out.append(f"{s1}@({ss})")
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index c6a634a05..8e3372113 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -50,7 +50,6 @@ from jax._src.interpreters import xla
 from jax._src.layout import AutoLayout, DeviceLocalLayout
 from jax._src.sharding import Sharding as JSharding
 from jax._src.sharding_impls import AUTO, NamedSharding
-from jax._src.partition_spec import UnconstrainedSingleton
 from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_extension
 from jax._src.lib.mlir import dialects, ir, passmanager
@@ -2599,8 +2598,7 @@ def lower_sharding_under_shit(ctx, op, aval, sharding_proto=None):
   if aval.sharding.mesh._any_axis_collective:
     unspecified_dims = set(range(aval.ndim))
   elif aval.sharding.mesh._any_axis_auto:
-    unspecified_dims = {i for i, s in enumerate(aval.sharding.spec)
-                        if isinstance(s, UnconstrainedSingleton)}
+    unspecified_dims = {i for i, s in enumerate(aval.sharding.spec) if s is None}
   return wrap_with_sharding_op(ctx, op, aval, proto, unspecified_dims)
 
 
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index 807dd0afe..2d9d9857f 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -2150,7 +2150,7 @@ def _discharge_refs_jaxpr(closed_jaxpr, in_shardings, in_layouts,
   return (closed_jaxpr, inout_aliases, mut, in_shardings, in_layouts,
           donated_invars, out_shardings, out_layouts)
 
-def _concretize_abstract_shardings(shardings, avals, device_assignment):
+def _concretize_abstract_out_shardings(shardings, avals, device_assignment):
   np_dev = np.vectorize(lambda i: device_assignment[i],
                         otypes=[object])(np.arange(len(device_assignment)))
 
@@ -2163,8 +2163,14 @@ def _concretize_abstract_shardings(shardings, avals, device_assignment):
   out = []
   for s, a in zip(shardings, avals):
     if isinstance(s, UnspecifiedValue) and a.sharding is not None:
-      out.append(NamedSharding(_abstract_to_concrete_mesh(a.sharding.mesh),
-                               a.sharding.spec))
+      if config.use_shardy_partitioner.value:
+        spec = a.sharding.spec
+      else:
+        spec = (PartitionSpec(*[PartitionSpec.UNCONSTRAINED if sp is None else sp
+                               for sp in a.sharding.spec])
+                if a.sharding.mesh._any_axis_auto else a.sharding.spec)
+      out.append(NamedSharding(
+          _abstract_to_concrete_mesh(a.sharding.mesh), spec))
     else:
       out.append(s)
   return tuple(out)
@@ -2243,7 +2249,7 @@ def lower_sharding_computation(
   unique_intermediate_shardings = [js for js, _ in unique_intermediate_shardings]
 
   if config.sharding_in_types.value:
-    out_shardings = _concretize_abstract_shardings(
+    out_shardings = _concretize_abstract_out_shardings(
         out_shardings, global_out_avals, device_assignment)
 
   # TODO(parkers): One _raw_platform has been unified with platform,


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 63620e5ad..77bcee296 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5589,11 +5589,11 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     @jax.jit
     def f(x, x2):
       y = x * 2
-      self.assertEqual(y.sharding.spec, P(P.UNCONSTRAINED, None))
+      self.assertEqual(y.sharding.spec, P(None, None))
       z = jnp.sin(y)
-      self.assertEqual(z.sharding.spec, P(P.UNCONSTRAINED, None))
+      self.assertEqual(z.sharding.spec, P(None, None))
       a = z @ x2
-      self.assertEqual(a.sharding.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+      self.assertEqual(a.sharding.spec, P(None, None))
       return a
 
     out = f(arr, arr.T)
@@ -5626,9 +5626,9 @@ class ShardingInTypesTest(jtu.JaxTestCase):
       arr = jax.device_put(arr, NamedSharding(mesh2, P('x', 'y')))
       arr2 = jax.device_put(np_inp.T, NamedSharding(mesh2, P('y', None)))
       out = f(arr, arr2)
-      self.assertEqual(out.sharding, NamedSharding(mesh2, P('x', None)))
+      self.assertEqual(out.sharding, NamedSharding(mesh2, P('x',)))
       lowered_text = f.lower(arr, arr2).as_text()
-      self.assertTrue(lowered_text.count("unspecified_dims") == 3)
+      self.assertTrue(lowered_text.count("unspecified_dims") == 5)
 
     mesh3 = jtu.create_mesh((2, 2), ('x', 'y'),
                             axis_types={mesh_lib.AxisTypes.User: 'y',
@@ -5669,11 +5669,11 @@ class ShardingInTypesTest(jtu.JaxTestCase):
           {mesh_lib.AxisTypes.Auto: ('x', 'y')})
       y = sharding_cast(y, y.sharding.with_mesh(auto_mesh))
       with mesh_lib.set_abstract_mesh(auto_mesh):
-        self.assertEqual(y.sharding.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+        self.assertEqual(y.sharding.spec, P(None, None))
         z = jnp.sin(y)
-        self.assertEqual(z.sharding.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+        self.assertEqual(z.sharding.spec, P(None, None))
         a = z @ z.T
-        self.assertEqual(a.sharding.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+        self.assertEqual(a.sharding.spec, P(None, None))
       a = sharding_cast(
           a, NamedSharding(mesh_lib.get_abstract_mesh(), P('x', None)))
       self.assertEqual(a.sharding.spec, P('x', None))
@@ -5707,7 +5707,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
         self.assertEqual(a.sharding.spec, P(None, None))
       a = sharding_cast(
           a, NamedSharding(mesh_lib.get_abstract_mesh(), P('x', None)))
-      self.assertEqual(a.sharding.spec, P(P.UNCONSTRAINED, None))
+      self.assertEqual(a.sharding.spec, P(None, None))
       return a
 
     out = f(arr)
@@ -5729,11 +5729,11 @@ class ShardingInTypesTest(jtu.JaxTestCase):
           {mesh_lib.AxisTypes.Auto: 'x', mesh_lib.AxisTypes.User: 'y'})
       y = sharding_cast(y, y.sharding.with_mesh(mix_mesh))
       with mesh_lib.set_abstract_mesh(mix_mesh):
-        self.assertEqual(y.sharding.spec, P(P.UNCONSTRAINED, 'y'))
+        self.assertEqual(y.sharding.spec, P(None, 'y'))
         z = jnp.sin(y)
-        self.assertEqual(z.sharding.spec, P(P.UNCONSTRAINED, 'y'))
+        self.assertEqual(z.sharding.spec, P(None, 'y'))
         a = z @ z.T
-        self.assertEqual(a.sharding.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+        self.assertEqual(a.sharding.spec, P(None, None))
       a = sharding_cast(
           a, NamedSharding(mesh_lib.get_abstract_mesh(), P('x', None)))
       self.assertEqual(a.sharding.spec, P('x', None))