comment on code
@@ -2256,6 +2256,8 @@ cpdef _prepare_mask_indexing_single(ndarray a, ndarray mask, int axis):
cdef int n_true
cdef tuple lshape, rshape, out_shape
+ # Get number of True in the mask to determine the shape of the array
+ # after masking.
if mask.size <= 2 ** 31 - 1:
mask_type = numpy.int32
else:
@@ -2266,10 +2268,11 @@ cpdef _prepare_mask_indexing_single(ndarray a, ndarray mask, int axis):
rshape = a.shape[axis + mask.ndim:]
masked_shape = lshape + (n_true,) + rshape
- # when mask covers the entire array, broadcasting is not necessary
+ # When mask covers the entire array, broadcasting is not necessary.
if mask.ndim == a.ndim and axis == 0:
return mask, mask_scanned._reshape(mask._shape), masked_shape
+ # The scan of the broadcasted array is used to index on kernel.
mask_br = mask._reshape(
axis * (1,) + mask.shape + (a.ndim - axis - mask.ndim) * (1,))
mask_br = broadcast_to(mask_br, a.shape)
