Incorrect indexing behavior when a numpy array is indexed with a jax.numpy array

[Sohl-Dickstein]

import numpy as onp
import jax.numpy as jnp 

oa = onp.zeros((7,3))
ob = onp.array([1,6])
ja = jnp.zeros((7,3))
jb = jnp.array([1,6])
print(oa[ob]) # works
print(ja[jb]) # works
print(ja[ob]) # works
print(oa[jb, :]) # works
print(oa[jb]) # throws an IndexError

[shoyer]

Yes, this is unfortunate. As far as we can tell, this is basically a NumPy bug (but at least should be automatically fixed once some deprecated behavior is removed). Let's keep discussion over in #620, though.