Repository: jax-ml/jax
Pull request #28869 (2025-05-21 00:44:00)

Support passing PartitionSpecs to ShapeDtypeStruct when there is a mesh in context.
Support passing PartitionSpecs to ShapeDtypeStruct when there is a mesh in context.



Patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 33802e494..059db1c92 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2826,9 +2826,10 @@ class ShapeDtypeStruct:
     if dtype is None:
       raise ValueError("ShapeDtypeStruct: dtype must be specified.")
     self.dtype = dtype if dtypes.issubdtype(dtype, dtypes.extended) else np.dtype(dtype)
-    if sharding is not None and not isinstance(sharding, (Sharding, Layout)):
+    if sharding is not None and not isinstance(sharding, (Sharding, Layout, P)):
       raise ValueError(
-          "sharding should be an instance of `jax.sharding.Sharding` or"
+          "sharding should be an instance of `jax.sharding.Sharding`, "
+          "`jax.sharding.PartitionSpec` or"
           f" `jax.experimental.layout.Layout`. Got {sharding} of type"
           f" {type(sharding)}.")
     if (isinstance(sharding, Layout) and
@@ -2836,7 +2837,19 @@ class ShapeDtypeStruct:
       raise TypeError(
           "`DeviceLocalLayout.AUTO` cannot be used in place of a device-local"
           f" layout in a `ShapeDtypeStruct`. Got {sharding}")
-    self.sharding = sharding.sharding if isinstance(sharding, Layout) else sharding
+    if isinstance(sharding, Layout):
+      self.sharding = sharding.sharding
+    elif isinstance(sharding, P):
+      # TODO(yashkatariya): Should this be abstract mesh?
+      cur_mesh = get_concrete_mesh()
+      if cur_mesh is None:
+        raise TypeError(
+            "When specifying PartitionSpec to `ShapeDtypeStruct`, the context"
+            " mesh cannot be empty. Please use `jax.sharding.use_mesh` to set"
+            " the mesh context.")
+      self.sharding = NamedSharding(cur_mesh, sharding)
+    else:
+      self.sharding = sharding
     self._dll = sharding.device_local_layout if isinstance(sharding, Layout) else None
     self.weak_type = weak_type
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index 3fe3d6fa7..9963e2603 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4525,13 +4525,6 @@ class APITest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TypeError, "applied to foo"):
       f_vjp(1.0, 1.0)
 
-  def test_shapedtypestruct_sharding_error(self):
-    with self.assertRaisesRegex(
-        ValueError,
-        "sharding should be an instance of `jax.sharding.Sharding`."):
-      jax.ShapeDtypeStruct((8, 2), np.float32,
-                           sharding=jax.sharding.PartitionSpec('x'))
-
   def test_make_jaxpr_weakref(self):
     class Foo(NamedTuple):
       x: int
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 7d87705e2..88ec58ec3 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5006,6 +5006,25 @@ class ArrayPjitTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(ValueError, "updating ShapeDtypeStruct"):
       s4.update(sharding=NamedSharding(mesh, P('x')))
 
+  @jtu.with_explicit_mesh((2, 1), ('x', 'y'), axis_types=(AxisType.Auto,) * 2)
+  def test_sds_pspec_input(self, mesh):
+    inp = jax.ShapeDtypeStruct((2, 2), np.float32, sharding=P('x'))
+    lowered = jax.jit(lambda x: x * 2).lower(inp)
+    self.assertIn('num_partitions = 2', lowered.as_text())
+
+    np_inp = np.arange(4, dtype=np.float32).reshape(2, 2)
+    arr = jax.device_put(np_inp, P('x'))
+    out = lowered.compile()(arr)
+    self.assertArraysEqual(out, np_inp * 2)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
+  def test_sds_pspec_no_mesh_ctx_error(self):
+    with self.assertRaisesRegex(
+        TypeError,
+        'When specifying PartitionSpec to `ShapeDtypeStruct`, the context mesh'
+        ' cannot be empty'):
+      jax.ShapeDtypeStruct((2, 2), np.float32, sharding=P('x'))
+
 
 def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")