Indexing with multidimensional boolean array and a number errors

[outofculture]

Short description

Combining a multidimensional boolean array with other indexing causes an IndexError.

Code to reproduce

data = cp.random.normal(size=(10, 10, 4), scale=8)
mask = data[..., 0] < -1
data[mask, 3] = 0

Expected behavior

Assign a bunch of zeros in the 4th elements of all innermost arrays whose 1st element is less than -1. The code above works when using numpy.

Real behavior

The following error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-6-b7f2e3dfd722> in <module>
----> 1 data[mask, 3] = 0

cupy/core/core.pyx in cupy.core.core.ndarray.__setitem__()

cupy/core/_routines_indexing.pyx in cupy.core._routines_indexing._ndarray_setitem()

cupy/core/_routines_indexing.pyx in cupy.core._routines_indexing._scatter_op()

cupy/core/_routines_indexing.pyx in cupy.core._routines_indexing._scatter_op_mask_single()

cupy/core/_routines_indexing.pyx in cupy.core._routines_indexing._prepare_mask_indexing_single()

IndexError: boolean index did not match

Conditions

OS                           : Linux-5.8.0-7642-generic-x86_64-with-debian-bullseye-sid
CuPy Version                 : 8.4.0
NumPy Version                : 1.20.1
SciPy Version                : None
Cython Build Version         : 0.29.21
CUDA Root                    : /usr/lib/cuda
CUDA Build Version           : 10020
CUDA Driver Version          : 11020
CUDA Runtime Version         : 10020
cuBLAS Version               : 10202
cuFFT Version                : 10102
cuRAND Version               : 10102
cuSOLVER Version             : (10, 3, 0)
cuSPARSE Version             : 10301
NVRTC Version                : (10, 2)
Thrust Version               : 100907
CUB Build Version            : 100800
cuDNN Build Version          : 8005
cuDNN Version                : 8004
NCCL Build Version           : 2708
NCCL Runtime Version         : 2708
cuTENSOR Version             : None
Device 0 Name                : GeForce GTX 1660 Ti
Device 0 Compute Capability  : 75

Additional context

Indexing with just the multidimensional array outputs the correct shape, so it's only when combined with other indexing that the internal understanding of its shape is incorrect. Only multi-dimensional arrays have this problem; size=(10, 4) is a shape that does not produce this error, size=(10, 10, 10, 4) does.

data[mask, :, 3] = 0 is effective in this case, as a workaround. data[mask, ..., 3] = 0 also works more generally.

Probably related to #1512

[kmaehashi]

Thank you for the feedback! Currently CuPy does not have a full support for boolean masks:

https://docs.cupy.dev/en/stable/overview.html

Most of Advanced indexing (except for some indexing patterns with boolean masks)

But I agree that it's better to add support for your boolean + int case as it's likely to be used frequently.
We also have to add more docs on what's supported and what's not.

[outofculture]

Thanks for your work making cupy better! Yeah, I recognize that the general case is pretty challenging. I was hoping the work-around I described would point toward a possible implementation in this case. E.g. splice in an :, to the indexing args for every ndim of the first arg. I'm not sure what other problems this might introduce, though.

[JonathanMaes]

It seems this has been solved in CuPy v10 with #6196.
The minimal example you give in the original post above works without any errors now.

I encountered this issue myself several months ago using an earlier version of CuPy. Since I recently upgraded to CuPy v10, I decided to give the boolean indexing another try for my own use case, which now also works without flaw.

[kmaehashi]

@JonathanMaes Thanks for the report! Yes, this is now implemented.