Add support for boolean mask indices inside jit about jax HOT 5 CLOSED

how does .at work with boolean indices when not compiled?

It converts them to integer indices, then passes those integer indices to scatter or gather. The conversion happens here:

You can see how tricky it is to get the semantics of boolean indices correct in numpy-style indexing expressions...

from jax.

Currently the update op (at[...].set(...)) lowers to a scatter, which needs to know statically how many values are being updated.

I agree that it would be nice to jit-compile updates with dynamic indices, but I'm not sure if there's a clean way to do it at the moment. We could conditionally lower the operation to a select when the operand is a scalar, which would allow it to be jit-compiled as described. This might not be very performant however, if we potentially have to materialize a mask over a large array.

Tagging @jakevdp for any additional opinions.

from jax.

I've thought about this on several occasions, but the added flexibility has never seemed worth the added complexity in the code.

As @justinjfu mentioned, arr.at[...].* is currently a way of spellinglax.gather and lax.scatter. These APIs do not support dynamic shapes, and so boolean indices cannot be used within JIT.

It's important to note that arr.at[mask].set(arr2) is inherently a dynamically-shaped operation, because semantically arr2 must be broadcast-compatible with the entries of arr selected by mask. It's true that there is a special case when arr2 has size 1, because a size-1 axis can broadcast to any array shape: but the shape semantics are still dynamic.

In the simplest case of a 1D array with a 1D mask, you could imagine forking the logic and lowering to select rather than scatter; but immediately someone would ask about arr.at[0, mask].set(arr2) or arr.at[indices, mask].set(arr2) or arr.at[..., mask1, :, 0, mask2].set(arr2), or all the other variants of combined indices that it may be possible to lower to select (depending on the shape of arr2). To get all of that logic correct would be tricky, and would add a lot of complexity both to the implementation, and to the mental overhead of the user trying to figure out when an operation may or may not be JIT-compatible.

By contrast, the current rule is very simple: boolean indices are not supported by JIT, and if you want to use select instead of gather/scatter, you can do so explicitly. In this case, I think simpler is better.

from jax.

Thanks for the answers. Just one more question - how does .at work with boolean indices when not compiled?

from jax.

ooh I see. Well, would still be nice to have at some point...

from jax.