refactor scatter_op_mask
@@ -2400,36 +2400,22 @@ cpdef _scatter_op_single(ndarray a, ndarray indices, v,
cpdef _scatter_op_mask_single(ndarray a, ndarray mask, v, int axis, op):
- cdef ndarray mask_scanned, mask_br, mask_br_scanned
- cdef int n_true
- cdef tuple lshape, rshape, v_shape
+ cdef ndarray mask_scanned
+ cdef tuple masked_shape
+
+ mask, mask_scanned, masked_shape = _prepare_mask_indexing_single(
+ a, mask, axis)
if not isinstance(v, ndarray):
v = array(v, dtype=a.dtype)
v = v.astype(a.dtype)
-
# broadcast v to shape determined by the mask
- if mask.size <= 2 ** 31 - 1:
- mask_type = numpy.int32
- else:
- mask_type = numpy.int64
- mask_scanned = scan(mask.astype(mask_type).ravel()) # starts with 1
- n_true = int(mask_scanned[-1])
- lshape = a.shape[:axis]
- rshape = a.shape[axis + mask.ndim:]
- v_shape = lshape + (n_true,) + rshape
- v = broadcast_to(v, v_shape)
-
- mask_br = mask._reshape(
- axis * (1,) + mask.shape + (a.ndim - axis - mask.ndim) * (1,))
- mask_br = broadcast_to(mask_br, a.shape)
- mask_br_scanned = scan(mask_br.astype(numpy.int32).ravel())
- mask_br_scanned = mask_br_scanned._reshape(mask_br._shape)
+ v = broadcast_to(v, masked_shape)
if op == 'update':
- _scatter_update_mask_kernel(v, mask_br, mask_br_scanned, a)
+ _scatter_update_mask_kernel(v, mask, mask_scanned, a)
elif op == 'add':
- _scatter_add_mask_kernel(v, mask_br, mask_br_scanned, a)
+ _scatter_add_mask_kernel(v, mask, mask_scanned, a)
else:
raise ValueError('provided op is not supported')
