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

PiperOrigin-RevId: 761322712

Author: yashk2810

jax/_src/api.py

@@ -2826,17 +2826,30 @@ def __init__(self, shape, dtype, *, sharding=None, weak_type=False):
     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
         isinstance(sharding.device_local_layout, AutoLayout)):
       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
 

tests/api_test.py

@@ -4525,13 +4525,6 @@ def foo(x):
     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

tests/pjit_test.py

@@ -5006,6 +5006,25 @@ def test_sds_update(self):
     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"\)")