diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py index cda3166a7013..7ac666eeb282 100644 --- a/python/mxnet/ndarray/ndarray.py +++ b/python/mxnet/ndarray/ndarray.py @@ -39,6 +39,7 @@ from ..base import ctypes2buffer from ..runtime import Features from ..context import Context, current_context +from ..util import is_np_array from . import _internal from . import op from ._internal import NDArrayBase @@ -111,7 +112,11 @@ _NDARRAY_BASIC_INDEXING = 0 _NDARRAY_ADVANCED_INDEXING = 1 _NDARRAY_EMPTY_TUPLE_INDEXING = 2 -_NDARRAY_BOOLEAN_INDEXING = 3 + +# Return code for 0-d boolean array handler +_NDARRAY_NO_ZERO_DIM_BOOL_ARRAY = -1 +_NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE = 0 +_NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE = 1 # Caching whether MXNet was built with INT64 support or not _INT64_TENSOR_SIZE_ENABLED = None @@ -521,7 +526,7 @@ def __setitem__(self, key, value): return else: - key = indexing_key_expand_implicit_axes(key, self.shape) + key, _ = indexing_key_expand_implicit_axes(key, self.shape) slc_key = tuple(idx for idx in key if idx is not None) if len(slc_key) < self.ndim: @@ -714,7 +719,7 @@ def __getitem__(self, key): # pylint: disable=too-many-return-statements elif key.step == 0: raise ValueError("slice step cannot be zero") - key = indexing_key_expand_implicit_axes(key, self.shape) + key, _ = indexing_key_expand_implicit_axes(key, self.shape) if len(key) == 0: raise ValueError('indexing key cannot be an empty tuple') @@ -2574,9 +2579,12 @@ def asscalar(self): >>> type(x.asscalar()) """ - if self.shape != (1,): + if self.size != 1: raise ValueError("The current array is not a scalar") - return self.asnumpy()[0] + if self.ndim == 1: + return self.asnumpy()[0] + else: + return self.asnumpy()[()] def astype(self, dtype, copy=True): """Returns a copy of the array after casting to a specified type. @@ -2943,9 +2951,23 @@ def _scatter_set_nd(self, value_nd, indices): lhs=self, rhs=value_nd, indices=indices, shape=self.shape, out=self ) +def check_boolean_array_dimension(array_shape, axis, bool_shape): + """ + Advanced boolean indexing is implemented through the use of `nonzero`. + Size check is necessary to make sure that the boolean array + has exactly as many dimensions as it is supposed to work with before the conversion + """ + for i, val in enumerate(bool_shape): + if array_shape[axis + i] != val: + raise IndexError('boolean index did not match indexed array along axis {};' + ' size is {} but corresponding boolean size is {}' + .format(axis + i, array_shape[axis + i], val)) def indexing_key_expand_implicit_axes(key, shape): - """Make implicit axes explicit by adding ``slice(None)``. + """ + Make implicit axes explicit by adding ``slice(None)`` + and convert boolean array to integer array through `nonzero`. + Examples -------- >>> shape = (3, 4, 5) @@ -2957,6 +2979,11 @@ def indexing_key_expand_implicit_axes(key, shape): (0, slice(None, None, None), slice(None, None, None)) >>> indexing_key_expand_implicit_axes(np.s_[:2, None, 0, ...], shape) (slice(None, 2, None), None, 0, slice(None, None, None)) + >>> bool_array = np.array([[True, False, True, False], + [False, True, False, True], + [True, False, True, False]], dtype=np.bool) + >>> indexing_key_expand_implicit_axes(np.s_[bool_array, None, 0:2], shape) + (array([0, 0, 1, 1, 2, 2], dtype=int64), array([0, 2, 1, 3, 0, 2], dtype=int64), None, slice(None, 2, None)) """ if not isinstance(key, tuple): key = (key,) @@ -2966,6 +2993,17 @@ def indexing_key_expand_implicit_axes(key, shape): ell_idx = None num_none = 0 nonell_key = [] + + # For 0-d boolean indices: A new axis is added, + # but at the same time no axis is "used". So if we have True, + # we add a new axis (a bit like with np.newaxis). If it is + # False, we add a new axis, but this axis has 0 entries. + # prepend is defined to handle this case. + # prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY/-1 means there is no 0-d boolean scalar + # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE/0 means an zero dim must be expanded + # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE/1 means a new axis must be expanded + prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY + axis = 0 for i, idx in enumerate(key): if idx is Ellipsis: if ell_idx is not None: @@ -2974,14 +3012,38 @@ def indexing_key_expand_implicit_axes(key, shape): ) ell_idx = i else: + # convert primitive type boolean value to mx.np.bool type + # otherwise will be treated as 1/0 + if isinstance(idx, bool): + idx = array(idx, dtype=np.bool_) if idx is None: num_none += 1 - if isinstance(idx, NDArrayBase) and idx.ndim == 0 and idx.dtype != np.bool_: + if isinstance(idx, NDArrayBase) and idx.ndim == 0 and idx.dtype == np.bool_: + if not idx: # array(False) has priority + prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE + else: + prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE + elif isinstance(idx, NDArrayBase) and idx.ndim == 0 and idx.dtype != np.bool_: # This handles ndarray of zero dim. e.g array(1) # while advoid converting zero dim boolean array - nonell_key.append(idx.item()) + # float type will be converted to int + nonell_key.append(int(idx.item())) + axis += 1 + elif isinstance(idx, NDArrayBase) and idx.dtype == np.bool_: + # Necessary size check before using `nonzero` + check_boolean_array_dimension(shape, axis, idx.shape) + # If the whole array is false and npx.set_np() is not set_up + # the program will throw infer shape error + if not is_np_array(): + raise ValueError('Cannot perform boolean indexing in legacy mode. Please activate' + ' numpy semantics by calling `npx.set_np()` in the global scope' + ' before calling this function.') + # Add the arrays from the nonzero result to the index + nonell_key.extend(idx.nonzero()) + axis += idx.ndim else: nonell_key.append(idx) + axis += 1 nonell_key = tuple(nonell_key) @@ -2995,7 +3057,7 @@ def indexing_key_expand_implicit_axes(key, shape): (slice(None),) * ell_ndim + nonell_key[ell_idx:]) - return expanded_key + return expanded_key, prepend def _int_to_slice(idx): @@ -3053,32 +3115,18 @@ def _is_advanced_index(idx): def get_indexing_dispatch_code(key): """Returns a dispatch code for calling basic or advanced indexing functions.""" assert isinstance(key, tuple) - num_bools = 0 - basic_indexing = True for idx in key: - if isinstance(idx, (NDArray, np.ndarray, list, tuple)): + if isinstance(idx, (NDArray, np.ndarray, list, tuple, range)): if isinstance(idx, tuple) and len(idx) == 0: return _NDARRAY_EMPTY_TUPLE_INDEXING - if getattr(idx, 'dtype', None) == np.bool_: - num_bools += 1 - basic_indexing = False - elif isinstance(idx, range): - basic_indexing = False + return _NDARRAY_ADVANCED_INDEXING elif not (isinstance(idx, (py_slice, integer_types)) or idx is None): raise ValueError( 'NDArray does not support slicing with key {} of type {}.' ''.format(idx, type(idx)) ) - if basic_indexing and num_bools == 0: - return _NDARRAY_BASIC_INDEXING - elif not basic_indexing and num_bools == 0: - return _NDARRAY_ADVANCED_INDEXING - elif num_bools == 1: - return _NDARRAY_BOOLEAN_INDEXING - else: - raise TypeError('ndarray indexing does not more than one boolean ndarray' - ' in a tuple of complex indices.') + return _NDARRAY_BASIC_INDEXING def _get_index_range(start, stop, length, step=1): diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 651699460348..fceaaf3a282f 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -76,10 +76,11 @@ _NDARRAY_BASIC_INDEXING = 0 _NDARRAY_ADVANCED_INDEXING = 1 _NDARRAY_EMPTY_TUPLE_INDEXING = 2 -_NDARRAY_BOOLEAN_INDEXING = 3 -_NDARRAY_INT_BOOLEAN_INDEXING = 4 -_NDARRAY_SLICE_BOOLEAN_INDEXING = 5 +# Return code for 0-d boolean array handler +_NDARRAY_NO_ZERO_DIM_BOOL_ARRAY = -1 +_NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE = 0 +_NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE = 1 # This function is copied from ndarray.py since pylint # keeps giving false alarm error of undefined-all-variable @@ -405,121 +406,43 @@ def _set_np_advanced_indexing(self, key, value): value_nd = self._prepare_value_nd(value, bcast_shape=vshape, squeeze_axes=new_axes) self._scatter_set_nd(value_nd, idcs) - def _check_boolean_indexing_type(self, key): - """Check boolean indexing type arr[bool, :, :], arr[1, bool, 4], or arr[:, bool, :] - return bool_type, bool_position""" - - dim = len(key) - rest_int = True - rest_full_slice = True - pos = None - for idx in range(dim): - if isinstance(key[idx], _np.ndarray) and key[idx].dtype == _np.bool_: - key[idx] = array(key[idx], dtype='bool', ctx=self.ctx) - if isinstance(key[idx], ndarray) and key[idx].dtype == _np.bool_: - pos = idx - elif isinstance(key[idx], integer_types): - rest_full_slice = False - elif isinstance(key[idx], py_slice) and key[idx] == slice(None, None, None): - rest_int = False - # not arr[:, bool, :] format slicing or not arr[3,bool,4] - else: - raise TypeError('ndarray boolean indexing does not support slicing ' - 'with key {} of type {}'.format(idx, type(idx)) - ) - - if rest_int: - return _NDARRAY_INT_BOOLEAN_INDEXING, pos - elif rest_full_slice: - return _NDARRAY_SLICE_BOOLEAN_INDEXING, pos - raise NotImplementedError("Do not support {} as key for boolean indexing".format(key)) - - @staticmethod - def _calculate_new_idx(key, shape, mask_pos, mask_ndim): # pylint: disable=redefined-outer-name - new_idx = 0 - step = 1 - for idx in range(len(key)-1, mask_pos, -1): - new_idx += key[idx]*step - step *= shape[idx+mask_ndim-1] - return new_idx - - def _get_np_boolean_indexing(self, key): - if not isinstance(key, tuple): - key = (key,) - bool_type, pos = self._check_boolean_indexing_type(key) - - from functools import reduce - mask_shape = key[pos].shape - mask_ndim = len(mask_shape) - ndim = len(self.shape) # pylint: disable=redefined-outer-name, unused-variable - for i in range(mask_ndim): - if key[pos].shape[i] != self.shape[pos + i]: - raise IndexError('boolean index did not match indexed array along axis {};' - ' size is {} but corresponding boolean size is {}' - .format(pos + i, self.shape[pos + i], key[pos].shape[i])) - remaining_idces = pos + mask_ndim - remaining_shapes = self.shape[remaining_idces:] - mask = _reshape_view(key[pos], -1) - - if bool_type == _NDARRAY_SLICE_BOOLEAN_INDEXING: - data = _reshape_view(self, -1, *remaining_shapes) - # if mask is at the begining, then the scale is one - scale = reduce(lambda x, y: x * y, self.shape[:pos], 1) - keys = mask if scale == 1 else _reshape_view(_npi.stack(*[mask for i in range(scale)]), -1) - all_shapes = self.shape[:pos] + remaining_shapes - return _reshape_view(_npi.boolean_mask(data, keys), -1, *all_shapes) - - elif bool_type == _NDARRAY_INT_BOOLEAN_INDEXING: - out = self - for idx in range(pos): - out = out[key[idx]] - data = _reshape_view(out, -1, *remaining_shapes) - after_mask = _reshape_view(_npi.boolean_mask(data, mask), -1, *remaining_shapes) - if pos == len(key) - 1: - return after_mask - # check boundary - for idx in range(pos+1, len(key)): - if key[idx] >= self.shape[idx+mask_ndim-1]: - raise IndexError('index {} on a dimension of {}' - .format(key[idx], self.shape[idx+mask_ndim-1])) - implicit_idces = len(key)+mask_ndim-1 # idces not explictly shown in the key - implicit_shape = self.shape[implicit_idces:] - new_dim = reduce(lambda x, y: x * y, self.shape[pos+mask_ndim:implicit_idces], 1) - new_idx = self._calculate_new_idx(key, self.shape, pos, mask_ndim) - after_reshape = _reshape_view(after_mask, -1, new_dim, *implicit_shape) - return _reshape_view(_npi.take(after_reshape, array([new_idx]), axis=1), -1, *implicit_shape) - - raise NotImplementedError("This boolean indexing type is not supported.") + # pylint: disable=redefined-outer-name + def _get_np_boolean_indexing(self, key, ndim, shape): + """ + There are two types of boolean indices (which are equivalent, + for the most part though). This function will handle single + boolean indexing for higher speed. + If this is not the case, it is instead expanded into (multiple) + integer array indices and will be handled by advanced indexing. + """ + key_shape = key.shape + key_ndim = len(key_shape) + if ndim < key_ndim: + raise IndexError('too many indices, whose ndim = {}, for array with ndim = {}' + .format(key_ndim, ndim)) + for i in range(key_ndim): + if key_shape[i] != shape[i]: + raise IndexError('boolean index did not match indexed array along dimension {};' + ' dimension is {} but corresponding boolean dimension is {}' + .format(i, shape[i], key_shape[i])) + remaining_dims = shape[key_ndim:] + data = _reshape_view(self, -1, *remaining_dims) + key = _reshape_view(key, -1) + return _reshape_view(_npi.boolean_mask(data, key), -1, *remaining_dims) def _set_np_boolean_indexing(self, key, value): - if not isinstance(key, tuple): - key = (key,) - bool_type, pos = self._check_boolean_indexing_type(key) - - mask = key[pos] - mask_shape = mask.shape - mask_ndim = len(mask_shape) - for i in range(mask_ndim): - if mask_shape[i] != self.shape[pos + i]: - raise IndexError('boolean index did not match indexed array along axis {};' - ' size is {} but corresponding boolean size is {}' - .format(pos + i, self.shape[pos + i], mask_shape[i])) - - data = self # when bool_type == _NDARRAY_SLICE_BOOLEAN_INDEXING - if bool_type == _NDARRAY_INT_BOOLEAN_INDEXING: - if pos != len(key) - 1: - raise NotImplementedError('only support boolean array at the end of the idces ' - 'when it is mixed with integers') - for idx in range(pos): - data = data[key[idx]] - pos -= 1 - + """ + There are two types of boolean indices (which are equivalent, + for the most part though). This function will handle single boolean assign for higher speed. + If this is not the case, it is instead expanded into (multiple) + integer array indices and will be handled by advanced assign. + """ if isinstance(value, numeric_types): - _npi.boolean_mask_assign_scalar(data=data, mask=mask, + _npi.boolean_mask_assign_scalar(data=self, mask=key, value=int(value) if isinstance(value, bool) else value, - start_axis=pos, out=data) + start_axis=0, out=self) elif isinstance(value, ndarray): - _npi.boolean_mask_assign_tensor(data=data, mask=mask, value=value, start_axis=pos, out=data) + _npi.boolean_mask_assign_tensor(data=self, mask=key, value=value, start_axis=0, out=self) else: raise NotImplementedError('type %s is not supported.'%(type(value))) @@ -658,12 +581,14 @@ def __getitem__(self, key): >>> x = np.array([1., -1., -2., 3]) >>> x[x < 0] array([-1., -2.]) + + For more imformation related to boolean indexing, please refer to + https://docs.scipy.org/doc/numpy-1.17.0/reference/arrays.indexing.html. """ - # handling possible boolean indexing first ndim = self.ndim # pylint: disable=redefined-outer-name shape = self.shape # pylint: disable=redefined-outer-name if isinstance(key, bool): # otherwise will be treated as 0 and 1 - key = array(key, dtype=_np.bool) + key = array(key, dtype=_np.bool, ctx=self.ctx) if isinstance(key, list): try: new_key = _np.array(key) @@ -674,11 +599,15 @@ def __getitem__(self, key): if isinstance(key, _np.ndarray) and key.dtype == _np.bool_: key = array(key, dtype='bool', ctx=self.ctx) - if ndim == 0: - if isinstance(key, ndarray) and key.dtype == _np.bool: - pass # will handle by function for boolean indexing - elif key != (): - raise IndexError('scalar tensor can only accept `()` as index') + # Handle single boolean index of matching dimensionality and size first for higher speed + # If the boolean array is mixed with other idices, it is instead expanded into (multiple) + # integer array indices and will be handled by advanced indexing. + # Come before the check self.dim == 0 as it also handle the 0-dim case. + if isinstance(key, ndarray) and key.dtype == _np.bool_: + return self._get_np_boolean_indexing(key, ndim, shape) + + if ndim == 0 and key != (): + raise IndexError('scalar tensor can only accept `()` as index') # Handle simple cases for higher speed if isinstance(key, tuple) and len(key) == 0: return self @@ -703,17 +632,33 @@ def __getitem__(self, key): elif key.step == 0: raise ValueError("slice step cannot be zero") - key_before_expaned = key - key = indexing_key_expand_implicit_axes(key, self.shape) + # For 0-d boolean indices: A new axis is added, + # but at the same time no axis is "used". So if we have True, + # we add a new axis (a bit like with np.newaxis). If it is + # False, we add a new axis, but this axis has 0 entries. + # prepend is defined to handle this case. + # prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY/-1 means there is no 0-d boolean scalar + # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE/0 means an zero dim must be expanded + # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE/1 means a new axis must be prepended + key, prepend = indexing_key_expand_implicit_axes(key, self.shape) indexing_dispatch_code = get_indexing_dispatch_code(key) - if indexing_dispatch_code == _NDARRAY_BASIC_INDEXING: - return self._get_np_basic_indexing(key) - elif indexing_dispatch_code == _NDARRAY_EMPTY_TUPLE_INDEXING: + if indexing_dispatch_code == _NDARRAY_EMPTY_TUPLE_INDEXING: + # won't be affected by zero-dim boolean indices return self._get_np_empty_tuple_indexing(key) + elif indexing_dispatch_code == _NDARRAY_BASIC_INDEXING: + if prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE: + return empty((0,) + self._get_np_basic_indexing(key).shape, + dtype=self.dtype, ctx=self.ctx) + if prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE: + key = (_np.newaxis,) + key + return self._get_np_basic_indexing(key) elif indexing_dispatch_code == _NDARRAY_ADVANCED_INDEXING: + if prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE: + return empty((0,) + self._get_np_adanced_indexing(key).shape, + dtype=self.dtype, ctx=self.ctx) + if prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE: + key = (_np.newaxis,) + key return self._get_np_advanced_indexing(key) - elif indexing_dispatch_code == _NDARRAY_BOOLEAN_INDEXING: - return self._get_np_boolean_indexing(key_before_expaned) else: raise RuntimeError @@ -764,16 +709,25 @@ def __setitem__(self, key, value): >>> x array([[ 6., 5., 5.], [ 6., 0., 4.]]) + + For imformation related to boolean indexing, please refer to + https://docs.scipy.org/doc/numpy-1.17.0/reference/arrays.indexing.html. """ if isinstance(value, NDArray) and not isinstance(value, ndarray): raise TypeError('Cannot assign mx.nd.NDArray to mxnet.numpy.ndarray') if isinstance(key, bool): # otherwise will be treated as 0 and 1 key = array(key, dtype=_np.bool) + + # Handle single boolean assign of matching dimensionality and size first for higher speed + # If the boolean array is mixed with other idices, it is instead expanded into (multiple) + # integer array indices and will be handled by advanced assign. + # Come before the check self.dim == 0 as it also handle the 0-dim case. + if isinstance(key, ndarray) and key.dtype == _np.bool: + return self._set_np_boolean_indexing(key, value) + # handle basic and advanced indexing if self.ndim == 0: - if isinstance(key, ndarray) and key.dtype == _np.bool: - pass # will be handled by boolean indexing - elif not isinstance(key, tuple) or len(key) != 0: + if not isinstance(key, tuple) or len(key) != 0: raise IndexError('scalar tensor can only accept `()` as index') if isinstance(value, numeric_types): self._full(value) @@ -788,8 +742,18 @@ def __setitem__(self, key, value): else: raise ValueError('setting an array element with a sequence.') else: - key_before_expaned = key - key = indexing_key_expand_implicit_axes(key, self.shape) + # For 0-d boolean indices: A new axis is added, + # but at the same time no axis is "used". So if we have True, + # we add a new axis (a bit like with np.newaxis). If it is + # False, we add a new axis, but this axis has 0 entries. + # prepend is defined to handle this case. + # prepend == _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY/-1 means there is no 0-d boolean scalar + # prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE/0 means an zero dim must be expanded + # prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE/1 means a new axis must be expanded + # prepend actually has no influence on __setitem__ + key, prepend = indexing_key_expand_implicit_axes(key, self.shape) + if prepend == _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE: + return # no action is needed slc_key = tuple(idx for idx in key if idx is not None) if len(slc_key) < self.ndim: raise RuntimeError( @@ -809,8 +773,6 @@ def __setitem__(self, key, value): pass # no action needed elif indexing_dispatch_code == _NDARRAY_ADVANCED_INDEXING: self._set_np_advanced_indexing(key, value) - elif indexing_dispatch_code == _NDARRAY_BOOLEAN_INDEXING: - return self._set_np_boolean_indexing(key_before_expaned, value) else: raise ValueError( 'Indexing NDArray with index {} of type {} is not supported' diff --git a/src/operator/numpy/np_nonzero_op.cc b/src/operator/numpy/np_nonzero_op.cc index 0eaf0878a24a..2b9de76e6a0d 100644 --- a/src/operator/numpy/np_nonzero_op.cc +++ b/src/operator/numpy/np_nonzero_op.cc @@ -66,7 +66,7 @@ void NonzeroForwardCPU(const nnvm::NodeAttrs& attrs, CHECK_LE(in.shape().ndim(), MAXDIM) << "ndim of input cannot larger than " << MAXDIM; // 0-dim if (0 == in.shape().ndim()) { - MSHADOW_TYPE_SWITCH(in.dtype(), DType, { + MSHADOW_TYPE_SWITCH_WITH_BOOL(in.dtype(), DType, { DType* in_dptr = in.data().dptr(); if (*in_dptr) { mxnet::TShape s(2, 1); diff --git a/src/operator/tensor/indexing_op.cc b/src/operator/tensor/indexing_op.cc index 2a2dd76ac95c..1303cb91b1c7 100644 --- a/src/operator/tensor/indexing_op.cc +++ b/src/operator/tensor/indexing_op.cc @@ -480,7 +480,7 @@ void GatherNDForwardCPU(const nnvm::NodeAttrs& attrs, strides[i] = stride; mshape[i] = dshape[i]; } - MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { // output data type switch + MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[0].type_flag_, DType, { // output data type switch MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // indices data type switch // check whether indices are out of bound IType* idx_ptr = inputs[1].dptr(); diff --git a/src/operator/tensor/indexing_op.cu b/src/operator/tensor/indexing_op.cu index 2d8789daa974..6904656b304b 100644 --- a/src/operator/tensor/indexing_op.cu +++ b/src/operator/tensor/indexing_op.cu @@ -486,7 +486,7 @@ void GatherNDForwardGPU(const nnvm::NodeAttrs& attrs, strides[i] = stride; mshape[i] = dshape[i]; } - MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { // output data type switch + MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[0].type_flag_, DType, { // output data type switch MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // indices data type switch // check whether indices are out of bound IType* idx_ptr = inputs[1].dptr(); diff --git a/src/operator/tensor/indexing_op.h b/src/operator/tensor/indexing_op.h index 5449fbe4afbd..2b048813a464 100644 --- a/src/operator/tensor/indexing_op.h +++ b/src/operator/tensor/indexing_op.h @@ -1447,8 +1447,8 @@ void ScatterNDForward(const nnvm::NodeAttrs& attrs, if (kWriteTo == req[0]) { Fill(s, outputs[0], req[0], 0); } - MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { // output data type switch - MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // indices data type switch + MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[0].type_flag_, DType, { // output data type switch + MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[1].type_flag_, IType, { // indices data type switch mxnet_op::Kernel::Launch( s, N, req[0], N, M, K, strides, outputs[0].dptr(), inputs[0].dptr(), inputs[1].dptr()); diff --git a/tests/python/unittest/test_numpy_ndarray.py b/tests/python/unittest/test_numpy_ndarray.py index 9ef5d79fa4be..3ce53c6a6e80 100644 --- a/tests/python/unittest/test_numpy_ndarray.py +++ b/tests/python/unittest/test_numpy_ndarray.py @@ -1194,6 +1194,24 @@ def test_boolean_indexing_tuple(): assert same(a[0, b].asnumpy(), _np_a[0, _np_b]) assert same(a[b, 1].asnumpy(), _np_a[_np_b, 1]) + a = np.arange(12).reshape(4,3) + b = np.array([1.,2.,3.]) + _np_a = a.asnumpy() + _np_b = b.asnumpy() + assert same(a[:, b > 2].shape, _np_a[:, _np_b > 2].shape) + assert same(a[:, b > 2].asnumpy(), _np_a[:, _np_b > 2]) + + a = np.array([[1,2,3],[3,4,5]]) + _np_a = a.asnumpy() + assert same(a[:,a[1,:] > 0].shape, _np_a[:,_np_a[1,: ] > 0].shape) + assert same(a[:,a[1,:] > 0].asnumpy(), _np_a[:,_np_a[1,: ] > 0]) + + a = np.ones((3,2), dtype='bool') + b = np.array([1,2,3]) + _np_a = a.asnumpy() + _np_b = b.asnumpy() + assert same(a[b > 1].asnumpy(), _np_a[_np_b > 1]) + def test_boolean_indexing_assign(): # test boolean indexing assign shape = (3, 2, 3) @@ -1208,11 +1226,11 @@ def test_boolean_indexing_assign(): np_data[np_mask] = 1 mx_data[mx_mask] = 1 assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) - # not supported at this moment - # only support boolean array at the end of the idces when it is mixed with integers - # np_data[np_mask, 1] = 2 - # mx_data[mx_mask, 1] = 2 - # assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) + + np_data[np_mask, 1] = 2 + mx_data[mx_mask, 1] = 2 + assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) + np_data[np_mask, :] = 3 mx_data[mx_mask, :] = 3 assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) @@ -1227,6 +1245,14 @@ def test_boolean_indexing_assign(): mx_data[:, mx_mask] = 6 assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) + np_data[0, True, True, np_mask] = 7 + mx_data[0, True, True, mx_mask] = 7 + assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) + + np_data[False, 1] = 8 + mx_data[False, 1] = 8 + assert_almost_equal(mx_data.asnumpy(), np_data, rtol=1e-3, atol=1e-5, use_broadcast=False) + def test_boolean_indexing_autograd(): a = np.random.uniform(size=(3, 4, 5)) a.attach_grad()