diff --git a/benchmark/python/ffi/benchmark_ffi.py b/benchmark/python/ffi/benchmark_ffi.py index 88af3cf3d55e..cf438f849d98 100644 --- a/benchmark/python/ffi/benchmark_ffi.py +++ b/benchmark/python/ffi/benchmark_ffi.py @@ -55,6 +55,8 @@ def prepare_workloads(): OpArgMngr.add_workload("cumsum", pool['3x2'], axis=0, out=pool['3x2']) OpArgMngr.add_workload("add", pool['2x2'], pool['2x2']) OpArgMngr.add_workload("random.uniform", low=0, high=1, size=1) + OpArgMngr.add_workload("roll", pool["2x2"], 1, axis=0) + OpArgMngr.add_workload("rot90", pool["2x2"], 2) def benchmark_helper(f, *args, **kwargs): diff --git a/python/mxnet/_numpy_op_doc.py b/python/mxnet/_numpy_op_doc.py index 8dfc0867cdb2..3d80ce03b44e 100644 --- a/python/mxnet/_numpy_op_doc.py +++ b/python/mxnet/_numpy_op_doc.py @@ -538,70 +538,6 @@ def _np_reshape(a, newshape, order='C', out=None): """ -def _np_roll(a, shift, axis=None): - """ - Roll array elements along a given axis. - - Elements that roll beyond the last position are re-introduced at - the first. - - Parameters - ---------- - a : ndarray - Input array. - shift : int or tuple of ints - The number of places by which elements are shifted. If a tuple, - then `axis` must be a tuple of the same size, and each of the - given axes is shifted by the corresponding number. If an int - while `axis` is a tuple of ints, then the same value is used for - all given axes. - axis : int or tuple of ints, optional - Axis or axes along which elements are shifted. By default, the - array is flattened before shifting, after which the original - shape is restored. - - Returns - ------- - res : ndarray - Output array, with the same shape as `a`. - - Notes - ----- - Supports rolling over multiple dimensions simultaneously. - - Examples - -------- - >>> x = np.arange(10) - >>> np.roll(x, 2) - array([8., 9., 0., 1., 2., 3., 4., 5., 6., 7.]) - >>> np.roll(x, -2) - array([2., 3., 4., 5., 6., 7., 8., 9., 0., 1.]) - - >>> x2 = np.reshape(x, (2,5)) - >>> x2 - array([[0., 1., 2., 3., 4.], - [5., 6., 7., 8., 9.]]) - >>> np.roll(x2, 1) - array([[9., 0., 1., 2., 3.], - [4., 5., 6., 7., 8.]]) - >>> np.roll(x2, -1) - array([[1., 2., 3., 4., 5.], - [6., 7., 8., 9., 0.]]) - >>> np.roll(x2, 1, axis=0) - array([[5., 6., 7., 8., 9.], - [0., 1., 2., 3., 4.]]) - >>> np.roll(x2, -1, axis=0) - array([[5., 6., 7., 8., 9.], - [0., 1., 2., 3., 4.]]) - >>> np.roll(x2, 1, axis=1) - array([[4., 0., 1., 2., 3.], - [9., 5., 6., 7., 8.]]) - >>> np.roll(x2, -1, axis=1) - array([[1., 2., 3., 4., 0.], - [6., 7., 8., 9., 5.]]) - """ - - def _np_trace(a, offset=0, axis1=0, axis2=1, out=None): """ Return the sum along diagonals of the array. diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index 3d30333d6da2..a698ab1c5c02 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -43,7 +43,7 @@ 'diag_indices_from', 'hanning', 'hamming', 'blackman', 'flip', 'flipud', 'fliplr', 'hypot', 'bitwise_and', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', - 'equal', 'not_equal', 'greater', 'less', 'greater_equal', 'less_equal', 'rot90', 'einsum', + 'equal', 'not_equal', 'greater', 'less', 'greater_equal', 'less_equal', 'roll', 'rot90', 'einsum', 'true_divide', 'nonzero', 'quantile', 'percentile', 'shares_memory', 'may_share_memory', 'diff', 'ediff1d', 'resize', 'polyval', 'nan_to_num', 'isnan', 'isinf', 'isposinf', 'isneginf', 'isfinite', 'where', 'bincount', 'pad', 'cumsum'] @@ -6308,6 +6308,72 @@ def less_equal(x1, x2, out=None): _npi.greater_equal_scalar, out) +@set_module('mxnet.ndarray.numpy') +def roll(a, shift, axis=None): + """ + Roll array elements along a given axis. + + Elements that roll beyond the last position are re-introduced at + the first. + + Parameters + ---------- + a : ndarray + Input array. + shift : int or tuple of ints + The number of places by which elements are shifted. If a tuple, + then `axis` must be a tuple of the same size, and each of the + given axes is shifted by the corresponding number. If an int + while `axis` is a tuple of ints, then the same value is used for + all given axes. + axis : int or tuple of ints, optional + Axis or axes along which elements are shifted. By default, the + array is flattened before shifting, after which the original + shape is restored. + + Returns + ------- + res : ndarray + Output array, with the same shape as `a`. + + Notes + ----- + Supports rolling over multiple dimensions simultaneously. + + Examples + -------- + >>> x = np.arange(10) + >>> np.roll(x, 2) + array([8., 9., 0., 1., 2., 3., 4., 5., 6., 7.]) + >>> np.roll(x, -2) + array([2., 3., 4., 5., 6., 7., 8., 9., 0., 1.]) + + >>> x2 = np.reshape(x, (2,5)) + >>> x2 + array([[0., 1., 2., 3., 4.], + [5., 6., 7., 8., 9.]]) + >>> np.roll(x2, 1) + array([[9., 0., 1., 2., 3.], + [4., 5., 6., 7., 8.]]) + >>> np.roll(x2, -1) + array([[1., 2., 3., 4., 5.], + [6., 7., 8., 9., 0.]]) + >>> np.roll(x2, 1, axis=0) + array([[5., 6., 7., 8., 9.], + [0., 1., 2., 3., 4.]]) + >>> np.roll(x2, -1, axis=0) + array([[5., 6., 7., 8., 9.], + [0., 1., 2., 3., 4.]]) + >>> np.roll(x2, 1, axis=1) + array([[4., 0., 1., 2., 3.], + [9., 5., 6., 7., 8.]]) + >>> np.roll(x2, -1, axis=1) + array([[1., 2., 3., 4., 0.], + [6., 7., 8., 9., 5.]]) + """ + return _api_internal.roll(a, shift, axis) + + @set_module('mxnet.ndarray.numpy') def rot90(m, k=1, axes=(0, 1)): """ @@ -6351,7 +6417,7 @@ def rot90(m, k=1, axes=(0, 1)): [[5., 7.], [4., 6.]]]) """ - return _npi.rot90(m, k=k, axes=axes) + return _api_internal.rot90(m, k, axes) @set_module('mxnet.ndarray.numpy') diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 09712c03ea96..04476919dbd0 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -66,7 +66,7 @@ 'flip', 'flipud', 'fliplr', 'around', 'round', 'round_', 'arctan2', 'hypot', 'bitwise_and', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', 'equal', 'not_equal', - 'greater', 'less', 'greater_equal', 'less_equal', 'rot90', 'einsum', 'true_divide', 'nonzero', + 'greater', 'less', 'greater_equal', 'less_equal', 'roll', 'rot90', 'einsum', 'true_divide', 'nonzero', 'quantile', 'percentile', 'shares_memory', 'may_share_memory', 'diff', 'ediff1d', 'resize', 'matmul', 'nan_to_num', 'isnan', 'isinf', 'isposinf', 'isneginf', 'isfinite', 'polyval', 'where', 'bincount', 'pad', 'cumsum'] @@ -8170,6 +8170,72 @@ def less_equal(x1, x2, out=None): return _mx_nd_np.less_equal(x1, x2, out) +@set_module('mxnet.numpy') +def roll(a, shift, axis=None): + """ + Roll array elements along a given axis. + + Elements that roll beyond the last position are re-introduced at + the first. + + Parameters + ---------- + a : ndarray + Input array. + shift : int or tuple of ints + The number of places by which elements are shifted. If a tuple, + then `axis` must be a tuple of the same size, and each of the + given axes is shifted by the corresponding number. If an int + while `axis` is a tuple of ints, then the same value is used for + all given axes. + axis : int or tuple of ints, optional + Axis or axes along which elements are shifted. By default, the + array is flattened before shifting, after which the original + shape is restored. + + Returns + ------- + res : ndarray + Output array, with the same shape as `a`. + + Notes + ----- + Supports rolling over multiple dimensions simultaneously. + + Examples + -------- + >>> x = np.arange(10) + >>> np.roll(x, 2) + array([8., 9., 0., 1., 2., 3., 4., 5., 6., 7.]) + >>> np.roll(x, -2) + array([2., 3., 4., 5., 6., 7., 8., 9., 0., 1.]) + + >>> x2 = np.reshape(x, (2,5)) + >>> x2 + array([[0., 1., 2., 3., 4.], + [5., 6., 7., 8., 9.]]) + >>> np.roll(x2, 1) + array([[9., 0., 1., 2., 3.], + [4., 5., 6., 7., 8.]]) + >>> np.roll(x2, -1) + array([[1., 2., 3., 4., 5.], + [6., 7., 8., 9., 0.]]) + >>> np.roll(x2, 1, axis=0) + array([[5., 6., 7., 8., 9.], + [0., 1., 2., 3., 4.]]) + >>> np.roll(x2, -1, axis=0) + array([[5., 6., 7., 8., 9.], + [0., 1., 2., 3., 4.]]) + >>> np.roll(x2, 1, axis=1) + array([[4., 0., 1., 2., 3.], + [9., 5., 6., 7., 8.]]) + >>> np.roll(x2, -1, axis=1) + array([[1., 2., 3., 4., 0.], + [6., 7., 8., 9., 5.]]) + """ + return _mx_nd_np.roll(a, shift, axis=axis) + + @set_module('mxnet.numpy') def rot90(m, k=1, axes=(0, 1)): """ diff --git a/python/mxnet/symbol/numpy/_symbol.py b/python/mxnet/symbol/numpy/_symbol.py index 11e574ae5fd9..717049a6a8bf 100644 --- a/python/mxnet/symbol/numpy/_symbol.py +++ b/python/mxnet/symbol/numpy/_symbol.py @@ -48,7 +48,7 @@ 'diag_indices_from', 'hanning', 'hamming', 'blackman', 'flip', 'flipud', 'fliplr', 'hypot', 'bitwise_and', 'bitwise_xor', 'bitwise_or', 'rad2deg', 'deg2rad', 'unique', 'lcm', 'tril', 'identity', 'take', 'ldexp', 'vdot', 'inner', 'outer', - 'equal', 'not_equal', 'greater', 'less', 'greater_equal', 'less_equal', 'rot90', 'einsum', + 'equal', 'not_equal', 'greater', 'less', 'greater_equal', 'less_equal', 'roll', 'rot90', 'einsum', 'true_divide', 'quantile', 'percentile', 'shares_memory', 'may_share_memory', 'diff', 'ediff1d', 'resize', 'polyval', 'nan_to_num', 'isnan', 'isinf', 'isposinf', 'isneginf', 'isfinite', 'where', 'bincount', 'pad', 'cumsum'] @@ -5841,6 +5841,41 @@ def less_equal(x1, x2, out=None): _npi.greater_equal_scalar, out) +@set_module('mxnet.symbol.numpy') +def roll(a, shift, axis=None): + """ + Roll array elements along a given axis. + + Elements that roll beyond the last position are re-introduced at + the first. + + Parameters + ---------- + a : _Symbol + Input array. + shift : int or tuple of ints + The number of places by which elements are shifted. If a tuple, + then `axis` must be a tuple of the same size, and each of the + given axes is shifted by the corresponding number. If an int + while `axis` is a tuple of ints, then the same value is used for + all given axes. + axis : int or tuple of ints, optional + Axis or axes along which elements are shifted. By default, the + array is flattened before shifting, after which the original + shape is restored. + + Returns + ------- + res : _Symbol + Output array, with the same shape as `a`. + + Notes + ----- + Supports rolling over multiple dimensions simultaneously. + """ + return _npi.roll(a, shift, axis=axis) + + @set_module('mxnet.symbol.numpy') def rot90(m, k=1, axes=(0, 1)): """ diff --git a/src/api/operator/numpy/np_matrix_op.cc b/src/api/operator/numpy/np_matrix_op.cc index 080cca867ecb..05b5c23f0f7b 100644 --- a/src/api/operator/numpy/np_matrix_op.cc +++ b/src/api/operator/numpy/np_matrix_op.cc @@ -22,8 +22,10 @@ * \brief Implementation of the API of functions in src/operator/tensor/matrix_op.cc */ #include +#include #include "../utils.h" #include "../../../operator/tensor/matrix_op-inl.h" +#include "../../../operator/numpy/np_matrix_op-inl.h" namespace mxnet { @@ -46,4 +48,58 @@ MXNET_REGISTER_API("_npi.expand_dims") *ret = ndoutputs[0]; }); +MXNET_REGISTER_API("_npi.roll") +.set_body([](runtime::MXNetArgs args, runtime::MXNetRetValue* ret) { + using namespace runtime; + static const nnvm::Op* op = Op::Get("_npi_roll"); + nnvm::NodeAttrs attrs; + op::NumpyRollParam param; + if (args[1].type_code() == kNull) { + param.shift = dmlc::nullopt; + } else if (args[1].type_code() == kDLInt) { + param.shift = TShape(1, args[1].operator int64_t()); + } else { + param.shift = TShape(args[1].operator ObjectRef()); + } + if (args[2].type_code() == kNull) { + param.axis = dmlc::nullopt; + } else if (args[2].type_code() == kDLInt) { + param.axis = TShape(1, args[2].operator int64_t()); + } else { + param.axis = TShape(args[2].operator ObjectRef()); + } + attrs.parsed = std::move(param); + attrs.op = op; + SetAttrDict(&attrs); + NDArray* inputs[] = {args[0].operator mxnet::NDArray*()}; + int num_inputs = 1; + int num_outputs = 0; + auto ndoutputs = Invoke(op, &attrs, num_inputs, inputs, &num_outputs, nullptr); + *ret = ndoutputs[0]; +}); + +MXNET_REGISTER_API("_npi.rot90") +.set_body([](runtime::MXNetArgs args, runtime::MXNetRetValue* ret) { + using namespace runtime; + static const nnvm::Op* op = Op::Get("_npi_rot90"); + nnvm::NodeAttrs attrs; + op::NumpyRot90Param param; + param.k = args[1].operator int(); + if (args[2].type_code() == kNull) { + param.axes = dmlc::nullopt; + } else if (args[2].type_code() == kDLInt) { + param.axes = TShape(1, args[2].operator int64_t()); + } else { + param.axes = TShape(args[2].operator ObjectRef()); + } + attrs.parsed = std::move(param); + attrs.op = op; + SetAttrDict(&attrs); + NDArray* inputs[] = {args[0].operator mxnet::NDArray*()}; + int num_inputs = 1; + int num_outputs = 0; + auto ndoutputs = Invoke(op, &attrs, num_inputs, inputs, &num_outputs, nullptr); + *ret = ndoutputs[0]; +}); + } // namespace mxnet diff --git a/src/operator/numpy/np_matrix_op-inl.h b/src/operator/numpy/np_matrix_op-inl.h index 593a698dc93a..0bbe263cfc76 100644 --- a/src/operator/numpy/np_matrix_op-inl.h +++ b/src/operator/numpy/np_matrix_op-inl.h @@ -302,6 +302,13 @@ struct NumpyRollParam : public dmlc::Parameter { .describe("Axis or axes along which elements are shifted. By default, the array is flattened" "before shifting, after which the original shape is restored."); } + void SetAttrDict(std::unordered_map* dict) { + std::ostringstream shift_s, axis_s; + shift_s << shift; + axis_s << axis; + (*dict)["shift"] = shift_s.str(); + (*dict)["axis"] = axis_s.str(); + } }; template @@ -605,6 +612,13 @@ struct NumpyRot90Param : public dmlc::Parameter { .set_default(dmlc::optional()) .describe(" The array is rotated in the plane defined by the axes. Axes must be different."); } + void SetAttrDict(std::unordered_map* dict) { + std::ostringstream k_s, axes_s; + k_s << k; + axes_s << axes; + (*dict)["k"] = k_s.str(); + (*dict)["axes"] = axes_s.str(); + } }; struct rot90reverse { @@ -852,6 +866,10 @@ inline void HSplitOpForward(const nnvm::NodeAttrs &attrs, } else { real_axis = 0; } + if (param.sections > 0) { + CHECK_EQ(inputs[0].shape_[real_axis] % param.sections, 0U) + << "ValueError: array split does not result in an equal division"; + } SplitOpForwardImpl(attrs, ctx, inputs, req, outputs, real_axis); } diff --git a/src/operator/numpy/np_matrix_op.cc b/src/operator/numpy/np_matrix_op.cc index b6b8ff332db8..e9d269dd54d6 100644 --- a/src/operator/numpy/np_matrix_op.cc +++ b/src/operator/numpy/np_matrix_op.cc @@ -1155,7 +1155,7 @@ inline bool NumpyRollShape(const nnvm::NodeAttrs& attrs, return ElemwiseShape<1, 1>(attrs, in_attrs, out_attrs); } -NNVM_REGISTER_OP(_np_roll) +NNVM_REGISTER_OP(_npi_roll) .set_num_inputs(1) .set_num_outputs(1) .set_attr_parser(ParamParser) @@ -1180,7 +1180,7 @@ NNVM_REGISTER_OP(_np_roll) os1 << dmlc::optional(shifts); std::ostringstream os2; os2 << param.axis; - return MakeNonlossGradNode("_np_roll", n, ograds, {}, + return MakeNonlossGradNode("_npi_roll", n, ograds, {}, {{"shift", os1.str()}, {"axis", os2.str()}}); }) .set_attr("FResourceRequest", diff --git a/src/operator/numpy/np_matrix_op.cu b/src/operator/numpy/np_matrix_op.cu index 335f4fd984d6..c9e896bc5b57 100644 --- a/src/operator/numpy/np_matrix_op.cu +++ b/src/operator/numpy/np_matrix_op.cu @@ -71,7 +71,7 @@ NNVM_REGISTER_OP(_npi_column_stack) NNVM_REGISTER_OP(_backward_np_column_stack) .set_attr("FCompute", NumpyColumnStackBackward); -NNVM_REGISTER_OP(_np_roll) +NNVM_REGISTER_OP(_npi_roll) .set_attr("FCompute", NumpyRollCompute); template<>