Repository: jax-ml/jax
Pull request #28739 (2025-05-14 14:08:03)

Rename `SdyArraySharding -> SdyArray` and `SdyDimSharding -> SdyDim` since these are not `Sharding` from a JAX POV.
Rename `SdyArraySharding -> SdyArray` and `SdyDimSharding -> SdyDim` since these are not `Sharding` from a JAX POV.



Patch:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index d23389af1..06b36ce5c 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -40,7 +40,7 @@ from jax._src.lax.control_flow.loops import map as lax_map
 from jax._src.lib import xla_client as xc
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
-from jax._src.sharding_impls import SdyArraySharding, SdyArrayShardingList, SingleDeviceSharding
+from jax._src.sharding_impls import SdyArray, SdyArrayList, SingleDeviceSharding
 from jax._src.typing import DeprecatedArg
 import numpy as np
 
@@ -155,11 +155,11 @@ def _callback_op_sharding(
       )
     if config.use_shardy_partitioner.value:
       assert len(avals_out) == 1
-      op_sharding = sharding_impls.SdyArrayShardingList([
-          sharding_impls.SdyArraySharding(
+      op_sharding = sharding_impls.SdyArrayList([
+          sharding_impls.SdyArray(
               mesh_shape=(),
               dimension_shardings=[
-                  sharding_impls.SdyDimSharding(axes=[], is_open=False)
+                  sharding_impls.SdyDim(axes=[], is_open=False)
               ] * avals_out[0].ndim,
               logical_device_ids=())])
     else:
@@ -197,8 +197,8 @@ def _callback_op_sharding(
       # number of result ops. If there are no result ops, we need 1 shardy
       # annotation.
       num_sdy_shardings = max(1, len(avals_out))
-      op_sharding = sharding_impls.SdyArrayShardingList(num_sdy_shardings * [
-          sharding_impls.SdyArraySharding(
+      op_sharding = sharding_impls.SdyArrayList(num_sdy_shardings * [
+          sharding_impls.SdyArray(
               mesh_shape=(),
               dimension_shardings=[],
               logical_device_ids=(device_index,))])
@@ -590,7 +590,7 @@ def send_to_host(
     operand: Any,
     name: str,
     *,
-    sharding: SdyArrayShardingList | xc.OpSharding | None = None,
+    sharding: SdyArrayList | xc.OpSharding | None = None,
 ) -> ir.Value:
   channel_handle = hlo.ChannelHandle.get(channel, mlir.SEND_TO_HOST_TYPE)
   send_op = hlo.SendOp([operand], token, channel_handle,
@@ -606,10 +606,10 @@ def send_to_host(
       # we need to create an equivalent sharding with no dimensions. If there
       # are multiple shardings, just grab the first one since all these
       # shardings should be the same.
-      assert isinstance(sharding, SdyArrayShardingList)
+      assert isinstance(sharding, SdyArrayList)
       assert len(sharding.shardings) >= 1
-      sharding = SdyArrayShardingList([
-          SdyArraySharding(
+      sharding = SdyArrayList([
+          SdyArray(
               mesh_shape=(), dimension_shardings=[],
               logical_device_ids=sharding.shardings[0].logical_device_ids)])
     mlir.set_sharding(send_op, sharding)
@@ -622,7 +622,7 @@ def receive_from_host(
     out_aval: core.ShapedArray,
     name: str,
     *,
-    sharding: SdyArrayShardingList | xc.OpSharding | None = None,
+    sharding: SdyArrayList | xc.OpSharding | None = None,
 ) -> tuple[ir.Value, ir.Value]:
   channel_handle = hlo.ChannelHandle.get(channel, mlir.RECV_FROM_HOST_TYPE)
   recv_op = hlo.RecvOp([mlir.aval_to_ir_type(out_aval),
@@ -634,7 +634,7 @@ def receive_from_host(
           _xla_host_transfer_rendezvous=ir.StringAttr.get(str(name))))
   if sharding is not None:
     if config.use_shardy_partitioner.value:
-      assert isinstance(sharding, SdyArrayShardingList)
+      assert isinstance(sharding, SdyArrayList)
       assert len(sharding.shardings) >= 1
        # `RecvOp`'s last argument is a `TokenType`. Since Shardy requires the
       # number of shardings to match the number of results, but JAX only sees
@@ -642,9 +642,9 @@ def receive_from_host(
       # Note that even if a function returns N results, we will end up with N
       # `RecvOp`s, so we only need to get the first sharding. All shardings are
       # the same anyways, operating on the same single device ID.
-      sharding = SdyArrayShardingList([
+      sharding = SdyArrayList([
           sharding.shardings[0],
-          SdyArraySharding(
+          SdyArray(
               mesh_shape=(), dimension_shardings=[],
               logical_device_ids=sharding.shardings[0].logical_device_ids)])
     mlir.set_sharding(recv_op, sharding)
@@ -683,7 +683,7 @@ def _emit_tpu_python_callback(
     result_avals: Sequence[core.ShapedArray],
     result_shapes: Sequence[xc.Shape],
     *,
-    sharding: SdyArrayShardingList | xc.OpSharding | None = None,
+    sharding: SdyArrayList | xc.OpSharding | None = None,
 ) -> tuple[Sequence[ir.Value], Any]:
   token = token or hlo.create_token()
   _wrapped_callback = callback
@@ -738,7 +738,7 @@ def emit_python_callback(
     *,
     has_side_effect: bool,
     partitioned: bool = False,
-    sharding: SdyArrayShardingList | xc.OpSharding | None = None,
+    sharding: SdyArrayList | xc.OpSharding | None = None,
 ) -> tuple[Sequence[mlir.IrValues], Any, Any]:
   """Emits MLIR that calls back to a provided Python function.
 
@@ -836,12 +836,12 @@ def emit_python_callback(
         config.use_shardy_partitioner.value
         and sharding is not None
         and len(ctx.avals_out) > 0
-        and isinstance(sharding, sharding_impls.SdyArrayShardingList)
+        and isinstance(sharding, sharding_impls.SdyArrayList)
     ):
       # Add a sharding annotation for the token if we have at least one
       # output. Otherwise, the single shardy annotation required of all ops
       # (even those without any results) can annotate the token.
-      sharding = sharding_impls.SdyArrayShardingList(
+      sharding = sharding_impls.SdyArrayList(
           [*sharding.shardings, sharding.shardings[-1]]
       )
     ctx = dataclasses.replace(
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 29cbb0151..63abcbef3 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -165,11 +165,11 @@ def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params)
       # program has per-device semantics, so we run the callback on each device.
       if config.use_shardy_partitioner.value:
         assert len(ctx.avals_out) == 1
-        sharding = sharding_impls.SdyArrayShardingList([
-            sharding_impls.SdyArraySharding(
+        sharding = sharding_impls.SdyArrayList([
+            sharding_impls.SdyArray(
                 mesh_shape=(),
                 dimension_shardings=[
-                    sharding_impls.SdyDimSharding(axes=[], is_open=False)
+                    sharding_impls.SdyDim(axes=[], is_open=False)
                 ] * ctx.avals_out[0].ndim,
                 logical_device_ids=())])
       else:
@@ -182,8 +182,8 @@ def debug_callback_lowering(ctx, *args, effect, partitioned, callback, **params)
     # program has bulk array semantics, so we run the callback with a MAXIMAL
     # sharding and hence execute it only once on the full logical value).
     if config.use_shardy_partitioner.value:
-      sharding = sharding_impls.SdyArrayShardingList([
-          sharding_impls.SdyArraySharding(
+      sharding = sharding_impls.SdyArrayList([
+          sharding_impls.SdyArray(
               mesh_shape=(), dimension_shardings=[], logical_device_ids=(0,))])
     else:
       sharding = xc.OpSharding()
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index f6ef5787c..94418f0b9 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -57,7 +57,7 @@ from jax._src.mesh import AxisType
 from jax._src.partition_spec import PartitionSpec
 from jax._src.sharding import Sharding as JSharding
 from jax._src.sharding_impls import ( AUTO, NamedSharding,
-                                     SdyArraySharding, SdyArrayShardingList,
+                                     SdyArray, SdyArrayList,
                                      modify_sdy_sharding_wrt_axis_types)
 from jax._src.state.types import AbstractRef
 from jax._src.util import foreach
@@ -1034,7 +1034,7 @@ def add_manual_axes(axis_ctx: sharding_impls.SPMDAxisContext, sharding, ndim):
 def _to_physical_op_sharding(
     ctx: ModuleContext,
     aval: core.AbstractValue, sharding: JSharding | AUTO | None,
-) -> xc.OpSharding | SdyArraySharding | None:
+) -> xc.OpSharding | SdyArray | None:
   if sharding is None:
     return None
   if all_unconstrained(sharding, aval):
@@ -1839,10 +1839,10 @@ def replicate_trailing_dims(ctx, val: ir.Value, aval) -> ir.Value:
   assert isinstance(aval, (core.ShapedArray, core.DShapedArray))
   if config.use_shardy_partitioner.value:
     physical_ndim = core.physical_aval(aval).ndim
-    s = SdyArraySharding(
+    s = SdyArray(
         mesh_shape=None,
         dimension_shardings=[
-            sharding_impls.SdyDimSharding(axes=[], is_open=i < aval.ndim)
+            sharding_impls.SdyDim(axes=[], is_open=i < aval.ndim)
             for i in range(physical_ndim)
         ])
     return wrap_with_sharding_op(ctx, val, aval, s)
@@ -2665,7 +2665,7 @@ def _wrap_with_spmd_op(name: str,
                        ctx: LoweringRuleContext,
                        x: ir.Value,
                        aval_out: core.AbstractValue,
-                       sharding: xc.OpSharding | SdyArraySharding,
+                       sharding: xc.OpSharding | SdyArray,
                        unspecified_dims: set[int] | None = None,
                        has_side_effect: bool = False,
                        allow_shardy_lowering: bool = False):
@@ -2730,7 +2730,7 @@ def lower_with_sharding_in_types(ctx, op, aval, sharding_proto=None):
     return wrap_with_sharding_op(ctx, op, aval, proto, unspecified_dims)
 
 
-def set_sharding(op, sharding: xc.OpSharding | SdyArraySharding | SdyArrayShardingList):
+def set_sharding(op, sharding: xc.OpSharding | SdyArray | SdyArrayList):
   if config.use_shardy_partitioner.value:
     op.attributes["sdy.sharding"] = get_sharding_attr(sharding)
   else:
@@ -2738,7 +2738,7 @@ def set_sharding(op, sharding: xc.OpSharding | SdyArraySharding | SdyArrayShardi
 
 
 def get_sharding_attr(
-    sharding: xc.OpSharding | SdyArraySharding | SdyArrayShardingList
+    sharding: xc.OpSharding | SdyArray | SdyArrayList
 ) -> ir.Attribute:
   if config.use_shardy_partitioner.value:
     return sharding.build()  # type: ignore
diff --git a/jax/_src/named_sharding.py b/jax/_src/named_sharding.py
index 3dfcbd29f..45ae1c124 100644
--- a/jax/_src/named_sharding.py
+++ b/jax/_src/named_sharding.py
@@ -41,10 +41,10 @@ class AUTO:
   def __init__(self, mesh: mesh_lib.Mesh):
     self.mesh = mesh
 
-  def _to_sdy_sharding(self, ndim: int) -> SdyArraySharding:
-    dim_shardings = [SdyDimSharding(axes=[], is_open=True)
+  def _to_sdy_sharding(self, ndim: int) -> SdyArray:
+    dim_shardings = [SdyDim(axes=[], is_open=True)
                      for _ in range(ndim)]
-    return SdyArraySharding(self.mesh.shape_tuple, dim_shardings)
+    return SdyArray(self.mesh.shape_tuple, dim_shardings)
 
 class UnspecifiedValue:
   def __repr__(self):
@@ -232,8 +232,8 @@ class NamedSharding(JSharding.Sharding):
   def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
     return named_sharding_to_xla_hlo_sharding(self, num_dimensions)
 
-  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
-    dim_shardings = [SdyDimSharding(axes=[], is_open=False)
+  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
+    dim_shardings = [SdyDim(axes=[], is_open=False)
                      for _ in range(num_dimensions)]
     for i, dim_spec in enumerate(self.spec):
       if dim_spec is PartitionSpec.UNCONSTRAINED:
@@ -244,8 +244,8 @@ class NamedSharding(JSharding.Sharding):
       else:
         dim_spec = dim_spec if isinstance(dim_spec, tuple) else (dim_spec,)
         dim_shardings[i].axes = dim_spec
-    return SdyArraySharding(self.mesh.shape_tuple, dim_shardings,
-                            self._logical_device_ids)
+    return SdyArray(self.mesh.shape_tuple, dim_shardings,
+                    self._logical_device_ids)
 
 NamedSharding.__module__ = 'jax.sharding'
 
@@ -264,7 +264,7 @@ def get_array_mapping(
   return d
 
 @dataclasses.dataclass
-class SdyDimSharding:
+class SdyDim:
   axes: Sequence[str]
   is_open: bool
   priority: int | None = None
@@ -275,7 +275,7 @@ class SdyDimSharding:
         is_closed=not self.is_open, priority=self.priority)
 
   def __repr__(self):
-    return f'SdyDimSharding({self._custom_repr()})'
+    return f'SdyDim({self._custom_repr()})'
 
   def _custom_repr(self):
     axes_repr = ', '.join(f"'{a}'" for a in self.axes)
@@ -287,9 +287,9 @@ class SdyDimSharding:
 
 
 @dataclasses.dataclass
-class SdyArraySharding:
+class SdyArray:
   mesh_shape: tuple[tuple[str, int], ...] | None
-  dimension_shardings: Sequence[SdyDimSharding]
+  dimension_shardings: Sequence[SdyDim]
   logical_device_ids: tuple[int, ...] | None = None
   replicated_axes: tuple[str, ...] = ()
 
@@ -314,7 +314,7 @@ class SdyArraySharding:
                       if self.logical_device_ids is not None else '')
     rar = (f', replicated_axes={self.replicated_axes}'
            if self.replicated_axes else '')
-    return f"SdyArraySharding([{dim_sharding_repr}]{device_id_repr}{rar})"
+    return f"SdyArray([{dim_sharding_repr}]{device_id_repr}{rar})"
 
 
 @cache(max_size=4096, trace_context_in_key=False)
diff --git a/jax/_src/shard_map.py b/jax/_src/shard_map.py
index b772a3de2..3eb46da89 100644
--- a/jax/_src/shard_map.py
+++ b/jax/_src/shard_map.py
@@ -773,7 +773,7 @@ def _valid_repeats(mesh: Mesh, vma: Set[AxisName], names: AxisNames) -> bool:
 
 def _shardy_shard_map_sharding(
     ctx: mlir.LoweringRuleContext, mesh, manual_axes, names, aval_in
-) -> sharding_impls.SdyArraySharding:
+) -> sharding_impls.SdyArray:
   axes = {name: i for i, ns in names.items() for name in ns}
   ns = _make_scoped_manual_sharding(ctx, mesh, axes)
   if dtypes.issubdtype(aval_in.dtype, dtypes.extended):
@@ -808,10 +808,10 @@ def _shard_map_lowering_shardy(
           dim_var_values=ctx.dim_var_values)
     return out_nodes
 
-  in_shardings = sharding_impls.SdyArrayShardingList(map(
+  in_shardings = sharding_impls.SdyArrayList(map(
       partial(_shardy_shard_map_sharding, ctx, mesh, manual_axes),
       in_names, ctx.avals_in)).build()
-  out_shardings = sharding_impls.SdyArrayShardingList(map(
+  out_shardings = sharding_impls.SdyArrayList(map(
       partial(_shardy_shard_map_sharding, ctx, mesh, manual_axes),
       out_names, ctx.avals_out)).build()
   output_types = map(mlir.aval_to_ir_type, ctx.avals_out)
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index 2394e9e18..f7f0ebd2c 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -35,7 +35,7 @@ from jax._src import mesh_utils
 from jax._src.lib import xla_client as xc
 from jax._src.lib.mlir.dialects import sdy
 from jax._src.named_sharding import (  # noqa: F401
-    SdyArraySharding, SdyDimSharding, UnspecifiedValue, AUTO,
+    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)
@@ -87,8 +87,8 @@ def device_replica_id_map(sharding, global_shape: Shape) -> Mapping[Device, int]
 
 
 @dataclasses.dataclass
-class SdyArrayShardingList:
-  shardings: Sequence[SdyArraySharding]
+class SdyArrayList:
+  shardings: Sequence[SdyArray]
 
   def build(self) -> sdy.TensorShardingPerValueAttr:
     return sdy.TensorShardingPerValueAttr.get(
@@ -97,12 +97,12 @@ class SdyArrayShardingList:
 
 # 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: SdyArraySharding, mesh):
+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(SdyDimSharding(axes=[], is_open=True)
+      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)
@@ -111,8 +111,8 @@ def modify_sdy_sharding_wrt_axis_types(sdy_sharding: SdyArraySharding, mesh):
     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 SdyArraySharding(sdy_sharding.mesh_shape, dim_shardings,
-                            sdy_sharding.logical_device_ids, replicated_axes)
+    return SdyArray(sdy_sharding.mesh_shape, dim_shardings,
+                    sdy_sharding.logical_device_ids, replicated_axes)
   return sdy_sharding
 
 
@@ -185,10 +185,10 @@ class SingleDeviceSharding(jsharding.Sharding):
   def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
     return replicated_hlo_sharding
 
-  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
-    sdy_dim_sharding = [SdyDimSharding(axes=[], is_open=False)
+  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
+    sdy_dim_sharding = [SdyDim(axes=[], is_open=False)
                         for _ in range(num_dimensions)]
-    return SdyArraySharding(None, sdy_dim_sharding)
+    return SdyArray(None, sdy_dim_sharding)
 
   @property
   def is_fully_replicated(self) -> bool:
@@ -330,8 +330,8 @@ class PmapSharding(jsharding.Sharding):
   def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
     raise NotImplementedError("pmap doesn't use OpSharding.")
 
-  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
-    raise NotImplementedError("pmap doesn't use SdyArraySharding.")
+  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
+    raise NotImplementedError("pmap doesn't use SdyArray.")
 
   @functools.cached_property
   def is_fully_replicated(self) -> bool:
@@ -540,9 +540,9 @@ class PositionalSharding(jsharding.Sharding):
   def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
     return _positional_sharding_to_xla_hlo_sharding(self, num_dimensions)
 
-  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
+  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
     raise NotImplementedError(
-        "PositionalSharding can't be converted to an SdyArraySharding.")
+        "PositionalSharding can't be converted to an SdyArray.")
 
   @functools.cached_property
   def is_fully_addressable(self) -> bool:
@@ -657,9 +657,9 @@ class GSPMDSharding(jsharding.Sharding):
   def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
     return self._hlo_sharding
 
-  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
+  def _to_sdy_sharding(self, num_dimensions: int) -> SdyArray:
     raise NotImplementedError(
-        "GSPMDSharding can't be converted to SdyArraySharding.")
+        "GSPMDSharding can't be converted to SdyArray.")
 
   @functools.cached_property
   def is_fully_replicated(self) -> bool:


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index b951f7f6b..0bab37c07 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -35,8 +35,8 @@ from jax._src.mesh import AxisType, AbstractMesh
 from jax._src.sharding import common_devices_indices_map
 from jax._src.sharding_impls import (
     _op_sharding_to_pos_sharding, pmap_sharding_devices_indices_map,
-    NamedSharding, GSPMDSharding, PositionalSharding, SdyDimSharding,
-    SdyArraySharding)
+    NamedSharding, GSPMDSharding, PositionalSharding, SdyDim,
+    SdyArray)
 from jax.experimental.pjit import pjit
 from jax.experimental import multihost_utils
 from jax.sharding import PartitionSpec as P
@@ -1476,12 +1476,12 @@ class ShardyShardingTest(jtu.JaxTestCase):
     sdy_sharding = s._to_sdy_sharding(3)
     self.assertEqual(
         sdy_sharding,
-        SdyArraySharding(
+        SdyArray(
             mesh.shape_tuple,
-            [SdyDimSharding(
+            [SdyDim(
              ('sequence', 'data'), False),
-             SdyDimSharding(('model',), False),
-             SdyDimSharding([], False)]))
+             SdyDim(('model',), False),
+             SdyDim([], False)]))
     with ir.Context() as ctx:
       dialects.sdy.register_dialect(ctx)
       self.assertEqual(
@@ -1496,11 +1496,11 @@ class ShardyShardingTest(jtu.JaxTestCase):
     sdy_sharding = s._to_sdy_sharding(3)
     self.assertEqual(
         sdy_sharding,
-        SdyArraySharding(
+        SdyArray(
             mesh.shape_tuple,
-            [SdyDimSharding([], False),
-             SdyDimSharding([], True),
-             SdyDimSharding(('x',), False)]))
+            [SdyDim([], False),
+             SdyDim([], True),
+             SdyDim(('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 22a4d4f70..9b658dd8c 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -8593,26 +8593,26 @@ class ShardyTest(jtu.JaxTestCase):
     self.assertIn('sdy.mesh @mesh = <["x"=8]>', lowered_str)
 
   def test_array_sharding_repr_with_priority(self):
-    sharding = sharding_impls.SdyArraySharding(
+    sharding = sharding_impls.SdyArray(
         mesh_shape=(('data', 4), ('model', 8), ('expert', 2)),
         dimension_shardings=[
-            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])")
+            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])")
 
   def test_array_sharding_repr_with_logical_ids(self):
     abstract_mesh = jax.sharding.AbstractMesh((4, 8, 2), ('x', 'y', 'z'))
     ns = NamedSharding(abstract_mesh, P(('x', 'y'), 'z', P.UNCONSTRAINED, None),
                        _logical_device_ids=[4, 5, 6, 7, 0, 1, 2, 3])
     self.assertEqual(repr(ns._to_sdy_sharding(4)),
-                     "SdyArraySharding([{'x', 'y'}, {'z'}, {?}, {}], "
+                     "SdyArray([{'x', 'y'}, {'z'}, {?}, {}], "
                      "device_ids=[4, 5, 6, 7, 0, 1, 2, 3])")
 
   def test_dimension_sharding_repr(self):
-    dim_sharding = sharding_impls.SdyDimSharding(
+    dim_sharding = sharding_impls.SdyDim(
         axes=['data', 'model'], is_open=True, priority=2)
     self.assertEqual(repr(dim_sharding),
-                     "SdyDimSharding({'data', 'model', ?}p2)")
+                     "SdyDim({'data', 'model', ?}p2)")
 
   def test_tensor_dialect(self):
     # While this doesn't emit any `mlir::TensorDialect` ops, some pass in the