Repository: google/jax

Issue:
np.take_along_axis gives wrong gradient
The following repro script shows the issue:
```
import jax.numpy as np
from jax import value_and_grad

idx = np.repeat(np.arange(3), 10).reshape((30, 1))

def f(x):
    y = x * np.arange(3.).reshape((1, 3))
    return np.take_along_axis(y, idx, -1).sum()

def g(x):
    y = x * np.arange(3.).reshape((1, 3))
    y = np.broadcast_to(y, (30, 3))
    return np.take_along_axis(y, idx, -1).sum()

print(value_and_grad(f)(1.))  # get 30, 0
print(value_and_grad(g)(1.))  # get 30, 30
```

I think `(30, 30)` is the correct answer.


Discussion:
Thanks for catching this.

By the way, I recommend using `check_grad` from `jax.test_util` for checking derivatives (because you don't have to worry about scalar output, and you can check fwd/rev and arbitrary order all at once):

```python
import jax.numpy as np
from jax import value_and_grad
from jax.test_util import check_grads

idx = np.arange(3).reshape(3, 1)

def h(x):
  y = x * np.arange(3.).reshape((1, 3))
  return np.take_along_axis(y, idx, -1)

check_grads(h, (1.,), order=1)
```

```
AssertionError:
0.0 !=
9.63398933411
```
It seems `check_grads(h, (1.,), order=1, modes=["fwd"])` succeeds but `check_grads(h, (1.,), order=1, modes=["rev"])` fails. The issue must be with the gather transpose rule (?).

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
@@ -2405,7 +2405,7 @@ def _take_along_axis(arr, indices, axis):
   if axis is None:
     if ndim(indices) != 1:
       msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
-      raise ValueError(msg.format(shape(indices)))
+      raise ValueError(msg.format(indices.shape))
     return take_along_axis(arr.ravel(), indices, 0)
   rank = ndim(arr)
   if rank != ndim(indices):
@@ -2413,11 +2413,19 @@ def _take_along_axis(arr, indices, axis):
     raise ValueError(msg.format(ndim(indices), ndim(arr)))
   axis = _canonicalize_axis(axis, rank)
 
-  arr_shape = list(shape(arr))
-  axis_size = arr_shape[axis]
-  arr_shape[axis] = 1
-  idx_shape = shape(indices)
-  out_shape = lax.broadcast_shapes(idx_shape, tuple(arr_shape))
+  def replace(tup, val):
+    lst = list(tup)
+    lst[axis] = val
+    return tuple(lst)
+
+  bcast_shape = lax.broadcast_shapes(replace(arr.shape, 1), replace(indices.shape, 1))
+  indices = broadcast_to(indices, replace(bcast_shape, indices.shape[axis]))
+  arr     = broadcast_to(arr,     replace(bcast_shape, arr.shape[axis]))
+
+  axis_size = arr.shape[axis]
+  arr_shape = replace(arr.shape, 1)
+  idx_shape = indices.shape
+  out_shape = lax.broadcast_shapes(idx_shape, arr_shape)
 
   index_dims = [i for i, idx in enumerate(idx_shape) if i == axis or idx != 1]
 
@@ -3144,4 +3152,4 @@ def _unstack(x):
   if x.ndim == 0:
     raise ValueError("Argument to _unstack must be non-scalar")
   return [lax.index_in_dim(x, i, keepdims=False) for i in range(x.shape[0])]
-setattr(DeviceArray, "_unstack", _unstack)
\ No newline at end of file
+setattr(DeviceArray, "_unstack", _unstack)


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
@@ -1985,6 +1985,16 @@ def testOpGrad(self, op, rng, shapes, dtype, order, tol):
   def testOpGradSpecialValue(self, op, special_value):
     check_grads(op, (special_value,), 2, ["fwd", "rev"])
 
+  def testTakeAlongAxisIssue1521(self):
+    # https://github.com/google/jax/issues/1521
+    idx = lnp.repeat(lnp.arange(3), 10).reshape((30, 1))
+
+    def f(x):
+      y = x * lnp.arange(3.).reshape((1, 3))
+      return lnp.take_along_axis(y, idx, -1).sum()
+
+    check_grads(f, (1.,), order=1)
+
 
 if __name__ == "__main__":
   absltest.main()