np.ndarray.__getitem__(jnp.ndarray) is incorrect

[trevorcai]

import jax
import jax.numpy as jnp
import numpy as np

arr = np.zeros((2, 8))
idxs = np.zeros((1, 4), dtype=np.int32)
jax_arr = jax.device_put(arr)
jax_idxs = jax.device_put(idxs)
print(arr[idxs].shape)
print(arr[jax_idxs].shape)
print(jax_arr[idxs].shape)
print(jax_arr[jax_idxs].shape)

prints:

(1, 4, 8)
(4, 8)
(1, 4, 8)
(1, 4, 8)

original numpy is mistreating jnp.ndarray for some reason; changing the shape of idxs can lead to various errors.

Continuing from above:
arr[jnp.zeros((2, 4), dtype=jnp.int32)].shape prints (4,)
arr[jnp.zeros((3, 4), dtype=jnp.int32)].shape raises: IndexError: too many indices for array
The other 3/4 combinations are all correct for these cases.

[trevorcai]

Dupe of #620