Repository: jax-ml/jax
Pull request #28634 (2025-05-09 02:00:59)

Add `.update` to ShapeDtypeStruct
Add `.update` to ShapeDtypeStruct



Patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 379f6d8d7..5c8a86c03 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2870,6 +2870,24 @@ class ShapeDtypeStruct:
     # https://github.com/jax-ml/jax/issues/8182
     return hash((self.shape, self.dtype, self.sharding, self.layout, self.weak_type))
 
+  def update(self, **kwargs):
+    if 'sharding' in kwargs:
+      s = kwargs['sharding']
+      if self._dll is not None and isinstance(s, Sharding):
+        raise ValueError(
+            f"You are updating ShapeDtypeStruct with a {type(s)} when the"
+            f" original ShapeDtypeStruct had a concrete layout {self.layout}."
+            " This might lead to bugs. If you want to do this, create a new"
+            " ShapeDtypeStruct via the constructor.")
+      sharding = s
+    else:
+      sharding = self.layout
+    return ShapeDtypeStruct(
+        shape=kwargs.pop('shape', self.shape),
+        dtype=kwargs.pop('dtype', self.dtype),
+        sharding=sharding,
+        weak_type=kwargs.pop('weak_type', self.weak_type))
+
 def _sds_aval_mapping(x):
   aval = ShapedArray(
       x.shape, dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True),


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index e4b91aa17..55c38f6c2 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -55,6 +55,7 @@ from jax._src.sharding_impls import (
     SingleDeviceSharding, parse_flatten_op_sharding)
 from jax._src.pjit import (pjit, mesh_cast, auto_axes, explicit_axes,
                            use_auto_axes, use_explicit_axes, reshard)
+from jax._src.layout import Layout, DeviceLocalLayout as DLL
 from jax._src.named_sharding import DuplicateSpecError
 from jax._src import mesh as mesh_lib
 from jax._src.mesh import AxisType
@@ -4973,6 +4974,40 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(out.shape, input_shape)
     jax.random.key_data(out)  # doesn't crash
 
+  def test_sds_update(self):
+    mesh = jtu.create_mesh((2, 1), ('x', 'y'))
+    s1 = jax.ShapeDtypeStruct((2, 2), jnp.int32)
+    s1_u = s1.update(shape=(4, 2), dtype=np.float32)
+    self.assertEqual(s1_u.shape, (4, 2))
+    self.assertEqual(s1_u.dtype, np.float32)
+    self.assertFalse(s1_u.weak_type)
+
+    s2 = jax.ShapeDtypeStruct((2, 2), jnp.int32)
+    s2_u = s2.update(shape=(4, 2), weak_type=True)
+    self.assertEqual(s2_u.shape, (4, 2))
+    self.assertEqual(s2_u.dtype, np.int32)
+    self.assertTrue(s2_u.weak_type)
+
+    s3 = jax.ShapeDtypeStruct((2, 2), jnp.int32,
+                              sharding=NamedSharding(mesh, P()))
+    s3_u = s3.update(sharding=NamedSharding(mesh, P('x')))
+    self.assertEqual(s3_u.sharding, NamedSharding(mesh, P('x')))
+
+    s32_u = s3.update(shape=(4, 2))
+    self.assertEqual(s32_u.shape, (4, 2))
+    self.assertEqual(s32_u.sharding, NamedSharding(mesh, P()))
+
+    sh = NamedSharding(mesh, P())
+    s4 = jax.ShapeDtypeStruct((2, 2), jnp.int32,
+                              sharding=Layout(DLL((0, 1)), sh))
+    new_layout = Layout(DLL((1, 0)), NamedSharding(mesh, P('x')))
+    s4_u = s4.update(sharding=new_layout)
+    self.assertEqual(s4_u.sharding, new_layout.sharding)
+    self.assertEqual(s4_u.layout, new_layout)
+
+    with self.assertRaisesRegex(ValueError, "updating ShapeDtypeStruct"):
+      s4.update(sharding=NamedSharding(mesh, P('x')))
+
 
 def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")