nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: jax-ml/jax
Pull request #27306 (2025-03-20 20:19:23)

[sharding_in_types] Fix a dynamic_slice bug where in the transpose, `DUS`'s operand was not sharded properly
[sharding_in_types] Fix a dynamic_slice bug where in the transpose, `DUS`'s operand was not sharded properly



Patch:
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
index c26de99c7..b4a8817fb 100644
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -1476,7 +1476,8 @@ def _dynamic_slice_transpose_rule(t, operand, *start_indices, slice_sizes):
   if type(t) is ad_util.Zero:
     return [ad_util.Zero(operand.aval)] + [None] * len(start_indices)
   else:
-    zeros = lax.full(operand_shape, 0, operand_dtype)
+    zeros = lax.full(operand_shape, 0, operand_dtype,
+                     sharding=operand.aval.sharding)
     return ([dynamic_update_slice_p.bind(zeros, t, *start_indices)] +
             [None] * len(start_indices))
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 6cf114949..6a1a73fe4 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7090,6 +7090,30 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     self.assertEqual(out.shape, expected_shape)
     self.assertEqual(out.sharding, NamedSharding(mesh, expected_spec))
 
+  @jtu.with_user_mesh((2,), ('x',))
+  def test_dynamic_slice(self, mesh):
+    np_inp = np.arange(16., dtype=np.float32)
+    s = NamedSharding(mesh, P('x'))
+    arr = jax.device_put(np_inp, s)
+
+    @jax.jit
+    def f(x):
+      y = lax.dynamic_slice_in_dim(x, jnp.array(1, dtype=np.int32), 2)
+      self.assertEqual(y.aval.sharding.spec, P('x'))
+      return y
+
+    out = f(arr)
+    self.assertEqual(out.sharding, s)
+
+    def g(x):
+      return jnp.sum(f(x))
+
+    out = jax.jit(jax.grad(g))(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+    out = jax.grad(g)(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
   def test_auto_axes_computation_follows_data_error(self):
     mesh = jtu.create_mesh((2,), ('x',), axis_types=(AxisType.Explicit,))
     s = NamedSharding(mesh, P('x'))