Repository navigation
Indexing numpy array with DeviceArray: index interpreted as tuple #620
Description
When you try to index a numpy ndarray with a DeviceArray, the numpy array tries to interpret the jax array as a tuple.
import numpy as onp
import jax.numpy as np
x = onp.zeros((5,7))
np_idx = onp.array([1,2,3])
jax_idx = np.array([1,2,3])
x[np_idx]
x[jax_idx] # <- raises IndexErrorWorkaround: put jax_idx in a singleton tuple x[(jax_idx,)]
This bug resulted in a confusing situation where my function worked when decorated by jax.jit but had a shape mismatch when called on a numpy array.
Reacted by George Dahl
Activity
mattjj commentedon Apr 17, 2019
This is related to #553, though in that case the user was only calling
jax.numpy(but the surprising behavior cropped up becausejax.numpy.arangeactually just callsonp.arange, producing a plain ndarray).To fix this, we'll either need to get
DeviceArraysto look like regular ndarrays well enough so that Numpy's indexing treats them like regular ndarrays, or we'll need to set up our own overrides, perhaps using ideas like #611.mattjj commentedon Apr 18, 2019
This is a great bug.
I asked our friendly neighborhood NumPy expert @shoyer about it, and after we poked around the NumPy code a bit, he noticed something awesome: this example would have worked if your index array had more than 32 elements in it. Indeed:
That breaks if you change the 11s to 10s, though.
@shoyer also pointed out that this example will work if you write
x[jax_idx, :], i.e. you include the colons.Not sure yet if we have a way to mitigate this issue, but I couldn't wait to report this NumPy oddity.
shoyer commentedon Apr 18, 2019
When I try this example, I also see a deprecation warning from NumPy:
/Users/shoyer/miniconda3/envs/jax-py36/bin/ipython:7: FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use
arr[tuple(seq)]instead ofarr[seq]. In the future this will be interpreted as an array index,arr[np.array(seq)], which will result either in an error or a different result.So hopefully that's a clue that something may be amiss.
mattjj commentedon Apr 18, 2019
Reading that warning again, it makes it sound like this use case will start working for us in the future, when the behavior is removed.
One path forward in the short term might be to add some code to this function in multiarray/mapping.c to check for an
__array__method (i.e. callPyArray_FromArrayAttr) on the indexing object (like our DeviceArray exposes) and call it. That is, maybe we can make this NumPy function respect that style of duck typing of ndarrays, not just proper subclassing. (EDIT: also apparently suggested in this comment.)shoyer commentedon Apr 18, 2019
Yes, once the deprecation cycle finishes this should start working for JAX.
levskaya commentedon Jul 10, 2020
Just remarking that I just helped an intern debug this issue again, so it's still happening circa July 2020.
nsavinov commentedon Oct 26, 2020
I have also just spent some debugging it.
Deprecation warning helped, but it was slightly confusing: rather than telling me the error already happened, it told that the behaviour won't be supported in the future.
nsavinov commentedon Oct 26, 2020
Maybe adding this gotcha to sharp bits would help new users? https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html
hawkinsp commentedon Feb 10, 2022
I sent a PR to fix this in upstream NumPy (numpy/numpy#21029). Let's see if it works!
hawkinsp commentedon Feb 14, 2022
The upstream NumPy PR was merged! So I'm going to declare this fixed, even though you're going to have to wait for NumPy 1.23 to get it...
Fixed 'ready' flag. Also note: sampling will throw an error if using …
Enable test_nvfp4_gradient_clip tests on ROCm (jax-ml#620)