jax.numpy array indexing has different out-of-bounds behavior to numpy

import jax
import jax.numpy as np

x = np.arange(10)
x = jax.device_put(x)
print(x[[13]])

This prints [3], but it should actually throw an out of bounds error like the original NumPy.

P.S.: why does np.arange return a host array? Is this intended behavior or shouldn’t it rather behave like np.array and return a device array?

[mattjj]

You might find jax.experimental.checkify useful for catching OOB checks.

[schmrlng]

Notably, even non-fancy indexing doesn’t raise a bounds error anymore (it wraps instead). Using the example in https://github.com/google/jax/issues/278#issuecomment-457829164:

import jax
import jax.numpy as np

x = np.arange(10)
x = jax.device_put(x)
print(x[13])  # 3