Boolean array indexing is not compatible with static memory allocation

[shoyer]

JAX is built on top of XLA which (currently) requires static memory allocation for operations. This means it is not possible to express operations like x[y] where y is a boolean, because the resulting array has value dependent size: https://data-apis.github.io/array-api/latest/API_specification/indexing.html#boolean-array-indexing

It's definitely possible that the static memory allocation requirement could be relaxed in the future, but dynamic memory allocation is always going to be harder to implement in a performant way. For example, I would guess Numba also struggles with this sort of operation. I don't think we should require it for array libraries implementing our standard, since it isn't needed for the majority of array operations.

[TomAugspurger]

Dask’s lazy evaluation struggles with this a bit too. The output shape of a Boolean mask will have NaN for the shape in that dimension, and .size will be NaN.

…

On Nov 11, 2020, at 17:11, Stephan Hoyer ***@***.***> wrote:  JAX is built on top of XLA which (currently) requires static memory allocation for operations. This means it is not possible to express operations like x[y] where y is a boolean, because the resulting array has a dynamic size: https://data-apis.github.io/array-api/latest/API_specification/indexing.html#boolean-array-indexing It's definitely possible that the static memory allocation requirement could be relaxed in the future, but dynamic memory allocation is always going to be harder to implement in a performant way. For example, I would guess Numba also struggles with this sort of operation. I don't think we should require it for array libraries implementing our standard, since it isn't needed for the majority of array operations. — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub, or unsubscribe.

[rgommers]

For example, I would guess Numba also struggles with this sort of operation

Numba does support boolean indexing, but only for 1-D arrays - and you cannot mix it with other advanced indexing.

[rgommers]

Related JAX issues with discussion:

• boolean indexing support jax-ml/jax#166
• Alternative to array-based boolean indexing for jax.jit jax-ml/jax#2765

It does seem that there's very limited support for boolean indexing (only static boolean arrays, which are not all that useful in practice), and none for jitted functions.

This led me to wonder about other functions where output shape depends on values in the input, such as unique. That is implemented, but also not jittable: jax-ml/jax#2760

TensorFlow is in a similar boat, but provides a boolean_mask function instead. And there's an open PR for JAX to add index_mask which looks similar.

EDIT: add the most informative comment from the JAX PR:
We sometimes need to perform boolean indexing inside of jax.jit, and found that it's useful to have a function which does that while padding the output to the same shape as the original input.

[rgommers]

nonzero and where are also problematic here. Both are implemented in jax.numpy but are not jittable.

[rgommers]

I think that's the complete list of functionality where shapes depend on values - so maybe that's a good set to separate out somehow. They are fairly heavily used functions, so would be good to dig up some examples - say in SciPy or scikit-learn - and figure out if those are easy to rewrite.

[rgommers]

Of course this can happen with basic indexing as well:

x = arange(5)
idx_max = random.randint(0, 5)
x2 = x[:idx_max]

So there needs to be something said about the principle of shapes depending on values.

[rgommers]

The summary of a discussion we had a few days ago:

• boolean indexing is fairly fundamental to how people use NumPy; leaving it out completely may cause people to simply ignore the standard.
• given that this type of data-dependent shape can be produced by several other functions, we need a way to clearly mark that as a property of those functions and of boolean indexing. I.e. with color-coding or a custom admonition in the API standard document.

[jakirkham]

Dask’s lazy evaluation struggles with this a bit too. The output shape of a Boolean mask will have NaN for the shape in that dimension, and .size will be NaN.
…

This is true. Though there are some subtleties (which I'm sure Tom also knows), but others here may not. For example, if we know the values in the mask beforehand (it's not a Dask Array) and raveling is not involved (it's applied along a single dimension), we can produce a Dask Array of known shape. However if one of those things is not true, then the Dask Array's shape is undetermined.

In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: X = da.ones((10, 11, 12),chunks=(2, 3, 4))

In [4]: Y = np.arange(10) > 5

In [5]: X[Y].chunks
Out[5]: ((2, 2), (3, 3, 3, 2), (4, 4, 4))

In [6]: Y2 = np.arange(10 * 11 * 12).reshape(10, 11, 12) > (10 * 11 * 12 / 2)

In [7]: X[Y2].chunks
Out[7]: ((nan, nan, nan, nan, nan, nan, nan, nan, nan, nan),)

In [8]: Y3 = da.asarray(Y, chunks=2)

In [9]: X[Y3].chunks
Out[9]: ((nan, nan, nan, nan, nan), (3, 3, 3, 2), (4, 4, 4))

[leofang]

Sorry, looks like I overlooked this discussion. I don't think CuPy has full support over boolean masks either: cupy/cupy#4693 (comment).

[rgommers]

From the CuPy issue: Combining a multidimensional boolean array with other indexing causes an IndexError.

https://data-apis.github.io/array-api/latest/API_specification/indexing.html#boolean-array-indexing doesn't say it's supported to combine boolean indexing with another indexing mode, however it's not stated clearly. The whole indexing section needs to be easier to understand.

For boolean indexing in particular, the way I read what's there now is: for a boolean array index ix, only x[ix] is supported (also for multi-dimensional ix), but this is not: x[ix, 0]. Which is odd of course, because writing it as x[:, 0][ix] is a more clumsy way to do the same thing and is supported.

I believe the intent was to allow combining a single boolean index with regular basic indexing.

[asmeurer]

Indeed. I had interpreted it to mean that boolean indices were only allowed as the sole index, and were not required to be supported as part of a larger multidimensional index. I believe this was intentional, because mixing a boolean array index with an slice and integer index in NumPy can trigger the weird "advanced indices separated by slice, ellipsis, or newaxis" rule (https://numpy.org/doc/stable/reference/arrays.indexing.html#combining-advanced-and-basic-indexing). Note that for the purposes of this rule, integer indices are treated like 0-dimensional integer array indices:

>>> b = np.array([[True, False], [True, True]])
>>> a = np.arange(1*2*2*4*5).reshape((1, 2, 2, 4, 5))
>>> a[:, b, :].shape
(1, 3, 4, 5)
>>> a[:, b, :, 0].shape
(3, 1, 4)

(note how the 3 dimension is moved to the front when the boolean and integer indices are separated by a slice).

My understanding is that we want to avoid making this rule part of the spec (Travis had expressed to me that he wished it to not be there, and that he would prefer if NumPy didn't do it in the first place). So if we do allow boolean arrays as part of larger multidimensional indices, it needs to be with caveats that disallow this case.

Another oddity I just noticed the boolean indexing section: it says "this is equivalent to A[nonzero(B)]", but that only works if we support integer array indexing, which we don't. So this should be rephrased in a way that doesn't assume that integer array indexing has been specified already.

[asmeurer]

(and to be clear, either way, the spec should be clarified, because I agree the way it's written now is a bit ambiguous)