Propagate unreduced for transpose and DUS correctly in shard_map

PiperOrigin-RevId: 836311442

Author: yashk2810

jax/_src/lax/lax.py

@@ -7498,7 +7498,12 @@ def _transpose_sharding_rule(operand, *, permutation):
   return operand.sharding.update(spec=o_spec.update(partitions=new_spec))
 
 def _transpose_unreduced_rule(out_s, operand, *, permutation):
-  return out_s
+  return out_s.update(spec=out_s.spec.update(
+      unreduced=operand.sharding.spec.unreduced))
+
+def _transpose_reduced_rule(out_s, operand, *, permutation):
+  return out_s.update(spec=out_s.spec.update(
+      reduced=operand.sharding.spec.reduced))
 
 def _transpose_batch_rule(batched_args, batch_dims, *, permutation):
   operand, = batched_args
@@ -7524,7 +7529,8 @@ def _transpose_lower(ctx, x, *, permutation):
     _transpose_shape_rule, input_dtype, 'transpose',
     sharding_rule=_transpose_sharding_rule,
     vma_rule=partial(core.standard_vma_rule, 'transpose'),
-    unreduced_rule=_transpose_unreduced_rule)
+    unreduced_rule=_transpose_unreduced_rule,
+    reduced_rule=_transpose_reduced_rule)
 ad.deflinear2(transpose_p,
               lambda t, _, permutation: [transpose(t, np.argsort(permutation))])
 batching.primitive_batchers[transpose_p] = _transpose_batch_rule

jax/_src/lax/slicing.py

@@ -1723,7 +1723,8 @@ def _dynamic_update_slice_unreduced_rule(out_s, operand, update, *start_indices)
         " same axes. Got operand sharding"
         f" {operand.str_short(mesh_axis_types=True)} and update sharding"
         f" {update.str_short(mesh_axis_types=True)}.")
-  return out_s
+  return out_s.update(spec=out_s.spec.update(
+      unreduced=operand.sharding.spec.unreduced))
 
 def _dynamic_update_slice_reduced_rule(out_s, operand, update, *start_indices):
   if operand.sharding.spec.reduced != update.sharding.spec.reduced:

tests/shard_map_test.py

@@ -4583,6 +4583,29 @@ def loss_fn(inputs, params, targets):
 
     jax.jit(jax.grad(loss_fn, argnums=1))(inputs, params, targets)  # doesn't crash
 
+  @jtu.with_explicit_mesh((2,), 'x')
+  def test_transpose_unreduced_shmap(self, mesh):
+    arr1 = jax.device_put(np.arange(8.).reshape(2, 4), P(reduced={'x'}))
+    arr2 = jax.device_put(np.arange(12.).reshape(2, 6), P(None, 'x'))
+
+    @jax.shard_map(out_specs=P(None, 'x'))
+    def f(x, y):
+      x_ = x.T
+      return jnp.dot(x_, y)
+
+    @jax.jit
+    def g(x, y):
+      return f(x, y).sum()
+
+    out = g(arr1, arr2)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P()))
+    self.assertArraysEqual(out, (arr1.T @ arr2).sum())
+
+    out1, out2 = jax.jit(jax.grad(g, argnums=(0, 1)))(arr1, arr2)
+    self.assertEqual(out1.sharding,
+                     NamedSharding(mesh, P(None, None, unreduced={'x'})))
+    self.assertEqual(out2.sharding, arr2.sharding)
+
 
 class FunSpec(NamedTuple):
   name: str