Repository: jax-ml/jax
Pull request #28713 (2025-05-13 17:24:06)

Lower unreduced to shardy. GSPMD doesn't support unreduced.
Lower unreduced to shardy. GSPMD doesn't support unreduced.



Patch:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index 06b36ce5c..bc233b634 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -158,7 +158,7 @@ def _callback_op_sharding(
       op_sharding = sharding_impls.SdyArrayList([
           sharding_impls.SdyArray(
               mesh_shape=(),
-              dimension_shardings=[
+              dim_shardings=[
                   sharding_impls.SdyDim(axes=[], is_open=False)
               ] * avals_out[0].ndim,
               logical_device_ids=())])
@@ -200,7 +200,7 @@ def _callback_op_sharding(
       op_sharding = sharding_impls.SdyArrayList(num_sdy_shardings * [
           sharding_impls.SdyArray(
               mesh_shape=(),
-              dimension_shardings=[],
+              dim_shardings=[],
               logical_device_ids=(device_index,))])
     else:
       op_sharding = xc.OpSharding()  # type: ignore[assignment]
@@ -610,7 +610,7 @@ def send_to_host(
       assert len(sharding.shardings) >= 1
       sharding = SdyArrayList([
           SdyArray(
-              mesh_shape=(), dimension_shardings=[],
+              mesh_shape=(), dim_shardings=[],
               logical_device_ids=sharding.shardings[0].logical_device_ids)])
     mlir.set_sharding(send_op, sharding)
   return send_op.result
@@ -645,7 +645,7 @@ def receive_from_host(
       sharding = SdyArrayList([
           sharding.shardings[0],
           SdyArray(
-              mesh_shape=(), dimension_shardings=[],
+              mesh_shape=(), dim_shardings=[],
               logical_device_ids=sharding.shardings[0].logical_device_ids)])
     mlir.set_sharding(recv_op, sharding)
   # Token should be at the end of the results
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 63abcbef3..c2febf752 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -168,7 +168,7 @@ def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params)
         sharding = sharding_impls.SdyArrayList([
             sharding_impls.SdyArray(
                 mesh_shape=(),
-                dimension_shardings=[
+                dim_shardings=[
                     sharding_impls.SdyDim(axes=[], is_open=False)
                 ] * ctx.avals_out[0].ndim,
                 logical_device_ids=())])
@@ -184,7 +184,7 @@ def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params)
     if config.use_shardy_partitioner.value:
       sharding = sharding_impls.SdyArrayList([
           sharding_impls.SdyArray(
-              mesh_shape=(), dimension_shardings=[], logical_device_ids=(0,))])
+              mesh_shape=(), dim_shardings=[], logical_device_ids=(0,))])
     else:
       sharding = xc.OpSharding()
       sharding.type = xc.OpSharding.Type.MAXIMAL
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 94418f0b9..0256057b8 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1841,7 +1841,7 @@ def replicate_trailing_dims(ctx, val: ir.Value, aval) -> ir.Value:
     physical_ndim = core.physical_aval(aval).ndim
     s = SdyArray(
         mesh_shape=None,
-        dimension_shardings=[
+        dim_shardings=[
             sharding_impls.SdyDim(axes=[], is_open=i < aval.ndim)
             for i in range(physical_ndim)
         ])
diff --git a/jax/_src/named_sharding.py b/jax/_src/named_sharding.py
index 45ae1c124..faf0b2a9f 100644
--- a/jax/_src/named_sharding.py
+++ b/jax/_src/named_sharding.py
@@ -44,7 +44,8 @@ class AUTO:
   def _to_sdy_sharding(self, ndim: int) -> SdyArray:
     dim_shardings = [SdyDim(axes=[], is_open=True)
                      for _ in range(ndim)]
-    return SdyArray(self.mesh.shape_tuple, dim_shardings)
+    return SdyArray(mesh_shape=self.mesh.shape_tuple,
+                    dim_shardings=dim_shardings)
 
 class UnspecifiedValue:
   def __repr__(self):
@@ -244,8 +245,10 @@ class NamedSharding(JSharding.Sharding):
       else:
         dim_spec = dim_spec if isinstance(dim_spec, tuple) else (dim_spec,)
         dim_shardings[i].axes = dim_spec
-    return SdyArray(self.mesh.shape_tuple, dim_shardings,
-                    self._logical_device_ids)
+    return SdyArray(mesh_shape=self.mesh.shape_tuple,
+                    dim_shardings=dim_shardings,
+                    logical_device_ids=self._logical_device_ids,
+                    unreduced_axes=self.spec.unreduced)
 
 NamedSharding.__module__ = 'jax.sharding'
 
@@ -285,13 +288,21 @@ class SdyDim:
     priority_repr = '' if self.priority is None else f'p{self.priority}'
     return f'{{{axes_repr}{open_repr}}}{priority_repr}'
 
+def _get_axes(axes, mesh_shape):
+  if not axes:
+    return ()
+  assert mesh_shape is not None
+  # Sort wrt mesh axis names so order is deterministic and doesn't hang in
+  # McJAX.
+  return tuple(n for n, _ in mesh_shape if n in axes)
 
-@dataclasses.dataclass
+@dataclasses.dataclass(kw_only=True)
 class SdyArray:
   mesh_shape: tuple[tuple[str, int], ...] | None
-  dimension_shardings: Sequence[SdyDim]
+  dim_shardings: Sequence[SdyDim]
   logical_device_ids: tuple[int, ...] | None = None
   replicated_axes: tuple[str, ...] = ()
+  unreduced_axes: tuple[str, ...] = ()
 
   def build(self) -> sdy.TensorShardingAttr:
     if self.mesh_shape is None:
@@ -302,14 +313,18 @@ class SdyArray:
       mesh_attr = sdy.MeshAttr.get(
           [sdy.MeshAxisAttr.get(name, size) for name, size in self.mesh_shape],
           ldi)
+
+    replicated_axes = _get_axes(self.replicated_axes, self.mesh_shape)
+    unreduced_axes = _get_axes(self.unreduced_axes, self.mesh_shape)
     return sdy.TensorShardingAttr.get(
         mesh_attr,
-        [dim_sharding.build() for dim_sharding in self.dimension_shardings],
-        replicated_axes=[sdy.AxisRefAttr.get(axis) for axis in self.replicated_axes])
+        [dim_sharding.build() for dim_sharding in self.dim_shardings],
+        replicated_axes=[sdy.AxisRefAttr.get(axis) for axis in replicated_axes],
+        unreduced_axes=[sdy.AxisRefAttr.get(axis) for axis in unreduced_axes])
 
   def __repr__(self):
     dim_sharding_repr = ', '.join(
-        d._custom_repr() for d in self.dimension_shardings)
+        d._custom_repr() for d in self.dim_shardings)
     device_id_repr = (f', device_ids={self.logical_device_ids}'
                       if self.logical_device_ids is not None else '')
     rar = (f', replicated_axes={self.replicated_axes}'
@@ -317,6 +332,26 @@ class SdyArray:
     return f"SdyArray([{dim_sharding_repr}]{device_id_repr}{rar})"
 
 
+# TODO(yashkatariya): Upstream this into `_to_sdy_sharding` maybe with an extra
+# parameter to it `_to_sdy_sharding(self, ndim, modify_wrt_axis_types=False)`
+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)
+      used_axes.extend(d.axes)
+    remaining_axes = set(mesh.axis_names) - set(used_axes)
+    replicated_axes = tuple(r for r in remaining_axes
+                            if mesh._name_to_type[r] == mesh_lib.AxisType.Explicit)
+    return SdyArray(mesh_shape=sdy_sharding.mesh_shape,
+                    dim_shardings=dim_shardings,
+                    logical_device_ids=sdy_sharding.logical_device_ids,
+                    replicated_axes=replicated_axes)
+  return sdy_sharding
+
+
 @cache(max_size=4096, trace_context_in_key=False)
 def named_sharding_to_xla_hlo_sharding(
     self, num_dimensions: int) -> xc.HloSharding:
diff --git a/jax/_src/shard_map.py b/jax/_src/shard_map.py
index 3eb46da89..abcc2ca0a 100644
--- a/jax/_src/shard_map.py
+++ b/jax/_src/shard_map.py
@@ -781,7 +781,7 @@ def _shardy_shard_map_sharding(
     aval_in = core.physical_aval(aval_in)
   sdy_sharding = ns._to_sdy_sharding(aval_in.ndim)
   if len(manual_axes) < len(mesh.axis_names):
-    for dim_sharding in sdy_sharding.dimension_shardings:
+    for dim_sharding in sdy_sharding.dim_shardings:
       dim_sharding.is_open = True
   return sdy_sharding
 
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index f7f0ebd2c..982af82c5 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -38,7 +38,8 @@ from jax._src.named_sharding import (  # noqa: F401
     SdyArray, SdyDim, UnspecifiedValue, AUTO,
     _check_unique_resources, NamedSharding, UNSPECIFIED,
     ArrayMapping, ArrayMappingOrAutoOrUnspecified, get_array_mapping,
-    array_mapping_to_axis_resources, named_sharding_to_xla_hlo_sharding)
+    array_mapping_to_axis_resources, named_sharding_to_xla_hlo_sharding,
+    modify_sdy_sharding_wrt_axis_types)
 from jax._src.op_shardings import (
     are_op_shardings_equal, get_num_ways_dim_sharded, is_op_sharding_replicated)
 from jax._src.partition_spec import PartitionSpec
@@ -95,27 +96,6 @@ class SdyArrayList:
         [sharding.build() for sharding in self.shardings])
 
 
-# TODO(yashkatariya): Upstream this into `_to_sdy_sharding` maybe with an extra
-# parameter to it `_to_sdy_sharding(self, ndim, modify_wrt_axis_types=False)`
-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.dimension_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)
-      used_axes.extend(d.axes)
-    remaining_axes = set(mesh.axis_names) - set(used_axes)
-    # Sort wrt mesh axis names so order is deterministic and doesn't hang in
-    # McJAX.
-    remaining_axes = [n for n in mesh.axis_names if n in remaining_axes]
-    replicated_axes = tuple(r for r in remaining_axes
-                            if mesh._name_to_type[r] == mesh_lib.AxisType.Explicit)
-    return SdyArray(sdy_sharding.mesh_shape, dim_shardings,
-                    sdy_sharding.logical_device_ids, replicated_axes)
-  return sdy_sharding
-
-
 replicated_hlo_sharding = xc.HloSharding.replicate()
 
 
@@ -188,7 +168,7 @@ class SingleDeviceSharding(jsharding.Sharding):
   def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
     sdy_dim_sharding = [SdyDim(axes=[], is_open=False)
                         for _ in range(num_dimensions)]
-    return SdyArray(None, sdy_dim_sharding)
+    return SdyArray(mesh_shape=None, dim_shardings=sdy_dim_sharding)
 
   @property
   def is_fully_replicated(self) -> bool:


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 0bab37c07..1691c3acc 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1477,8 +1477,8 @@ class ShardyShardingTest(jtu.JaxTestCase):
     self.assertEqual(
         sdy_sharding,
         SdyArray(
-            mesh.shape_tuple,
-            [SdyDim(
+            mesh_shape=mesh.shape_tuple,
+            dim_shardings=[SdyDim(
              ('sequence', 'data'), False),
              SdyDim(('model',), False),
              SdyDim([], False)]))
@@ -1497,8 +1497,8 @@ class ShardyShardingTest(jtu.JaxTestCase):
     self.assertEqual(
         sdy_sharding,
         SdyArray(
-            mesh.shape_tuple,
-            [SdyDim([], False),
+            mesh_shape=mesh.shape_tuple,
+            dim_shardings=[SdyDim([], False),
              SdyDim([], True),
              SdyDim(('x',), False)]))
     with ir.Context() as ctx:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 9b658dd8c..e374b0b15 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7735,6 +7735,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     f(arr)  # doesn't crash
     jax.jit(f)(arr)  # doesn't crash
 
+  @config.use_shardy_partitioner(True)
   @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
   def test_unreduced_basic(self, mesh):
     np_inp = np.arange(16).reshape(8, 2)
@@ -7758,7 +7759,10 @@ class ShardingInTypesTest(jtu.JaxTestCase):
       self.assertEqual(out.aval.sharding.spec, P('x', None))
       return out
 
-    f.trace(x, y, a, b)  # doesn't crash
+    traced = f.trace(x, y, a, b)
+    lowered_text = traced.lower().as_text()
+    self.assertIn('unreduced={"y"}', lowered_text)
+    self.assertTrue(lowered_text.count('unreduced={"y"}') == 3)
 
   @jtu.with_explicit_mesh((2, 2, 1), ('x', 'y', 'z'))
   def test_dot_general_unreduced_error(self, mesh):
@@ -8595,7 +8599,7 @@ class ShardyTest(jtu.JaxTestCase):
   def test_array_sharding_repr_with_priority(self):
     sharding = sharding_impls.SdyArray(
         mesh_shape=(('data', 4), ('model', 8), ('expert', 2)),
-        dimension_shardings=[
+        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])")