diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index d31681113933..295c83ef47a0 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -35,11 +35,11 @@ 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'histogram', 'eye', 'linspace', 'logspace', 'expand_dims', 'tile', 'arange', 'split', 'vsplit', 'concatenate', 'append', 'stack', 'vstack', 'column_stack', 'dstack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', - 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'hanning', 'hamming', 'blackman', 'flip', - 'around', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', - 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less', 'greater_equal', - 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'nonzero', 'shares_memory', 'may_share_memory', - 'diff', 'resize', 'nan_to_num', 'where'] + 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'unravel_index', 'hanning', 'hamming', 'blackman', + 'flip', 'around', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', + 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less', + 'greater_equal', 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'nonzero', 'shares_memory', + 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where'] @set_module('mxnet.ndarray.numpy') def shape(a): @@ -3828,6 +3828,44 @@ def ravel(x, order='C'): raise TypeError('type {} not supported'.format(str(type(x)))) +def unravel_index(indices, shape, order='C'): # pylint: disable=redefined-outer-name + """ + Converts a flat index or array of flat indices into a tuple of coordinate arrays. + + Parameters: + ------------- + indices : array_like + An integer array whose elements are indices into the flattened version of an array of dimensions shape. + Before version 1.6.0, this function accepted just one index value. + shape : tuple of ints + The shape of the array to use for unraveling indices. + + Returns: + ------------- + unraveled_coords : ndarray + Each row in the ndarray has the same shape as the indices array. + Each column in the ndarray represents the unravelled index + + Examples: + ------------- + >>> np.unravel_index([22, 41, 37], (7,6)) + ([3. 6. 6.] + [4. 5. 1.]) + >>> np.unravel_index(1621, (6,7,8,9)) + (3, 1, 4, 1) + """ + if order == 'C': + if isinstance(indices, numeric_types): + return _np.unravel_index(indices, shape) + ret = _npi.unravel_index_fallback(indices, shape=shape) + ret_list = [] + for item in ret: + ret_list += [item] + return tuple(ret_list) + else: + raise NotImplementedError('Do not support column-major (Fortran-style) order at this moment') + + @set_module('mxnet.ndarray.numpy') def hanning(M, dtype=_np.float32, ctx=None): r"""Return the Hanning window. diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 3ad9254eb0e6..1c6873d342d1 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -53,11 +53,11 @@ 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'append', 'tensordot', 'histogram', 'eye', 'linspace', 'logspace', 'expand_dims', 'tile', 'arange', 'split', 'vsplit', 'concatenate', 'stack', 'vstack', 'column_stack', 'dstack', 'mean', 'maximum', 'minimum', - 'swapaxes', 'clip', 'argmax', 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'hanning', 'hamming', - 'blackman', 'flip', 'around', 'arctan2', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', - 'unique', 'lcm', 'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', - 'greater', 'less', 'greater_equal', 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'nonzero', - 'shares_memory', 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where'] + 'swapaxes', 'clip', 'argmax', 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'unravel_index', + 'hanning', 'hamming', 'blackman', 'flip', 'around', 'arctan2', 'hypot', 'bitwise_xor', 'bitwise_or', + 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', + 'equal', 'not_equal', 'greater', 'less', 'greater_equal', 'less_equal', 'hsplit', 'rot90', 'einsum', + 'true_divide', 'nonzero', 'shares_memory', 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where'] # Return code for dispatching indexing function call _NDARRAY_UNSUPPORTED_INDEXING = -1 @@ -5759,6 +5759,36 @@ def ravel(x, order='C'): return _mx_nd_np.ravel(x, order) +def unravel_index(indices, shape, order='C'): # pylint: disable=redefined-outer-name + """ + Converts a flat index or array of flat indices into a tuple of coordinate arrays. + + Parameters: + ------------- + indices : array_like + An integer array whose elements are indices into the flattened version of an array of dimensions shape. + Before version 1.6.0, this function accepted just one index value. + shape : tuple of ints + The shape of the array to use for unraveling indices. + order : Only row-major is supported currently. + + Returns: + ------------- + unraveled_coords : ndarray + Each row in the ndarray has the same shape as the indices array. + Each column in the ndarray represents the unravelled index + + Examples: + ------------- + >>> np.unravel_index([22, 41, 37], (7,6)) + [[3. 6. 6.] + [4. 5. 1.]] + >>> np.unravel_index(1621, (6,7,8,9)) + [3, 1, 4, 1] + """ + return _mx_nd_np.unravel_index(indices, shape, order=order) + + @set_module('mxnet.numpy') def hanning(M, dtype=_np.float32, ctx=None): r"""Return the Hanning window. diff --git a/python/mxnet/numpy_dispatch_protocol.py b/python/mxnet/numpy_dispatch_protocol.py index 13e186f4008a..1f68ca3c522a 100644 --- a/python/mxnet/numpy_dispatch_protocol.py +++ b/python/mxnet/numpy_dispatch_protocol.py @@ -121,6 +121,7 @@ def _run_with_array_ufunc_proto(*args, **kwargs): 'tile', 'transpose', 'unique', + 'unravel_index', 'var', 'vdot', 'vstack', diff --git a/python/mxnet/numpy_op_fallback.py b/python/mxnet/numpy_op_fallback.py index b98a211c7169..f3fd2bf6ae6c 100644 --- a/python/mxnet/numpy_op_fallback.py +++ b/python/mxnet/numpy_op_fallback.py @@ -80,3 +80,37 @@ def infer_shape(self, in_shape): def create_operator(self, ctx, in_shapes, in_dtypes): return Resize(self._new_shape) + + +@use_np +class Unravel_index(operator.CustomOp): + """Fallback to NumPy Unravel_index operator.""" + def __init__(self, shape): + super(Unravel_index, self).__init__() + self._shape = shape + + def forward(self, is_train, req, in_data, out_data, aux): + out = np.unravel_index(in_data[0].asnumpy(), self._shape) + self.assign(out_data[0], req[0], _mx_np.array(out, dtype=out[0].dtype, ctx=out_data[0].ctx)) + + def backward(self, req, out_grad, in_data, out_data, in_grad, aux): + raise NotImplementedError('Operator Unravel_index does not support gradient computation') + + +@register('unravel_index_fallback') +class Unravel_indexProp(operator.CustomOpProp): + """Fallback unravel_index operator properties.""" + def __init__(self, shape): + super(Unravel_indexProp, self).__init__(need_top_grad=True) + self._shape = ast.literal_eval(shape) + + def list_arguments(self): + return ['indices'] + + def infer_shape(self, in_shape): + dim_list = (1,) if np.isscalar(self._shape) else (len(self._shape),) + out_shape = dim_list + tuple(in_shape[0]) + return (in_shape[0],), (out_shape,), () + + def create_operator(self, ctx, in_shapes, in_dtypes): + return Unravel_index(self._shape) diff --git a/python/mxnet/symbol/numpy/_symbol.py b/python/mxnet/symbol/numpy/_symbol.py index 734242f4e86d..3790078bd397 100644 --- a/python/mxnet/symbol/numpy/_symbol.py +++ b/python/mxnet/symbol/numpy/_symbol.py @@ -37,11 +37,11 @@ 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'histogram', 'eye', 'linspace', 'logspace', 'expand_dims', 'tile', 'arange', 'split', 'vsplit', 'concatenate', 'append', 'stack', 'vstack', 'column_stack', 'dstack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', - 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'hanning', 'hamming', 'blackman', 'flip', - 'around', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', - 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less', 'greater_equal', - 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'shares_memory', 'may_share_memory', 'diff', - 'resize', 'nan_to_num', 'where'] + 'argmin', 'std', 'var', 'indices', 'copysign', 'ravel', 'unravel_index', 'hanning', 'hamming', 'blackman', + 'flip', 'around', 'hypot', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', + 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', 'greater', 'less', + 'greater_equal', 'less_equal', 'hsplit', 'rot90', 'einsum', 'true_divide', 'shares_memory', + 'may_share_memory', 'diff', 'resize', 'nan_to_num', 'where'] def _num_outputs(sym): return len(sym.as_nd_ndarray()) @@ -3657,6 +3657,38 @@ def ravel(x, order='C'): raise TypeError('type {} not supported'.format(str(type(x)))) +def unravel_index(indices, shape, order='C'): # pylint: disable=redefined-outer-name + """ + Converts a flat index or array of flat indices into a tuple of coordinate arrays. + + Parameters: + ------------- + indices : array_like + An integer array whose elements are indices into the flattened version of an array of dimensions shape. + Before version 1.6.0, this function accepted just one index value. + shape : tuple of ints + The shape of the array to use for unraveling indices. + + Returns: + ------------- + unraveled_coords : ndarray + Each row in the ndarray has the same shape as the indices array. + Each column in the ndarray represents the unravelled index + + Examples: + ------------- + >>> np.unravel_index([22, 41, 37], (7,6)) + ([3. 6. 6.] + [4. 5. 1.]) + >>> np.unravel_index(1621, (6,7,8,9)) + (3, 1, 4, 1) + """ + if order == 'C': + return _npi.unravel_index_fallback(indices, shape=shape) + else: + raise NotImplementedError('Don not support column-major (Fortran-style) order at this moment') + + @set_module('mxnet.symbol.numpy') def hanning(M, dtype=_np.float32, ctx=None): r"""Return the Hanning window. diff --git a/tests/python/unittest/test_numpy_interoperability.py b/tests/python/unittest/test_numpy_interoperability.py index 206686f68de8..930ad5260430 100644 --- a/tests/python/unittest/test_numpy_interoperability.py +++ b/tests/python/unittest/test_numpy_interoperability.py @@ -57,6 +57,15 @@ def get_workloads(name): return OpArgMngr._args.get(name, None) +def _add_workload_unravel_index(): + OpArgMngr.add_workload('unravel_index', indices=np.array([2],dtype=_np.int64), shape=(2, 2)) + OpArgMngr.add_workload('unravel_index', np.array([(2*3 + 1)*6 + 4], dtype=_np.int64), (4, 3, 6)) + OpArgMngr.add_workload('unravel_index', np.array([22, 41, 37], dtype=_np.int32), (7, 6)) + OpArgMngr.add_workload('unravel_index', np.array([1621],dtype=_np.uint8), (6, 7, 8, 9)) + OpArgMngr.add_workload('unravel_index', np.array([],dtype=_np.int64), (10, 3, 5)) + OpArgMngr.add_workload('unravel_index', np.array([3], dtype=_np.int32), (2,2)) + + def _add_workload_diag(): def get_mat(n): data = _np.arange(n) @@ -1310,6 +1319,7 @@ def _prepare_workloads(): _add_workload_copy() _add_workload_cumsum() _add_workload_ravel() + _add_workload_unravel_index() _add_workload_diag() _add_workload_diagflat() _add_workload_dot() diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 608116d893f5..9b7f7036bcda 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -4998,6 +4998,58 @@ def hybrid_forward(self, F, x): assert_almost_equal(y.asnumpy(), expected, use_broadcast=False) +@with_seed() +@use_np +def test_np_unravel_index(): + class TestUnravel_index(HybridBlock): + def __init__(self, shape, order='C') : + super(TestUnravel_index, self).__init__() + self._shape = shape + self._order = order + + def hybrid_forward(self, F, a): + return F.np.unravel_index(a, self._shape, self._order) + + in_shapes = [ + 2, 5, + (), (1,), (4,), + (2, 2), (2, 4), (3, 5), + (2, 2, 2), (2, 3, 2), (2, 3, 4), + ] + unravel_shapes = [ + 10, (15,), + (3, 4), (4, 5), + (2,3,4) + ] + dtypes = [np.uint8, np.int8, np.int32, np.int64] + for hybridize, ishape, dtype, rshape in itertools.product([False, True], in_shapes, dtypes, unravel_shapes): + rtol = 1e-2 if dtype == np.float16 else 1e-3 + atol = 1e-4 if dtype == np.float16 else 1e-5 + test_unravel_index = TestUnravel_index(rshape) + if hybridize: + test_unravel_index.hybridize() + if type(ishape) == int and hybridize: + x = np.array([ishape], dtype=dtype) + np_out = _np.unravel_index(x.asnumpy(), rshape) + else: + x = np.random.uniform(0, 8, size=ishape).astype(dtype) + np_out = _np.unravel_index(x.asnumpy(), rshape) + mx_out = test_unravel_index(x) + assert len(mx_out) == len(np_out) + for elem_mx, elem_np in zip(mx_out, np_out): + assert elem_mx.asnumpy().shape == elem_np.shape + assert_almost_equal(elem_mx.asnumpy(), elem_np, rtol=rtol, atol=atol) + # no backward function for unravel_index operator + + # Test imperative once again + mx_out = np.unravel_index(x, rshape) + np_out = _np.unravel_index(x.asnumpy(), rshape) + assert len(mx_out) == len(np_out) + for elem_mx, elem_np in zip(mx_out, np_out): + assert elem_mx.asnumpy().shape == elem_np.shape + assert_almost_equal(elem_mx.asnumpy(), elem_np, rtol=rtol, atol=atol) + + if __name__ == '__main__': import nose nose.runmodule()