diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index 964d533b2387..ca960ff6edec 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -33,7 +33,7 @@ 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean', - 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] + 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'flip'] @set_module('mxnet.ndarray.numpy') @@ -2363,3 +2363,71 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): # pylint: 0.2025 """ return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + +@set_module('mxnet.ndarray.numpy') +def flip(x, axis=None, out=None, **kwargs): + r""" + flip(x, axis=None, out=None) + + Reverse the order of elements in an array along the given axis. + + The shape of the array is preserved, but the elements are reordered. + + Parameters + ---------- + m : ndarray or scalar + Input array. + axis : None or int or tuple of ints, optional + Axis or axes along which to flip over. The default, + axis=None, will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + out : ndarray or scalar, optional + Alternative output array in which to place the result. It must have + the same shape and type as the expected output. + + Returns + ------- + out : ndarray or scalar + A view of `m` with the entries of axis reversed. Since a view is + returned, this operation is done in constant time. + + Examples + -------- + >>> A = np.arange(8).reshape((2,2,2)) + >>> A + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + >>> np.flip(A, 0) + array([[[4, 5], + [6, 7]], + [[0, 1], + [2, 3]]]) + >>> np.flip(A, 1) + array([[[2, 3], + [0, 1]], + [[6, 7], + [4, 5]]]) + >>> np.flip(A) + array([[[7, 6], + [5, 4]], + [[3, 2], + [1, 0]]]) + >>> np.flip(A, (0, 2)) + array([[[5, 4], + [7, 6]], + [[1, 0], + [3, 2]]]) + """ + from ...numpy import ndarray + if isinstance(x, numeric_types): + return _np.flip(x, axis, **kwargs) + elif isinstance(x, ndarray): + return _npi.flip(x, axis, out=out, **kwargs) + else: + raise TypeError('type {} not supported'.format(str(type(x)))) diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 7d6e81a2d4a5..41efcc8022e4 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -52,7 +52,7 @@ 'degrees', 'log2', 'log1p', 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', - 'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] + 'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'flip'] # Return code for dispatching indexing function call _NDARRAY_UNSUPPORTED_INDEXING = -1 @@ -3808,3 +3808,65 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None): 0.2025 """ return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) + + +@set_module('mxnet.numpy') +def flip(x, axis=None, out=None, **kwargs): + r""" + flip(x, axis=None, out=None) + + Reverse the order of elements in an array along the given axis. + + The shape of the array is preserved, but the elements are reordered. + + Parameters + ---------- + m : ndarray or scalar + Input array. + axis : None or int or tuple of ints, optional + Axis or axes along which to flip over. The default, + axis=None, will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + out : ndarray or scalar, optional + Alternative output array in which to place the result. It must have + the same shape and type as the expected output. + + Returns + ------- + out : ndarray or scalar + A view of `m` with the entries of axis reversed. Since a view is + returned, this operation is done in constant time. + + Examples + -------- + >>> A = np.arange(8).reshape((2,2,2)) + >>> A + array([[[0, 1], + [2, 3]], + [[4, 5], + [6, 7]]]) + >>> np.flip(A, 0) + array([[[4, 5], + [6, 7]], + [[0, 1], + [2, 3]]]) + >>> np.flip(A, 1) + array([[[2, 3], + [0, 1]], + [[6, 7], + [4, 5]]]) + >>> np.flip(A) + array([[[7, 6], + [5, 4]], + [[3, 2], + [1, 0]]]) + >>> np.flip(A, (0, 2)) + array([[[5, 4], + [7, 6]], + [[1, 0], + [3, 2]]]) + """ + return _mx_nd_np.flip(x, axis, out=out, **kwargs) diff --git a/python/mxnet/symbol/numpy/_symbol.py b/python/mxnet/symbol/numpy/_symbol.py index ec4f6a4dd741..6e37168895d5 100644 --- a/python/mxnet/symbol/numpy/_symbol.py +++ b/python/mxnet/symbol/numpy/_symbol.py @@ -35,7 +35,7 @@ 'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean', - 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var'] + 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'flip'] def _num_outputs(sym): @@ -2678,4 +2678,42 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=False): # pylint: return _npi.var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, out=out) +@set_module('mxnet.symbol.numpy') +def flip(x, axis=None, out=None, **kwargs): + r""" + flip(x, axis=None, out=None) + + Reverse the order of elements in an array along the given axis. + + The shape of the array is preserved, but the elements are reordered. + + Parameters + ---------- + m : _Symbol or scalar + Input array. + axis : None or int or tuple of ints, optional + Axis or axes along which to flip over. The default, + axis=None, will flip over all of the axes of the input array. + If axis is negative it counts from the last to the first axis. + + If axis is a tuple of ints, flipping is performed on all of the axes + specified in the tuple. + out : _Symbol or scalar, optional + Alternative output array in which to place the result. It must have + the same shape and type as the expected output. + + Returns + ------- + out : _Symbol or scalar + A view of `m` with the entries of axis reversed. Since a view is + returned, this operation is done in constant time. + """ + if isinstance(x, numeric_types): + return _np.flip(x, axis, **kwargs) + elif isinstance(x, _Symbol): + return _npi.flip(x, axis, out=out, **kwargs) + else: + raise TypeError('type {} not supported'.format(str(type(x)))) + + _set_np_symbol_class(_Symbol) diff --git a/src/operator/numpy/np_matrix_op-inl.h b/src/operator/numpy/np_matrix_op-inl.h index 6d3d9ea5ec85..10b29d50ac18 100644 --- a/src/operator/numpy/np_matrix_op-inl.h +++ b/src/operator/numpy/np_matrix_op-inl.h @@ -60,6 +60,81 @@ void NumpyTranspose(const nnvm::NodeAttrs& attrs, } } +struct FlipParam : public dmlc::Parameter { + mxnet::Tuple axis; + DMLC_DECLARE_PARAMETER(FlipParam) { + DMLC_DECLARE_FIELD(axis) + .describe("The axis which to flip elements."); + } +}; + +struct flip0dim_shared_kernel { + template + MSHADOW_XINLINE static void Map(int i, + DType* out_data, + const DType* in_data) { + out_data[i] = in_data[i]; + } +}; + +#define FLIP_MAX_DIM 10 +#define FLIP_MIN_DIM -1 + +template +void NumpyFlipForwardImpl(const OpContext& ctx, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& stride_, + const std::vector& trailing_, + const index_t& flip_index); + +template +void NumpyFlipForward(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + const FlipParam& param = nnvm::get(attrs.parsed); + mxnet::Tuple axistemp; + CHECK_EQ(inputs[0].type_flag_, outputs[0].type_flag_); + CHECK_LT(param.axis.ndim(), FLIP_MAX_DIM); + CHECK_GE(param.axis.ndim(), FLIP_MIN_DIM); + if (param.axis.ndim() == FLIP_MIN_DIM) { + if (inputs[0].shape_.ndim() == 0) { + mshadow::Stream *s = ctx.get_stream(); + MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { + mxnet_op::Kernel::Launch(s, inputs[0].Size(), + outputs[0].dptr(), inputs[0].dptr()); + }); + return; + } + std::vector temp; + for (int i = 0; i < inputs[0].shape_.ndim(); i++) { + temp.push_back(i); + } + axistemp.assign(temp.begin(), temp.end()); + } else { + axistemp = param.axis; + } + + const mxnet::TShape& ishape = inputs[0].shape_; + if (ishape.ProdShape(0, ishape.ndim()) == 0) { + return; // zero shape + } + std::vector stride_(axistemp.ndim()); + std::vector trailing_(axistemp.ndim()); + index_t flip_index = 0; + for (int axis : axistemp) { + CHECK_LT(axis, ishape.ndim()); + stride_[flip_index] = ishape[axis]; + trailing_[flip_index] = 1; + for (int i2 = axis + 1; i2 < ishape.ndim(); ++i2) { + trailing_[flip_index] *= ishape[i2]; + } + flip_index++; + } + NumpyFlipForwardImpl(ctx, inputs, outputs, stride_, trailing_, flip_index); +} } // namespace op } // namespace mxnet diff --git a/src/operator/numpy/np_matrix_op.cc b/src/operator/numpy/np_matrix_op.cc index 5ad6c8908017..1176bcdb5cb3 100644 --- a/src/operator/numpy/np_matrix_op.cc +++ b/src/operator/numpy/np_matrix_op.cc @@ -345,5 +345,51 @@ Examples:: .add_argument("data", "NDArray-or-Symbol[]", "List of arrays to stack") .add_arguments(StackParam::__FIELDS__()); +template<> +void NumpyFlipForwardImpl(const OpContext& ctx, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& stride_, + const std::vector& trailing_, + const index_t& flip_index) { + mshadow::Stream *s = ctx.get_stream(); + MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { + mxnet_op::Kernel::Launch(s, inputs[0].Size(), flip_index, + inputs[0].dptr(), outputs[0].dptr(), + stride_.data(), trailing_.data()); + }); +} + +DMLC_REGISTER_PARAMETER(FlipParam); + +NNVM_REGISTER_OP(_npi_flip) +.set_num_outputs(1) +.set_num_inputs(1) +.set_attr_parser(ParamParser) +.set_attr("FListInputNames", +[](const NodeAttrs& attrs) { + return std::vector {"data"}; +}) +.set_attr("FResourceRequest", +[](const NodeAttrs& attrs) { + return std::vector {ResourceRequest::kTempSpace}; +}) +.set_attr("FInferShape", ElemwiseShape<1, 1>) +.set_attr("FInferType", ElemwiseType<1, 1>) +.set_attr("FCompute", NumpyFlipForward) +.set_attr("FGradient", ElemwiseGradUseNone{"_backward_npi_flip"}) +.add_argument("data", "NDArray-or-Symbol", "Input data array") +.add_arguments(FlipParam::__FIELDS__()); + +NNVM_REGISTER_OP(_backward_npi_flip) +.set_num_inputs(1) +.set_num_outputs(1) +.set_attr_parser(ParamParser) +.set_attr("TIsBackward", true) +.set_attr("FResourceRequest", +[](const NodeAttrs& attrs) { + return std::vector {ResourceRequest::kTempSpace}; +}) +.set_attr("FCompute", NumpyFlipForward); } // namespace op } // namespace mxnet diff --git a/src/operator/numpy/np_matrix_op.cu b/src/operator/numpy/np_matrix_op.cu index 4ba527deca09..5b53f12137ee 100644 --- a/src/operator/numpy/np_matrix_op.cu +++ b/src/operator/numpy/np_matrix_op.cu @@ -46,5 +46,39 @@ NNVM_REGISTER_OP(_backward_np_concat) NNVM_REGISTER_OP(_npi_stack) .set_attr("FCompute", StackOpForward); +template<> +void NumpyFlipForwardImpl(const OpContext& ctx, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& stride_, + const std::vector& trailing_, + const index_t& flip_index) { + mshadow::Stream *s = ctx.get_stream(); + mshadow::Tensor workspace = + ctx.requested[0].get_space_typed( + mshadow::Shape1(flip_index * sizeof(index_t) * 2), s); + + auto stride_workspace = workspace.dptr_; + auto trailing_workspace = workspace.dptr_ + flip_index * sizeof(index_t); + + cudaMemcpyAsync(stride_workspace, thrust::raw_pointer_cast(stride_.data()), + stride_.size() * sizeof(index_t), + cudaMemcpyHostToDevice, mshadow::Stream::GetStream(s)); + cudaMemcpyAsync(trailing_workspace, thrust::raw_pointer_cast(trailing_.data()), + trailing_.size() * sizeof(index_t), + cudaMemcpyHostToDevice, mshadow::Stream::GetStream(s)); + + MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { + mxnet_op::Kernel::Launch(s, inputs[0].Size(), flip_index, + inputs[0].dptr(), outputs[0].dptr(), + reinterpret_cast(stride_workspace), reinterpret_cast(trailing_workspace)); + }); +} + +NNVM_REGISTER_OP(_npi_flip) +.set_attr("FCompute", NumpyFlipForward); + +NNVM_REGISTER_OP(_backward_npi_flip) +.set_attr("FCompute", NumpyFlipForward); } // namespace op } // namespace mxnet diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 1c9ebbb32866..dc13d4d26c7d 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -1744,6 +1744,66 @@ def test_indexing_mode(sampler, set_size, samples_size, replace, weight=None): test_indexing_mode(test_choice_weighted, num_classes, num_classes // 2, replace, weight) +@with_seed() +@use_np +def test_np_flip(): + class TestFlip(HybridBlock): + def __init__(self, axis): + super(TestFlip, self).__init__() + self.axis = axis + + def hybrid_forward(self, F, x): + return F.np.flip(x, self.axis) + + shapes1 = [(1, 2, 3), (1, 0), (3, 0, 2), (), (1,)] + shapes2 = [(2, 2, 3, 3), (1, 1, 2, 4)] + types = ['int32', 'int64', 'float16', 'float32', 'double'] + for hybridize in [True, False]: + for oneType in types: + rtol=1e-3 + atol=1e-5 + for shape in shapes1: + axis = None + test_flip = TestFlip(axis) + if hybridize: + test_flip.hybridize() + x = rand_ndarray(shape, dtype=oneType).as_np_ndarray() + x.attach_grad() + np_out = _np.flip(x.asnumpy(), axis) + with mx.autograd.record(): + mx_out = test_flip(x) + assert mx_out.shape == np_out.shape + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol) + mx_out.backward() + np_backward = _np.ones(np_out.shape) + assert_almost_equal(x.grad.asnumpy(), np_backward, rtol=rtol, atol=atol) + + # Test imperative once again + mx_out = np.flip(x, axis) + np_out = _np.flip(x.asnumpy(), axis) + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol) + for shape in shapes2: + for axis in [None, 2, (0, 1), (1, 3), (0, 2)]: + test_flip = TestFlip(axis) + if hybridize: + test_flip.hybridize() + x = rand_ndarray(shape, dtype=oneType).as_np_ndarray() + x.attach_grad() + np_out = _np.flip(x.asnumpy(), axis) + with mx.autograd.record(): + mx_out = test_flip(x) + assert mx_out.shape == np_out.shape + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol) + mx_out.backward() + np_backward = _np.ones(np_out.shape) + assert_almost_equal(x.grad.asnumpy(), np_backward, rtol=rtol, atol=atol) + + # Test imperative once again + mx_out = np.flip(x, axis) + np_out = _np.flip(x.asnumpy(), axis) + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=rtol, atol=atol) + + if __name__ == '__main__': import nose nose.runmodule()