Repository: google/jax

Issue:
Support for non-scalar np.repeat arguments.
Hello, 

I noticed that jax.numpy.repeat does not support non-scalar repeats.

```
np.repeat(np.array([[1, 2, 3]]), np.array([2, 3, 1]))

NotImplementedError                       Traceback (most recent call last)
<ipython-input-120-c0c19bd58e3a> in <module>()
----> 1 np.repeat(np.array([[1, 2, 3]]), np.array([2, 3, 1]))

google3/third_party/py/jax/numpy/lax_numpy.py in repeat(a, repeats, axis)
   1539   if not isscalar(repeats):
   1540     raise NotImplementedError(
-> 1541         "np.repeat implementation only supports scalar repeats")
   1542   if axis is None or isscalar(a):
   1543     a = ravel(a)
```

The expected results with numpy:
```
array([1, 1, 2, 2, 2, 3])
```

Note that the documentation indicates that this is supported:
https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.repeat.html

Is it possible to add support for this in jax, and/or is there a workaround?


Discussion:
Certainly doable! Anyone who wants to get started on this (maybe you!) should start by looking at [the current `np.repeat` implementation in lax_numpy.py.](https://github.com/google/jax/blob/d7894198a1ad0e54de42450c27ad5e715cb59aa1/jax/numpy/lax_numpy.py#L1596-L1622).
If no one has picked this up I will try to have a go at this. 

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1682,11 +1682,11 @@ def ix_(*args):
   return tuple(output)
 
 
-@_wraps(onp.repeat)
-def repeat(a, repeats, axis=None):
+
+def _repeat_scalar(a, repeats, axis=None):
   if not isscalar(repeats):
     raise NotImplementedError(
-        "np.repeat implementation only supports scalar repeats")
+        "_repeat_scalar implementation only supports scalar repeats")
   if axis is None or isscalar(a):
     a = ravel(a)
     axis = 0
@@ -1710,6 +1710,55 @@ def repeat(a, repeats, axis=None):
       lax.broadcast_in_dim(a, broadcast_shape, broadcast_dims),
       a_shape)
 
+@_wraps(onp.repeat)
+def repeat(a, repeats, axis=None):
+  '''
+  :param repeats: int or array of ints
+  '''
+  # use `_repeat_scalar` when possible
+  if isscalar(repeats):
+    return _repeat_scalar(a, repeats, axis)
+  repeats_raveled = ravel(array(repeats)) # make sure it's jax's array type
+  if size(repeats_raveled) == 1:
+    return _repeat_scalar(a, list(repeats_raveled)[0], axis)
+
+  if axis is None or isscalar(a):
+    a = ravel(a)
+    axis = 0
+
+  # repeats must match the dimension along the requested axis
+  a_shape = list(a.shape)
+  n = a_shape[axis]
+  if size(repeats_raveled) != n:
+    raise ValueError("repeats shape {} does not match the dimension on axis {}".format(
+      repeats_raveled.shape, n
+    ))
+
+  # calculating the new shape
+  total = sum(repeats_raveled)
+
+  new_shape = a_shape[:]
+  new_shape[axis] = total
+
+  a_flattened = ravel(a)
+
+  '''
+  main algorithm:
+  first break down raveled input array into list of chunks; each chunk is the unit of repeat
+  then tile the repeats to have same length as the list of chunks
+  finally repeat each unit x number of times according to the tiled repeat list
+  '''
+  chunks = product(a_shape[:axis+1]).item()
+  a_splitted = split(a_flattened, chunks)
+  repeats_tiled = tile(repeats_raveled, chunks // len(repeats_raveled))
+
+  ret = array([], dtype=a.dtype)
+  for i, repeat in enumerate(repeats_tiled):
+    if not isinstance(repeat, int):
+      repeat = repeat.item()
+    ret = concatenate((ret, tile(a_splitted[i], repeat)))
+
+  return reshape(ret, new_shape)
 
 @_wraps(onp.tri)
 def tri(N, M=None, k=0, dtype=None):


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -729,6 +729,36 @@ def testRepeat(self, axis, shape, dtype, repeats, rng):
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  def testIssue1233(self):
+    '''
+    Following numpy test suite from `test_repeat` at https://github.com/numpy/numpy/blob/master/numpy/core/tests/test_multiarray.py
+    '''
+    tol = 1e-5
+    
+    def test_single(m, args_maker, repeats, axis):
+      lax_ans = lnp.repeat(m, repeats, axis)
+      numpy_ans = onp.repeat(m, repeats, axis)
+
+      self.assertAllClose(lax_ans, numpy_ans, check_dtypes=True, rtol=tol, atol=tol)
+
+      lnp_fun = lambda arg: lnp.repeat(arg, repeats = repeats, axis=axis)
+      self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
+    m = lnp.array([1,2,3,4,5,6])
+    args_maker = lambda: [m]
+
+    for repeats in [2, [1,3,2,1,1,2], [2], lnp.array([1,3,2,1,1,2]), lnp.array([2])]:
+      test_single(m, args_maker, repeats, None)
+
+    m_rect = m.reshape((2,3))
+    args_maker = lambda: [m_rect]
+
+    for repeats in [2, [2,1], [2], lnp.array([2,1]), lnp.array([2])]:
+      test_single(m_rect, args_maker, repeats, axis=0)
+
+    for repeats in [2, [1,3,2], [2], lnp.array([1,3,2]), lnp.array([2])]:
+      test_single(m_rect, args_maker, repeats, axis=1)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "op={}_shape=[{}]_axis={}_out_dtype={}".format(
           op, jtu.format_shape_dtype_string(shape, dtype), axis, out_dtype),