From c4d4cc5e5235329f5aaf9b17dc6c2a3a78e8346a Mon Sep 17 00:00:00 2001 From: barry-jin Date: Fri, 22 Oct 2021 12:10:49 -0700 Subject: [PATCH 1/5] [API] Add keepdims=False to argmax/argmin --- python/mxnet/ndarray/numpy/_op.py | 26 ++++++-- python/mxnet/numpy/multiarray.py | 34 +++++++--- src/operator/numpy/np_broadcast_reduce_op.cc | 6 +- src/operator/numpy/np_broadcast_reduce_op.h | 6 +- tests/python/unittest/test_numpy_op.py | 66 +++++++++++--------- 5 files changed, 95 insertions(+), 43 deletions(-) diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index cf1bd52c2228..e0b57e9f5399 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -5373,7 +5373,7 @@ def tril_indices(n, k=0, m=None): @set_module('mxnet.ndarray.numpy') -def argmax(a, axis=None, out=None): +def argmax(a, axis=None, out=None, keepdims=False): r""" Returns the indices of the maximum values along an axis. @@ -5388,6 +5388,11 @@ def argmax(a, axis=None, out=None): A location into which the result is stored. If provided, it must have the same shape and dtype as input ndarray. If not provided or `None`, a freshly-allocated array is returned. + keepdims : bool + If True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with + the input array. Otherwise, if False, the reduced axes (dimensions) must + not be included in the result. Default: False . Returns ------- @@ -5397,6 +5402,10 @@ def argmax(a, axis=None, out=None): Notes ----- + ``keepdims`` param is part of request `in data-api-standard + >> b array([2., 2.]) """ - return _api_internal.argmax(a, axis, False, out) + return _api_internal.argmax(a, axis, keepdims, out) @set_module('mxnet.ndarray.numpy') -def argmin(a, axis=None, out=None): +def argmin(a, axis=None, out=None, keepdims=False): r""" Returns the indices of the maximum values along an axis. @@ -5456,6 +5465,11 @@ def argmin(a, axis=None, out=None): out : ndarray or None, optional If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. + keepdims : bool + If True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with + the input array. Otherwise, if False, the reduced axes (dimensions) must + not be included in the result. Default: False . Returns ------- @@ -5465,6 +5479,10 @@ def argmin(a, axis=None, out=None): Notes ----- + ``keepdims`` param is part of request `in data-api-standard + >> b array([0., 0.]) """ - return _api_internal.argmin(a, axis, False, out) + return _api_internal.argmin(a, axis, keepdims, out) @set_module('mxnet.ndarray.numpy') diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index a58f1faf5587..d112322bd8ec 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -1568,10 +1568,10 @@ def copyto(self, other): def asscalar(self): raise AttributeError('mxnet.numpy.ndarray object has no attribute asscalar') - def argmax(self, axis=None, out=None): # pylint: disable=arguments-differ + def argmax(self, axis=None, out=None, keepdims=False): # pylint: disable=arguments-differ """Return indices of the maximum values along the given axis. Refer to `mxnet.numpy.argmax` for full documentation.""" - return argmax(self, axis, out) + return argmax(self, axis, out, keepdims) def as_in_context(self, context): """This function has been deprecated. Please refer to ``ndarray.as_in_ctx``.""" @@ -1909,10 +1909,10 @@ def argmax_channel(self, *args, **kwargs): """ raise AttributeError('mxnet.numpy.ndarray object has no attribute argmax_channel') - def argmin(self, axis=None, out=None): # pylint: disable=arguments-differ + def argmin(self, axis=None, out=None, keepdims=False): # pylint: disable=arguments-differ """Return indices of the minium values along the given axis. Refer to `mxnet.numpy.argmin` for full documentation.""" - return argmin(self, axis, out) + return argmin(self, axis, out, keepdims) def clip(self, min=None, max=None, out=None): # pylint: disable=arguments-differ """Return an array whose values are limited to [min, max]. @@ -7758,7 +7758,7 @@ def clip(a, a_min, a_max, out=None): @set_module('mxnet.numpy') -def argmax(a, axis=None, out=None): +def argmax(a, axis=None, out=None, keepdims=False): r""" Returns the indices of the maximum values along an axis. @@ -7772,6 +7772,11 @@ def argmax(a, axis=None, out=None): out : ndarray or None, optional If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. + keepdims : bool + If True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with + the input array. Otherwise, if False, the reduced axes (dimensions) must + not be included in the result. Default: False . Returns ------- @@ -7780,6 +7785,10 @@ def argmax(a, axis=None, out=None): with the dimension along `axis` removed. .. note:: + ``keepdims`` param is part of request `in data-api-standard + >> b array([2., 2.]) """ - return _mx_nd_np.argmax(a, axis, out) + return _mx_nd_np.argmax(a, axis, out, keepdims) @set_module('mxnet.numpy') -def argmin(a, axis=None, out=None): +def argmin(a, axis=None, out=None, keepdims=False): r""" Returns the indices of the minimum values along an axis. @@ -7841,6 +7850,11 @@ def argmin(a, axis=None, out=None): out : ndarray or None, optional If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. + keepdims : bool + If True, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with + the input array. Otherwise, if False, the reduced axes (dimensions) must + not be included in the result. Default: False . Returns ------- @@ -7849,6 +7863,10 @@ def argmin(a, axis=None, out=None): with the dimension along `axis` removed. .. note:: + ``keepdims`` param is part of request `in data-api-standard + >> b array([0., 0.]) """ - return _mx_nd_np.argmin(a, axis, out) + return _mx_nd_np.argmin(a, axis, out, keepdims) @set_module('mxnet.numpy') diff --git a/src/operator/numpy/np_broadcast_reduce_op.cc b/src/operator/numpy/np_broadcast_reduce_op.cc index 4ddfe0516e4f..56d0132ca16a 100644 --- a/src/operator/numpy/np_broadcast_reduce_op.cc +++ b/src/operator/numpy/np_broadcast_reduce_op.cc @@ -66,7 +66,11 @@ void NumpyArgMinMaxRTCCompute::operator()(const nnvm::NodeAttrs& attrs, axes = dmlc::optional>(t); } TShape small; - small = NumpyReduceAxesShapeImpl(in.shape_, axes, true); + if (param.keepdims) { + small = outputs[0].shape_; + } else { + small = NumpyReduceAxesShapeImpl(in.shape_, axes, true); + } mxnet::TShape src_shape, dst_shape; BroadcastReduceShapeCompact(in.shape_, small, &src_shape, &dst_shape); const TBlob in_data = in.reshape(src_shape); diff --git a/src/operator/numpy/np_broadcast_reduce_op.h b/src/operator/numpy/np_broadcast_reduce_op.h index 822ca1ffef27..21c8957ac816 100644 --- a/src/operator/numpy/np_broadcast_reduce_op.h +++ b/src/operator/numpy/np_broadcast_reduce_op.h @@ -550,7 +550,11 @@ void NumpyArgMinMaxCompute(const nnvm::NodeAttrs& attrs, axes = dmlc::optional>(t); } TShape small; - small = NumpyReduceAxesShapeImpl(in.shape_, axes, true); + if (param.keepdims) { + small = outputs[0].shape_; + } else { + small = NumpyReduceAxesShapeImpl(in.shape_, axes, true); + } mxnet::TShape src_shape, dst_shape; BroadcastReduceShapeCompact(in.shape_, small, &src_shape, &dst_shape); const TBlob in_data = in.reshape(src_shape); diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 1010475c605d..6b34f2fe97a6 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -4460,50 +4460,58 @@ def test_np_argmin_argmax(): ops = ['argmin', 'argmax'] class TestArgExtreme(HybridBlock): - def __init__(self, op_name, axis=None): + def __init__(self, op_name, axis=None, keepdims=False): super(TestArgExtreme, self).__init__() self._op_name = op_name self._axis = axis + self.keepdims = keepdims def forward(self, x): - return getattr(x, self._op_name)(self._axis) + return getattr(x, self._op_name)(self._axis, keepdims=self.keepdims) for op_name in ops: - for shape, axis, throw_exception in workloads: - for dtype in dtypes: - a = np.random.uniform(low=0, high=100, size=shape).astype(dtype) - if throw_exception: - # Cannot use assert_exception because sometimes the main thread - # proceeds to `assert False` before the exception is thrown - # in the worker thread. Have to use mx.nd.waitall() here - # to block the main thread. - try: - getattr(np, op_name)(a, axis) - mx.nd.waitall() - assert False - except mx.MXNetError: - pass - else: - mx_ret = getattr(np, op_name)(a, axis=axis) - np_ret = getattr(onp, op_name)(a.asnumpy(), axis=axis) - assert mx_ret.dtype == np_ret.dtype - assert same(mx_ret.asnumpy(), np_ret) - - for hybridize in [False, True]: - net = TestArgExtreme(op_name, axis) - if hybridize: - net.hybridize() + for keepdims in (True, False): + for shape, axis, throw_exception in workloads: + for dtype in dtypes: + a = np.random.uniform(low=0, high=100, size=shape).astype(dtype) if throw_exception: + # Cannot use assert_exception because sometimes the main thread + # proceeds to `assert False` before the exception is thrown + # in the worker thread. Have to use mx.nd.waitall() here + # to block the main thread. try: - net(a) + getattr(np, op_name)(a, axis) mx.nd.waitall() assert False except mx.MXNetError: pass else: - mx_ret = net(a) + mx_ret = getattr(np, op_name)(a, axis=axis, keepdims=keepdims) + np_ret = getattr(onp, op_name)(a.asnumpy(), axis=axis) assert mx_ret.dtype == np_ret.dtype - assert same(mx_ret.asnumpy(), np_ret) + if keepdims: + assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) + else: + assert same(mx_ret.asnumpy(), np_ret) + + for hybridize in [False, True]: + net = TestArgExtreme(op_name, axis, keepdims) + if hybridize: + net.hybridize() + if throw_exception: + try: + net(a) + mx.nd.waitall() + assert False + except mx.MXNetError: + pass + else: + mx_ret = net(a) + assert mx_ret.dtype == np_ret.dtype + if keepdims: + assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) + else: + assert same(mx_ret.asnumpy(), np_ret) @use_np From b09930ae667a70a61cedf3eefdda0627c1708401 Mon Sep 17 00:00:00 2001 From: barry-jin Date: Sun, 24 Oct 2021 21:25:00 -0700 Subject: [PATCH 2/5] fix docstring --- python/mxnet/ndarray/numpy/_op.py | 4 ++-- python/mxnet/numpy/multiarray.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index e0b57e9f5399..5315c5c2159c 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -5402,7 +5402,7 @@ def argmax(a, axis=None, out=None, keepdims=False): Notes ----- - ``keepdims`` param is part of request `in data-api-standard + ``keepdims`` param is part of request in data-api-standard Date: Mon, 25 Oct 2021 09:29:00 -0700 Subject: [PATCH 3/5] fix link --- python/mxnet/ndarray/numpy/_op.py | 4 ++-- python/mxnet/numpy/multiarray.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index 5315c5c2159c..d4cb72313724 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -5403,7 +5403,7 @@ def argmax(a, axis=None, out=None, keepdims=False): Notes ----- ``keepdims`` param is part of request in data-api-standard - `_, which is not the parameter in official NumPy In case of multiple occurrences of the maximum values, the indices @@ -5480,7 +5480,7 @@ def argmin(a, axis=None, out=None, keepdims=False): Notes ----- ``keepdims`` param is part of request in data-api-standard - `_, which is not the parameter in official NumPy In case of multiple occurrences of the maximum values, the indices diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 72df4fa3a971..72e70e7a4b8d 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -7786,7 +7786,7 @@ def argmax(a, axis=None, out=None, keepdims=False): .. note:: ``keepdims`` param is part of request in data-api-standard - `_, which is not the parameter in official NumPy In case of multiple occurrences of the maximum values, the indices @@ -7864,7 +7864,7 @@ def argmin(a, axis=None, out=None, keepdims=False): .. note:: ``keepdims`` param is part of request in data-api-standard - `_, which is not the parameter in official NumPy In case of multiple occurrences of the minimum values, the indices From fcf0083eb4b564f7773a07b5d8f237857f060519 Mon Sep 17 00:00:00 2001 From: barry-jin Date: Thu, 28 Oct 2021 14:55:23 -0700 Subject: [PATCH 4/5] Parametrize argmin/argmax test --- tests/python/unittest/test_numpy_op.py | 131 +++++++++++++------------ 1 file changed, 67 insertions(+), 64 deletions(-) diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 6b34f2fe97a6..71490488e544 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -4437,28 +4437,28 @@ def GetDimSize(shp, axis): @use_np -def test_np_argmin_argmax(): - workloads = [ - ((), 0, False), - ((), -1, False), - ((), 1, True), - ((5, 3), None, False), - ((5, 3), -1, False), - ((5, 3), 1, False), - ((5, 3), 3, True), - ((5, 0, 3), 0, False), - ((5, 0, 3), -1, False), - ((5, 0, 3), None, True), - ((5, 0, 3), 1, True), - ((3, 5, 7), None, False), - ((3, 5, 7), 0, False), - ((3, 5, 7), 1, False), - ((3, 5, 7), 2, False), - ((3, 5, 7, 9, 11), -3, False), - ] - dtypes = ['float16', 'float32', 'float64', 'bool', 'int32'] - ops = ['argmin', 'argmax'] - +@pytest.mark.parametrize('shape,axis,throw_exception', [ + ((), 0, False), + ((), -1, False), + ((), 1, True), + ((5, 3), None, False), + ((5, 3), -1, False), + ((5, 3), 1, False), + ((5, 3), 3, True), + ((5, 0, 3), 0, False), + ((5, 0, 3), -1, False), + ((5, 0, 3), None, True), + ((5, 0, 3), 1, True), + ((3, 5, 7), None, False), + ((3, 5, 7), 0, False), + ((3, 5, 7), 1, False), + ((3, 5, 7), 2, False), + ((3, 5, 7, 9, 11), -3, False), +]) +@pytest.mark.parametrize('dtype', ['float16', 'float32', 'float64', 'bool', 'int32']) +@pytest.mark.parametrize('op_name', ['argmin', 'argmax']) +@pytest.mark.parametrize('keepdims', [True, False]) +def test_np_argmin_argmax(shape, axis, throw_exception, dtype, op_name, keepdims): class TestArgExtreme(HybridBlock): def __init__(self, op_name, axis=None, keepdims=False): super(TestArgExtreme, self).__init__() @@ -4469,49 +4469,52 @@ def __init__(self, op_name, axis=None, keepdims=False): def forward(self, x): return getattr(x, self._op_name)(self._axis, keepdims=self.keepdims) - for op_name in ops: - for keepdims in (True, False): - for shape, axis, throw_exception in workloads: - for dtype in dtypes: - a = np.random.uniform(low=0, high=100, size=shape).astype(dtype) - if throw_exception: - # Cannot use assert_exception because sometimes the main thread - # proceeds to `assert False` before the exception is thrown - # in the worker thread. Have to use mx.nd.waitall() here - # to block the main thread. - try: - getattr(np, op_name)(a, axis) - mx.nd.waitall() - assert False - except mx.MXNetError: - pass - else: - mx_ret = getattr(np, op_name)(a, axis=axis, keepdims=keepdims) - np_ret = getattr(onp, op_name)(a.asnumpy(), axis=axis) - assert mx_ret.dtype == np_ret.dtype - if keepdims: - assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) - else: - assert same(mx_ret.asnumpy(), np_ret) + a = np.random.uniform(low=0, high=100, size=shape).astype(dtype) + if throw_exception: + # Cannot use assert_exception because sometimes the main thread + # proceeds to `assert False` before the exception is thrown + # in the worker thread. Have to use mx.nd.waitall() here + # to block the main thread. + try: + getattr(np, op_name)(a, axis) + mx.nd.waitall() + assert False + except mx.MXNetError: + pass + else: + mx_ret = getattr(np, op_name)(a, axis=axis, keepdims=keepdims) + np_ret = getattr(onp, op_name)(a.asnumpy(), axis=axis) + assert mx_ret.dtype == np_ret.dtype + if keepdims: + # if shape == (): + # assert mx_ret.shape == shape + # # elif not axis: + # # assert mx_ret.size == 0 + # else: + # print(mx_ret.shape, np_ret.shape, shape, axis) + assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) + # assert same(mx_ret.asnumpy(), np_ret.reshape((mx_ret.shape))) + else: + assert same(mx_ret.asnumpy(), np_ret) - for hybridize in [False, True]: - net = TestArgExtreme(op_name, axis, keepdims) - if hybridize: - net.hybridize() - if throw_exception: - try: - net(a) - mx.nd.waitall() - assert False - except mx.MXNetError: - pass - else: - mx_ret = net(a) - assert mx_ret.dtype == np_ret.dtype - if keepdims: - assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) - else: - assert same(mx_ret.asnumpy(), np_ret) + for hybridize in [False, True]: + net = TestArgExtreme(op_name, axis, keepdims) + if hybridize: + net.hybridize() + if throw_exception: + try: + net(a) + mx.nd.waitall() + assert False + except mx.MXNetError: + pass + else: + mx_ret = net(a) + assert mx_ret.dtype == np_ret.dtype + if keepdims: + assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) + else: + assert same(mx_ret.asnumpy(), np_ret) @use_np From 02b2ce09a623413bb420b4fa527eceae8fcf5052 Mon Sep 17 00:00:00 2001 From: barry-jin Date: Fri, 29 Oct 2021 11:08:53 -0700 Subject: [PATCH 5/5] update test_np_argmin_argmax --- tests/python/unittest/test_numpy_op.py | 51 +++++++++----------------- 1 file changed, 17 insertions(+), 34 deletions(-) diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 71490488e544..86145873f1e7 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -4458,7 +4458,8 @@ def GetDimSize(shp, axis): @pytest.mark.parametrize('dtype', ['float16', 'float32', 'float64', 'bool', 'int32']) @pytest.mark.parametrize('op_name', ['argmin', 'argmax']) @pytest.mark.parametrize('keepdims', [True, False]) -def test_np_argmin_argmax(shape, axis, throw_exception, dtype, op_name, keepdims): +@pytest.mark.parametrize('hybridize', [True, False]) +def test_np_argmin_argmax(shape, axis, throw_exception, dtype, op_name, keepdims, hybridize): class TestArgExtreme(HybridBlock): def __init__(self, op_name, axis=None, keepdims=False): super(TestArgExtreme, self).__init__() @@ -4471,50 +4472,32 @@ def forward(self, x): a = np.random.uniform(low=0, high=100, size=shape).astype(dtype) if throw_exception: - # Cannot use assert_exception because sometimes the main thread - # proceeds to `assert False` before the exception is thrown - # in the worker thread. Have to use mx.nd.waitall() here - # to block the main thread. - try: + with pytest.raises(MXNetError): getattr(np, op_name)(a, axis) - mx.nd.waitall() - assert False - except mx.MXNetError: - pass + mx.npx.waitall() else: mx_ret = getattr(np, op_name)(a, axis=axis, keepdims=keepdims) np_ret = getattr(onp, op_name)(a.asnumpy(), axis=axis) assert mx_ret.dtype == np_ret.dtype if keepdims: - # if shape == (): - # assert mx_ret.shape == shape - # # elif not axis: - # # assert mx_ret.size == 0 - # else: - # print(mx_ret.shape, np_ret.shape, shape, axis) assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) - # assert same(mx_ret.asnumpy(), np_ret.reshape((mx_ret.shape))) else: assert same(mx_ret.asnumpy(), np_ret) - for hybridize in [False, True]: - net = TestArgExtreme(op_name, axis, keepdims) - if hybridize: - net.hybridize() - if throw_exception: - try: - net(a) - mx.nd.waitall() - assert False - except mx.MXNetError: - pass + net = TestArgExtreme(op_name, axis, keepdims) + if hybridize: + net.hybridize() + if throw_exception: + with pytest.raises(MXNetError): + getattr(np, op_name)(a, axis) + mx.npx.waitall() + else: + mx_ret = net(a) + assert mx_ret.dtype == np_ret.dtype + if keepdims: + assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) else: - mx_ret = net(a) - assert mx_ret.dtype == np_ret.dtype - if keepdims: - assert same(np.squeeze(mx_ret, axis=axis).asnumpy(), np_ret) - else: - assert same(mx_ret.asnumpy(), np_ret) + assert same(mx_ret.asnumpy(), np_ret) @use_np