Indexing with multidimensional boolean array and a number errors #4693
Description
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] = 0Expected 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