Repository: jax-ml/jax
Pull request #28600 (2025-05-08 02:10:33)

Add `out_sharding` to `jnp.repeat`. Drop into auto mode if out_sharding is provided.
Add `out_sharding` to `jnp.repeat`. Drop into auto mode if out_sharding is provided.

In cases where axis is None or the input is sharded on the `axis` we are going to repeat on.

If the input is not sharded on the repeat axis, forward the input sharding to the output.

Fixes https://github.com/jax-ml/jax/issues/28538



Patch:
diff --git a/jax/_src/basearray.pyi b/jax/_src/basearray.pyi
index cf64afdac..a98cc0120 100644
--- a/jax/_src/basearray.pyi
+++ b/jax/_src/basearray.pyi
@@ -189,7 +189,8 @@ class Array(metaclass=abc.ABCMeta):
   @property
   def real(self) -> Array: ...
   def repeat(self, repeats: ArrayLike, axis: int | None = None, *,
-             total_repeat_length: int | None = None) -> Array: ...
+             total_repeat_length: int | None = None,
+             out_sharding: NamedSharding | P | None = None) -> Array: ...
   def reshape(self, *args: Any, order: str = "C",
               out_sharding: NamedSharding | P | None = ...) -> Array: ...
   def round(self, decimals: int = 0, out: None = None) -> Array: ...
diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
index a3dbc0f9f..73f916537 100644
--- a/jax/_src/numpy/array_methods.py
+++ b/jax/_src/numpy/array_methods.py
@@ -293,12 +293,15 @@ def _real_property(self: Array) -> Array:
   return ufuncs.real(self)
 
 def _repeat(self: Array, repeats: ArrayLike, axis: int | None = None, *,
-            total_repeat_length: int | None = None) -> Array:
+            total_repeat_length: int | None = None,
+            out_sharding: NamedSharding | PartitionSpec | None = None) -> Array:
   """Construct an array from repeated elements.
 
   Refer to :func:`jax.numpy.repeat` for the full documentation.
   """
-  return lax_numpy.repeat(self, repeats=repeats, axis=axis, total_repeat_length=total_repeat_length)
+  return lax_numpy.repeat(self, repeats=repeats, axis=axis,
+                          total_repeat_length=total_repeat_length,
+                          out_sharding=out_sharding)
 
 def _reshape(self: Array, *args: Any, order: str = "C", out_sharding=None
              ) -> Array:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index b662f1d6f..2a1ec7227 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -65,7 +65,9 @@ from jax._src.util import (
     NumpyComplexWarning, canonicalize_axis as _canonicalize_axis,
     ceil_of_ratio, safe_zip, set_module, unzip2)
 from jax.sharding import Sharding
-from jax._src.sharding_impls import (NamedSharding, PartitionSpec as P)
+from jax._src.sharding_impls import NamedSharding, PartitionSpec as P
+from jax._src.mesh import get_abstract_mesh
+from jax._src.pjit import auto_axes
 from jax.tree_util import tree_flatten, tree_map
 import numpy as np
 
@@ -6630,7 +6632,8 @@ def indices(dimensions: Sequence[int], dtype: DTypeLike | None = None,
 
 @export
 def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
-           total_repeat_length: int | None = None) -> Array:
+           total_repeat_length: int | None = None,
+           out_sharding: NamedSharding | P | None = None) -> Array:
   """Construct an array from repeated elements.
 
   JAX implementation of :func:`numpy.repeat`.
@@ -6694,6 +6697,31 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
     Array([[1, 1, 2, 2, 2, 2, 2],
            [3, 3, 4, 4, 4, 4, 4]], dtype=int32)
   """
+  if out_sharding is not None:
+    return auto_axes(
+        partial(_repeat, axis=axis, total_repeat_length=total_repeat_length),
+        out_sharding=out_sharding)(a, repeats)
+  ctx_mesh = get_abstract_mesh()
+  if ctx_mesh._are_all_axes_explicit:
+    aval = core.typeof(a)
+    if axis is None or aval.sharding.spec[axis] is not None:
+      raise ValueError(
+          "Please pass sharding to `jnp.repeat` via `out_sharding` parameter.")
+    assert axis is not None and aval.sharding.spec[axis] is None
+    out_sharding = (NamedSharding(ctx_mesh, P())
+                    if aval.sharding.mesh.empty else aval.sharding)
+    return auto_axes(
+        partial(_repeat, axis=axis, total_repeat_length=total_repeat_length),
+        out_sharding=out_sharding)(a, repeats)
+  try:
+    return _repeat(a, repeats, axis=axis,
+                   total_repeat_length=total_repeat_length)
+  except core.ShardingTypeError as e:
+    raise ValueError(
+        "Please pass sharding to `jnp.repeat` via `out_sharding` parameter.")
+
+def _repeat(a: ArrayLike, repeats: ArrayLike, *, axis: int | None = None,
+            total_repeat_length: int | None = None) -> Array:
   if core.is_dim(repeats):
     util.check_arraylike("repeat", a)
   else:
diff --git a/jax/numpy/__init__.pyi b/jax/numpy/__init__.pyi
index a8de717a0..c52ce2628 100644
--- a/jax/numpy/__init__.pyi
+++ b/jax/numpy/__init__.pyi
@@ -808,7 +808,8 @@ def reciprocal(x: ArrayLike, /) -> Array: ...
 register_jax_array_methods: Any
 def remainder(x: ArrayLike, y: ArrayLike, /) -> Array: ...
 def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = ..., *,
-           total_repeat_length: int | None = ...) -> Array: ...
+           total_repeat_length: int | None = ...,
+           out_sharding: NamedSharding | P | None = None) -> Array: ...
 def reshape(
     a: ArrayLike, shape: DimSize | Shape, order: str = ..., *, copy: bool | None = ...,
     out_sharding: NamedSharding | P | None = ...,


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 373e3ceff..0c43df76c 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7569,6 +7569,19 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = f(arr)
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
 
+  @jtu.with_explicit_mesh((2,), ('x',))
+  def test_jnp_repeat(self, mesh):
+    out = jnp.repeat(np.eye(3), np.array((2,2,2,)) - 1, axis=0)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P(None, None)))
+
+    a = jnp.eye(3)
+    out = jnp.repeat(a, np.array((2,2,2,)) - 1, axis=0)
+    self.assertEqual(out.sharding, a.sharding)
+
+    a = jax.device_put(jnp.eye(4), P('x'))
+    out = jnp.repeat(a, np.array((2,2,2,2)) - 1, axis=0, out_sharding=P('x'))
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
+
   @jtu.with_explicit_mesh((2,), ('x',))
   def test_scatter_gather(self, mesh):
     x = np.random.uniform(size=(mesh.size * 2, 3))