diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index a129f738c566..0570aed74db1 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -27,7 +27,7 @@ from ..ndarray import NDArray __all__ = ['zeros', 'ones', 'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'tensordot', - 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack'] + 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'ldexp'] @set_module('mxnet.ndarray.numpy') @@ -732,3 +732,38 @@ def get_list(arrays): arrays = get_list(arrays) return _npi.stack(*arrays, axis=axis, out=out) + + +@set_module('mxnet.ndarray.numpy') +def ldexp(x1, x2, out=None): + """ + ldexp(x1, x2, out=None) + Returns x1 * 2**x2, element-wise. + The mantissas `x1` and twos exponents `x2` are used to construct + floating point numbers ``x1 * 2**x2``. + Parameters + ---------- + x1 : ndarray or scalar + Array of multipliers. + x2 : ndarray or scalar, int + Array of twos exponents. + out : ndarray, optional + A location into which the result is stored. If provided, it must have + a shape that the inputs broadcast to. If not, a freshly-allocated array is returned. + Returns + ------- + y : ndarray or scalar + The result of ``x1 * 2**x2``. + This is a scalar if both `x1` and `x2` are scalars. + Notes + ----- + Complex dtypes are not supported, they will raise a TypeError. + Different from numpy, we allow x2 to be float besides int. + `ldexp` is useful as the inverse of `frexp`, if used by itself it is + more clear to simply use the expression ``x1 * 2**x2``. + Examples + -------- + >>> np.ldexp(5, np.arange(4)) + array([ 5., 10., 20., 40.]) + """ + return _ufunc_helper(x1, x2, _npi.ldexp, _np.ldexp, _npi.ldexp_scalar, _npi.rldexp_scalar, out) diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 316d88ae0a81..e98b6428a675 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -45,7 +45,7 @@ __all__ = ['ndarray', 'empty', 'array', 'zeros', 'ones', 'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', - 'concatenate', 'stack'] + 'concatenate', 'stack', 'ldexp'] # This function is copied from ndarray.py since pylint @@ -1898,3 +1898,38 @@ def stack(arrays, axis=0, out=None): stacked : ndarray The stacked array has one more dimension than the input arrays.""" return _mx_nd_np.stack(arrays, axis=axis, out=out) + + +@set_module('mxnet.numpy') +def ldexp(x1, x2, out=None): + """ + ldexp(x1, x2, out=None) + Returns x1 * 2**x2, element-wise. + The mantissas `x1` and twos exponents `x2` are used to construct + floating point numbers ``x1 * 2**x2``. + Parameters + ---------- + x1 : ndarray or scalar + Array of multipliers. + x2 : ndarray or scalar, int + Array of twos exponents. + out : ndarray, optional + A location into which the result is stored. If provided, it must have + a shape that the inputs broadcast to. If not, a freshly-allocated array is returned. + Returns + ------- + y : ndarray or scalar + The result of ``x1 * 2**x2``. + This is a scalar if both `x1` and `x2` are scalars. + Notes + ----- + Complex dtypes are not supported, they will raise a TypeError. + Different from numpy, we allow x2 to be float besides int. + `ldexp` is useful as the inverse of `frexp`, if used by itself it is + more clear to simply use the expression ``x1 * 2**x2``. + Examples + -------- + >>> np.ldexp(5, np.arange(4)) + array([ 5., 10., 20., 40.]) + """ + return _mx_nd_np.ldexp(x1, x2, out) diff --git a/python/mxnet/symbol/numpy/_symbol.py b/python/mxnet/symbol/numpy/_symbol.py index 5537e6305b83..60c6aa39d97a 100644 --- a/python/mxnet/symbol/numpy/_symbol.py +++ b/python/mxnet/symbol/numpy/_symbol.py @@ -30,7 +30,7 @@ from . import _internal as _npi __all__ = ['zeros', 'ones', 'add', 'subtract', 'multiply', 'divide', 'mod', 'power', 'tensordot', - 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack'] + 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'ldexp'] def _num_outputs(sym): @@ -1362,4 +1362,33 @@ def get_list(arrays): return _npi.stack(*arrays, axis=axis, out=out) +@set_module('mxnet.symbol.numpy') +def ldexp(x1, x2, out=None): + """ + ldexp(x1, x2, out=None) + Returns x1 * 2**x2, element-wise. + The mantissas `x1` and twos exponents `x2` are used to construct + floating point numbers ``x1 * 2**x2``. + Parameters + ---------- + x1 : _Symbol + Array of multipliers. + x2 : _Symbol + Array of twos exponents. + out : _Symbol or None + Dummy parameter to keep the consistency with the ndarray counterpart. + Returns + ------- + y : _Symbol + The result of ``x1 * 2**x2``. + Notes + ----- + Complex dtypes are not supported, they will raise a TypeError. + Different from numpy, we allow x2 to be float besides int. + `ldexp` is useful as the inverse of `frexp`, if used by itself it is + more clear to simply use the expression ``x1 * 2**x2``. + """ + return _ufunc_helper(x1, x2, _npi.ldexp, _np.ldexp, _npi.ldexp_scalar, _npi.rldexp_scalar, out) + + _set_np_symbol_class(_Symbol) diff --git a/src/operator/mshadow_op.h b/src/operator/mshadow_op.h index ab53e7733066..f7cae55eb264 100644 --- a/src/operator/mshadow_op.h +++ b/src/operator/mshadow_op.h @@ -357,6 +357,17 @@ MXNET_UNARY_MATH_OP(reciprocal_cube_root, 1.0f / math::cbrt(a)); MXNET_UNARY_MATH_OP(reciprocal_cube_root_grad, -1.0f / (3.0f * math::cbrt(a) * math::id(a))); +/*! \brief used for generate element of ldexp */ +MXNET_BINARY_MATH_OP(ldexp, math::id(a) * math::pow(2.0f, b)); + +MXNET_BINARY_MATH_OP(ldexp_grad, math::pow(2.0f, b)); + +MXNET_BINARY_MATH_OP(ldexp_rgrad, math::id(a) * math::pow(2.0f, b) * math::log(2.0f)); + +MXNET_BINARY_MATH_OP(rldexp, math::id(b) * math::pow(2.0f, a)); // swap a and b if a is scalar. + +MXNET_BINARY_MATH_OP(rldexp_grad, math::id(b) * math::pow(2.0f, a) * math::log(2.0f)); + /*! \brief used for generate element of round */ MXNET_SIMPLE_UNARY_MATH_OP(round); diff --git a/src/operator/numpy/np_elemwise_broadcast_op.cc b/src/operator/numpy/np_elemwise_broadcast_op.cc index c36423dff9fd..70ea9f75747d 100644 --- a/src/operator/numpy/np_elemwise_broadcast_op.cc +++ b/src/operator/numpy/np_elemwise_broadcast_op.cc @@ -182,5 +182,79 @@ MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rpower_scalar) .set_attr("FCompute", BinaryScalarOp::Compute) .set_attr("FGradient", ElemwiseGradUseOut{"_backward_rpower_scalar"}); +MXNET_OPERATOR_REGISTER_BINARY_BROADCAST(_npi_ldexp) +.describe(R"code( + ldexp(x1, x2, out=None) + + Returns x1 * 2**x2, element-wise. + + The mantissas `x1` and twos exponents `x2` are used to construct + floating point numbers ``x1 * 2**x2``. + + Parameters + ---------- + x1 : ndarray or scalar + Array of multipliers. + x2 : ndarray or scalar, int + Array of twos exponents. + out : ndarray, optional + A location into which the result is stored. If provided, it must have + a shape that the inputs broadcast to. If not, a freshly-allocated array is returned. + + Returns + ------- + y : ndarray or scalar + The result of ``x1 * 2**x2``. + This is a scalar if both `x1` and `x2` are scalars. + + Notes + ----- + Complex dtypes are not supported, they will raise a TypeError. + Different from numpy, we allow x2 to be float besides int. + + `ldexp` is useful as the inverse of `frexp`, if used by itself it is + more clear to simply use the expression ``x1 * 2**x2``. + + Examples + -------- + >>> np.ldexp(5, np.arange(4)) + array([ 5., 10., 20., 40.]) +)code" ADD_FILELINE) +.set_attr("FCompute", BinaryBroadcastCompute) +.set_attr("FGradient", ElemwiseGradUseIn{"_backward_npi_ldexp"}); + +MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_ldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Compute) +.set_attr("FGradient", ElemwiseGradUseIn{"_backward_npi_ldexp_scalar"}); + +MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Compute) +.set_attr("FGradient", ElemwiseGradUseIn{"_backward_npi_rldexp_scalar"}); + +NNVM_REGISTER_OP(_backward_npi_ldexp) +.set_num_inputs(3) +.set_num_outputs(2) +.set_attr("TIsBackward", true) +.set_attr("FInplaceOption", + [](const NodeAttrs& attrs){ + return std::vector >{{0, 1}}; + }) +.set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) +.set_attr("FCompute", BinaryBroadcastBackwardUseIn); + +MXNET_OPERATOR_REGISTER_BINARY(_backward_npi_ldexp_scalar) +.add_argument("scalar", "float", "scalar value") +.set_attr_parser([](NodeAttrs *attrs) { attrs->parsed = std::stod(attrs->dict["scalar"]); }) +.set_attr("FCompute", BinaryScalarOp::Backward); + +MXNET_OPERATOR_REGISTER_BINARY(_backward_npi_rldexp_scalar) +.add_argument("scalar", "float", "scalar value") +.set_attr_parser([](NodeAttrs *attrs) { attrs->parsed = std::stod(attrs->dict["scalar"]); }) +.set_attr("FCompute", BinaryScalarOp::Backward); + } // namespace op } // namespace mxnet diff --git a/src/operator/numpy/np_elemwise_broadcast_op.cu b/src/operator/numpy/np_elemwise_broadcast_op.cu index c858b3a4987a..525a12f28561 100644 --- a/src/operator/numpy/np_elemwise_broadcast_op.cu +++ b/src/operator/numpy/np_elemwise_broadcast_op.cu @@ -78,5 +78,24 @@ NNVM_REGISTER_OP(_npi_maximum_scalar) NNVM_REGISTER_OP(_npi_minimum_scalar) .set_attr("FCompute", BinaryScalarOp::Compute); +NNVM_REGISTER_OP(_npi_ldexp) +.set_attr("FCompute", BinaryBroadcastCompute); + +NNVM_REGISTER_OP(_npi_ldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Compute); + +NNVM_REGISTER_OP(_npi_rldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Compute); + +NNVM_REGISTER_OP(_backward_npi_ldexp) +.set_attr("FCompute", BinaryBroadcastBackwardUseIn); + +NNVM_REGISTER_OP(_backward_npi_ldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Backward); + +NNVM_REGISTER_OP(_backward_npi_rldexp_scalar) +.set_attr("FCompute", BinaryScalarOp::Backward); + } // namespace op } // namespace mxnet diff --git a/src/operator/operator_tune.cc b/src/operator/operator_tune.cc index 98ce14e7bf05..9776203ef2cb 100644 --- a/src/operator/operator_tune.cc +++ b/src/operator/operator_tune.cc @@ -362,6 +362,11 @@ IMPLEMENT_BINARY_WORKLOAD_BWD(mxnet::op::mshadow_op::smooth_l1_gradient); // NO IMPLEMENT_BLANK_WORKLOAD_FWD(mxnet::op::mxnet_op::set_to_int<0>); // NOLINT() IMPLEMENT_BLANK_WORKLOAD_FWD(mxnet::op::mxnet_op::set_to_int<1>); // NOLINT() IMPLEMENT_BLANK_WORKLOAD_FWD(mxnet::op::PopulateFullIdxRspKernel); // NOLINT() +IMPLEMENT_BINARY_WORKLOAD_FWD(mxnet::op::mshadow_op::ldexp); // NOLINT() +IMPLEMENT_BINARY_WORKLOAD_FWD(mxnet::op::mshadow_op::rldexp); // NOLINT() +IMPLEMENT_BINARY_WORKLOAD_BWD(mxnet::op::mshadow_op::ldexp_grad); // NOLINT() +IMPLEMENT_BINARY_WORKLOAD_BWD(mxnet::op::mshadow_op::ldexp_rgrad); // NOLINT() +IMPLEMENT_BINARY_WORKLOAD_BWD(mxnet::op::mshadow_op::rldexp_grad); // NOLINT() /*! * \brief Tuner objects, *not* automatically generated */ diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 6a1a6a6dfd5d..91d346509d95 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -213,6 +213,67 @@ def test_np_dot(): assert False +@with_seed() +@use_np +def test_np_ldexp(): + class TestLdexp(HybridBlock): + def __init__(self): + super(TestLdexp, self).__init__() + + def hybrid_forward(self, F, x1, x2): + return F.np.ldexp(x1, x2) + + def _np_ldexp(x1, x2): + return x1 * _np.power(2.0, x2) + + def dldx(x1, x2): + grad_a = _np.power(2.0, x2) + grad_b = _np_ldexp(x1, x2) * _np.log(2.0) + if len(x1) == 1: + grad_a = _np.sum(grad_a) + if len(x2) == 1: + grad_b = _np.sum(grad_b) + return [grad_a, grad_b] + + shapes = [ + ((3, 1), (3, 1)), + ((3, 1, 2), (3, 1, 2)), + ((1, ),(1, )), + ((1, ), (2, )), + ((3, ), (1, )), + ((3, 0), (3, 0)), # zero-size shape + ((0, 1), (0, 1)), # zero-size shape + ((2, 0, 2), (2, 0, 2)), # zero-size shape + ] + + for hybridize in [True, False]: + for shape1, shape2 in shapes: + for dtype in [_np.float16, _np.float32, _np.float64]: + test_ldexp = TestLdexp() + if hybridize: + test_ldexp.hybridize() + x1 = rand_ndarray(shape = shape1, dtype = dtype).as_np_ndarray() + x1.attach_grad() + x2 = rand_ndarray(shape = shape2, dtype = dtype).as_np_ndarray() + x2.attach_grad() + + np_out = _np_ldexp(x1.asnumpy(), x2.asnumpy()) + with mx.autograd.record(): + mx_out = test_ldexp(x1, x2) + assert mx_out.shape == np_out.shape + assert_almost_equal(mx_out.asnumpy(), np_out, rtol = 1e-1, atol = 1e-1) + + mx_out.backward() + np_backward = dldx(x1.asnumpy(), x2.asnumpy()) + assert_almost_equal(x1.grad.asnumpy(), np_backward[0], atol=1e-1, rtol=1e-1) + assert_almost_equal(x2.grad.asnumpy(), np_backward[1], atol=1e-1, rtol=1e-1) + + # Test imperative once again + mx_out = np.ldexp(x1, x2) + np_out = _np_ldexp(x1.asnumpy(), x2.asnumpy()) + assert_almost_equal(mx_out.asnumpy(), np_out, rtol=1e-1, atol=1e-1) + + @with_seed() @use_np def test_np_sum():