Repository: jax-ml/jax
Pull request #28745 (2025-05-14 17:13:52)

Partial in `repeats` if `total_repeat_length is None` since repeats is expected to be a constant.
Partial in `repeats` if `total_repeat_length is None` since repeats is expected to be a constant.

This fixes an error in explicit sharding mode where we were converting repeats to a Tracer instead of it being a concrete value.



Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 266aad495..0bd287dad 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -6698,9 +6698,8 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
            [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)
+    return _auto_repeat(_repeat, a, repeats, axis, total_repeat_length,
+                        out_sharding)
   ctx_mesh = get_abstract_mesh()
   if ctx_mesh._are_all_axes_explicit:
     aval = core.typeof(a)
@@ -6710,17 +6709,26 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
     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)
+    return _auto_repeat(_repeat, a, repeats, axis, total_repeat_length,
+                        out_sharding)
   try:
-    return _repeat(a, repeats, axis=axis,
+    return _repeat(a, repeats=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,
+def _auto_repeat(fun, a, repeats, axis, total_repeat_length, out_sharding):
+  if total_repeat_length is None:
+    return auto_axes(partial(fun, repeats=repeats, axis=axis,
+                             total_repeat_length=total_repeat_length),
+                     out_sharding=out_sharding)(a)
+  else:
+    return auto_axes(
+        partial(fun, axis=axis, total_repeat_length=total_repeat_length),
+        out_sharding=out_sharding)(a, repeats=repeats)
+
+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)


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index e374b0b15..d1c8ec7f0 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7617,6 +7617,12 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     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)))
 
+    a = jax.device_put(jnp.eye(16).reshape(16, 16), P('x'))
+    @jax.jit
+    def f(x):
+      return jnp.repeat(x, 3, axis=-1)
+    f(a)
+
   @jtu.with_explicit_mesh((2,), ('x',))
   def test_scatter_gather(self, mesh):
     x = np.random.uniform(size=(mesh.size * 2, 3))