Repository: jax-ml/jax
Pull request #27975 (2025-04-12 00:15:09)

Rename is_closed to is_open in the shardy shardings
Rename is_closed to is_open in the shardy shardings



Patch:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index 20cdec781..630539ab0 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -167,7 +167,7 @@ def _callback_op_sharding(
           sharding_impls.SdyArraySharding(
               mesh_shape=(),
               dimension_shardings=[
-                  sharding_impls.SdyDimSharding(axes=[], is_closed=True)
+                  sharding_impls.SdyDimSharding(axes=[], is_open=False)
               ] * avals_out[0].ndim,
               logical_device_ids=())])
     else:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 79b67219b..dc140c226 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -170,7 +170,7 @@ def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params)
             sharding_impls.SdyArraySharding(
                 mesh_shape=(),
                 dimension_shardings=[
-                    sharding_impls.SdyDimSharding(axes=[], is_closed=True)
+                    sharding_impls.SdyDimSharding(axes=[], is_open=False)
                 ] * ctx.avals_out[0].ndim,
                 logical_device_ids=())])
       else:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 84bbd6852..becdfd46b 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1795,7 +1795,7 @@ def replicate_trailing_dims(ctx, val: ir.Value, aval) -> ir.Value:
     s = SdyArraySharding(
         mesh_shape=None,
         dimension_shardings=[
-            sharding_impls.SdyDimSharding(axes=[], is_closed=i >= aval.ndim)
+            sharding_impls.SdyDimSharding(axes=[], is_open=i < aval.ndim)
             for i in range(physical_ndim)
         ])
     return wrap_with_sharding_op(ctx, val, aval, s)
diff --git a/jax/_src/named_sharding.py b/jax/_src/named_sharding.py
index b712e69dd..9bdadd8a8 100644
--- a/jax/_src/named_sharding.py
+++ b/jax/_src/named_sharding.py
@@ -42,7 +42,7 @@ class AUTO:
     self.mesh = mesh
 
   def _to_sdy_sharding(self, ndim: int) -> SdyArraySharding:
-    dim_shardings = [SdyDimSharding(axes=[], is_closed=False)
+    dim_shardings = [SdyDimSharding(axes=[], is_open=True)
                      for _ in range(ndim)]
     return SdyArraySharding(self.mesh.shape_tuple, dim_shardings)
 
@@ -242,11 +242,11 @@ class NamedSharding(JSharding.Sharding):
     return named_sharding_to_xla_hlo_sharding(self, num_dimensions)
 
   def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
-    dim_shardings = [SdyDimSharding(axes=[], is_closed=True)
+    dim_shardings = [SdyDimSharding(axes=[], is_open=False)
                      for _ in range(num_dimensions)]
     for i, dim_spec in enumerate(self.spec):
       if dim_spec is PartitionSpec.UNCONSTRAINED:
-        dim_shardings[i].is_closed = False
+        dim_shardings[i].is_open = True
       elif dim_spec is None:
         # Already empty and closed sharding.
         pass
@@ -274,14 +274,13 @@ def get_array_mapping(
 @dataclasses.dataclass
 class SdyDimSharding:
   axes: Sequence[str]
-  is_closed: bool
+  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=self.is_closed,
-        priority=self.priority)
+        is_closed=not self.is_open, priority=self.priority)
 
   def __repr__(self):
     return f'SdyDimSharding({self._custom_repr()})'
@@ -289,7 +288,7 @@ class SdyDimSharding:
   def _custom_repr(self):
     axes_repr = ', '.join(f"'{a}'" for a in self.axes)
     open_repr = ''
-    if not self.is_closed:
+    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}'
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index d95b12f24..e462ee2f0 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -101,8 +101,8 @@ def modify_sdy_sharding_wrt_axis_types(sdy_sharding: SdyArraySharding, mesh):
     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(SdyDimSharding(axes=[], is_closed=False)
-                           if not d.axes and d.is_closed else d)
+      dim_shardings.append(SdyDimSharding(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
@@ -185,7 +185,7 @@ class SingleDeviceSharding(jsharding.Sharding):
     return replicated_hlo_sharding
 
   def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
-    sdy_dim_sharding = [SdyDimSharding(axes=[], is_closed=True)
+    sdy_dim_sharding = [SdyDimSharding(axes=[], is_open=False)
                         for _ in range(num_dimensions)]
     return SdyArraySharding(None, sdy_dim_sharding)
 
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index 2d114c6c3..a6998d29e 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -706,7 +706,7 @@ def _shardy_shard_map_sharding(
     for dim_sharding in sdy_sharding.dimension_shardings:
       # Only allow dimensions which have no sharding to be auto-sharded.
       if not dim_sharding.axes:
-        dim_sharding.is_closed = False
+        dim_sharding.is_open = True
   return sdy_sharding
 
 


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 87227d4d6..1780213bc 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1464,9 +1464,9 @@ class ShardyShardingTest(jtu.JaxTestCase):
         SdyArraySharding(
             mesh.shape_tuple,
             [SdyDimSharding(
-             ('sequence', 'data'), True),
-             SdyDimSharding(('model',), True),
-             SdyDimSharding([], True)]))
+             ('sequence', 'data'), False),
+             SdyDimSharding(('model',), False),
+             SdyDimSharding([], False)]))
     with ir.Context() as ctx:
       dialects.sdy.register_dialect(ctx)
       self.assertEqual(
@@ -1483,9 +1483,9 @@ class ShardyShardingTest(jtu.JaxTestCase):
         sdy_sharding,
         SdyArraySharding(
             mesh.shape_tuple,
-            [SdyDimSharding([], True),
-             SdyDimSharding([], False),
-             SdyDimSharding(('x',), True)]))
+            [SdyDimSharding([], False),
+             SdyDimSharding([], True),
+             SdyDimSharding(('x',), False)]))
     with ir.Context() as ctx:
       dialects.sdy.register_dialect(ctx)
       self.assertEqual(
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index f90b02cbd..0e7867cc2 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -8188,8 +8188,8 @@ class ShardyTest(jtu.JaxTestCase):
     sharding = sharding_impls.SdyArraySharding(
         mesh_shape=(('data', 4), ('model', 8), ('expert', 2)),
         dimension_shardings=[
-            sharding_impls.SdyDimSharding(axes=['data', 'expert'], is_closed=True),
-            sharding_impls.SdyDimSharding(axes=['model'], is_closed=False, priority=2)])
+            sharding_impls.SdyDimSharding(axes=['data', 'expert'], is_open=False),
+            sharding_impls.SdyDimSharding(axes=['model'], is_open=True, priority=2)])
     self.assertEqual(repr(sharding), "SdyArraySharding([{'data', 'expert'}, {'model', ?}p2])")
 
   def test_array_sharding_repr_with_logical_ids(self):
@@ -8202,7 +8202,7 @@ class ShardyTest(jtu.JaxTestCase):
 
   def test_dimension_sharding_repr(self):
     dim_sharding = sharding_impls.SdyDimSharding(
-        axes=['data', 'model'], is_closed=False, priority=2)
+        axes=['data', 'model'], is_open=True, priority=2)
     self.assertEqual(repr(dim_sharding),
                      "SdyDimSharding({'data', 'model', ?}p2)")