Merge pull request jax-ml#27836 from jakevdp:fix-repeat

PiperOrigin-RevId: 745215941

Author: Google-ML-Automation

jax/_src/numpy/lax_numpy.py

@@ -6989,9 +6989,10 @@ def repeat(a: ArrayLike, repeats: ArrayLike, axis: int | None = None, *,
            [3, 3, 4, 4, 4, 4, 4]], dtype=int32)
   """
   if core.is_dim(repeats):
-    arr = util.ensure_arraylike("repeat", a)
+    util.check_arraylike("repeat", a)
   else:
-    arr, repeats = util.ensure_arraylike("repeat", a, repeats)
+    util.check_arraylike("repeat", a, repeats)
+  arr = asarray(a)
 
   if axis is None:
     arr = arr.ravel()

tests/array_extensibility_test.py

@@ -442,7 +442,7 @@ def __getitem__(self, shape) -> jax.ShapeDtypeStruct:
   NumPyAPI.sig(jnp.real, Complex[5]),
   NumPyAPI.sig(jnp.reciprocal, Float[5]),
   NumPyAPI.sig(jnp.remainder, Float[5], Float[5]),
-  NumPyAPI.sig(jnp.repeat, Float[5], Int[5]),
+  NumPyAPI.sig(jnp.repeat, Float[5], repeats=np.array([2, 3, 1, 5, 4])),
   NumPyAPI.sig(jnp.reshape, Float[6], shape=(2, 3)),
   NumPyAPI.sig(jnp.resize, Float[6], new_shape=(2, 3)),
   NumPyAPI.sig(jnp.right_shift, Int[5], Int[5]),