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.

[mattjj]

My excuse is that I’m rusty because I’ve been on paternity leave for almost 3 weeks 😃

[mattjj]

Hah, okay, I read @fehiepsi’s OP again and I realized he basically figured out the issue for us. I was just too sloppy to read it!

I think we just need to broadcast arr against indices in our implementation of take_along_axis. Will attempt that now…