Repository: jax-ml/jax
Pull request #25021 (2024-11-21 00:32:19)

[sharding_in_types] Add `pad_p` support to sharding_in_types to handle transpose to slice correctly.
[sharding_in_types] Add `pad_p` support to sharding_in_types to handle transpose to slice correctly.



Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 4a9925ce3..9a2746090 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4490,6 +4490,25 @@ def _pad_shape_rule(operand, padding_value, *, padding_config):
     raise ValueError(msg)
   return result
 
+def _pad_sharding_rule(operand, padding_value, *, padding_config):
+  # TODO(yashkatariya): Once JAX supports uneven sharding at the top level,
+  # change this logic to `return operand.sharding` directly.
+  out_shape = _pad_shape_rule(operand, padding_value,
+                              padding_config=padding_config)
+  mesh = operand.sharding.mesh
+  new_spec = []
+  for op_sh, out_sh, op_spec in safe_zip(
+      operand.shape, out_shape, operand.sharding.spec):
+    if (op_sh != out_sh and op_spec is not None and
+        out_sh % slicing._get_sub_spec_size(mesh, op_spec) != 0):
+      raise NotImplementedError(
+          f"padding on sharded dims where out dim ({out_sh}) is not divisble by"
+          f" mesh axes ({slicing._get_sub_spec_size(mesh, op_spec)}) with spec"
+          f" ({op_spec}) is not implemented.")
+    new_spec.append(op_spec)
+  return NamedSharding(mesh, P(*new_spec))
+
+
 def _pad_transpose(t, operand, padding_value, *, padding_config):
   if type(t) is ad_util.Zero:
     t_operand = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
@@ -4529,14 +4548,18 @@ def _pad_batch_rule(batched_args, batch_dims, *, padding_config):
                                          (operand_bdim,))
   return select(mask, x, broadcasted_padding), operand_bdim
 
-pad_p = standard_primitive(_pad_shape_rule, _pad_dtype_rule, 'pad')
+pad_p = standard_primitive(_pad_shape_rule, _pad_dtype_rule, 'pad',
+                           sharding_rule=_pad_sharding_rule)
 ad.deflinear2(pad_p, _pad_transpose)
 batching.primitive_batchers[pad_p] = _pad_batch_rule
 
 def _pad_lower(ctx, x, padding_value, *, padding_config):
   aval_out, = ctx.avals_out
   low, high, interior = util.unzip3(padding_config)
-  return [mlir.pad(ctx, aval_out, x, padding_value, low, high, interior)]
+  out = mlir.pad(ctx, aval_out, x, padding_value, low, high, interior)
+  if config.sharding_in_types.value:
+    return [mlir.lower_sharding_under_shit(ctx, out, aval_out)]
+  return [out]
 
 mlir.register_lowering(pad_p, _pad_lower)
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index e52c805ef..372d6c334 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5287,7 +5287,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
 
   def test_slice(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))
-    np_inp = np.arange(16).reshape(4, 4)
+    np_inp = np.arange(16.).reshape(4, 4)
     arr = jax.device_put(np_inp, NamedSharding(mesh, P('x', None)))
 
     @jax.jit
@@ -5300,6 +5300,16 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
     self.assertIn('@Sharding', f.lower(arr).as_text())
 
+    def g(x):
+      out = f(x)
+      return jnp.square(jnp.sum(out))
+
+    out = jax.grad(g)(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+    out = jax.jit(jax.grad(g))(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
     with self.assertRaisesRegex(NotImplementedError, "slicing on sharded dims"):
       f(jax.device_put(np_inp, NamedSharding(mesh, P('x', 'y'))))
 
@@ -5308,7 +5318,7 @@ class ShardingInTypesTest(jtu.JaxTestCase):
 
   def test_squeeze(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))
-    np_inp = np.arange(16).reshape(4, 4, 1)
+    np_inp = np.arange(16.).reshape(4, 4, 1)
     arr = jax.device_put(np_inp, NamedSharding(mesh, P('x', None, None)))
 
     @jax.jit
@@ -5322,6 +5332,58 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     self.assertIn('@Sharding', f.lower(arr).as_text())
     self.assertArraysEqual(out, np.squeeze(np_inp, axis=2))
 
+    def g(x):
+      out = f(x)
+      return jnp.square(jnp.sum(out))
+
+    out = jax.grad(g)(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+    out = jax.jit(jax.grad(g))(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+  def test_pad(self):
+    mesh = jtu.create_mesh((2, 1), ('x', 'y'))
+    np_inp = np.arange(8.)
+    arr = jax.device_put(np_inp, NamedSharding(mesh, P('x')))
+
+    @partial(jax.jit, static_argnums=(1, 2))
+    def f(x, padding_config, spec):
+      y = lax.pad(x, 0., padding_config)
+      self.assertEqual(y.sharding.spec, spec)
+      return y
+
+    out = f(arr, ((2, 2, 0),), P('x'))
+    self.assertArraysEqual(out, np.pad(np_inp, 2))
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+    self.assertIn('@Sharding', f.lower(arr, ((2, 2, 0),), P('x')).as_text())
+
+    out = f(arr, ((0, 0, 0),), P('x'))
+    self.assertArraysEqual(out, np_inp)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
+    f(arr, ((0, 3, 1), ), P('x'))  # doesn't crash
+
+    def g(x):
+      out = f(x, ((2, 2, 0),), P('x'))
+      return jnp.square(jnp.sum(out))
+
+    out = jax.grad(g)(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+    out = jax.jit(jax.grad(g))(arr)
+    self.assertEqual(out.sharding, arr.sharding)
+
+    with self.assertRaisesRegex(NotImplementedError, "padding on sharded dims"):
+      f(arr, ((2, 3, 0), ), None)
+
+    with self.assertRaisesRegex(NotImplementedError, "padding on sharded dims"):
+      f(arr, ((0, 3, 0), ), None)
+
+    with self.assertRaisesRegex(NotImplementedError, "padding on sharded dims"):
+      arr = jax.device_put(np_inp, NamedSharding(mesh, P(('x', 'y'))))
+      f(arr, ((4, 4, 1),), None)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):