np.arange weirdness with jit.

I found an issue related to jit compilation with np.arange when it is indexed dynamically. In particular, the code below produces the produces the following error.

def fun(x):
  r = np.arange(x.shape[0])[x]
  return r

jit(fun)(np.array([0, 1, 2], dtype=np.int32))

[mattjj]

You can probably also call np.take here. So long as we’re not dispatching on a raw ndarray things are okay.

[sschoenholz]

Ah thanks Matt, makes sense! It turns out one can also do

import jax.numpy as np
from jax.api import device_put

def fun(x):
  r = device_put(np.arange(x.shape[0]))[x]
  return r

jit(fun)(np.array([0, 1, 2], dtype=np.int32))

which might be useful, if less efficient, in cases where the more advanced functionality of np.arange is required.