Indexing bug with multiple index arrays

From the numpy docs:

import numpy as np
y = np.arange(35).reshape(5,7)
y[np.array([0,2,4]), np.array([0,1,2])]

yields array([ 0, 15, 30])

However this doesn’t work with jax:

import jax.numpy as np
y = np.reshape(np.arange(35), (5, 7))
y[np.array([0,2,4]), np.array([0,1,2])]

yields TypeError: No abstraction handler for type: <class 'tuple'>

Is this a bug, or am I misunderstanding jax’s intended capabilities?

[mattjj]

It’s a bug! We want to handle NumPy indexing exactly (except for boolean indexing, which we’ll add sooner or later but can’t be supported under @jit because it results in dynamic shapes).

Thanks for raising this, and for the clear report.

[mattjj]

Since we can’t repro the issue, and since we have a test checked in for it that passes, I’m going to close this issue as fixed. Please re-open if you run into it again!