From 6568c6b022e12ab6d71e7548314009ced6ccabe9 Mon Sep 17 00:00:00 2001 From: Gagandeep Singh Date: Wed, 7 Jul 2021 22:00:03 +0530 Subject: ENH: Adding keepdims to np.argmin,np.argmax (#19211) * keepdims added to np.argmin,np.argmax * Added release notes entry * tested for axis=None,keepdims=True * Apply suggestions from code review * updated interface * updated interface * API changed, implementation to be done * Added reshape approach to C implementation * buggy implementation without reshape * TestArgMax, TestArgMin fixed, comments added * Fixed for matrix * removed unrequired changes * fixed CI failure * fixed linting issue * PyArray_ArgMaxKeepdims now only modifies shape and strides * Comments added to PyArray_ArgMaxKeepdims * Updated implementation of PyArray_ArgMinKeepdims to match with PyArray_ArgMaxKeepdims * Testing complete for PyArray_ArgMinKeepdims and PyArray_ArgMaxKeepdims * PyArray_ArgMinWithKeepdims both keepdims=True and keepdims=False * matched implementation of PyArray_ArgMaxKeepdims and PyArray_ArgMinKeepdims * simplified implementation * Added missing comment * removed unwanted header * addressed all the reviews * Removing unwanted changes * fixed docs * Added new lines * restored annotations * parametrized test * Apply suggestions from code review Co-authored-by: Sebastian Berg * keyword handling now done in np.argmin/np.argmax * corrected indendation * used with pytest.riases(ValueError) * fixed release notes * removed PyArray_ArgMaxWithKeepdims and PyArray_ArgMinWithKeepdims from public C-API * Apply suggestions from code review Co-authored-by: Eric Wieser * Apply suggestions from code review Co-authored-by: Eric Wieser Co-authored-by: Sebastian Berg Co-authored-by: Eric Wieser --- numpy/core/fromnumeric.py | 42 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 34 insertions(+), 8 deletions(-) (limited to 'numpy/core/fromnumeric.py') diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index ee93da901..7164a2c28 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -1114,12 +1114,12 @@ def argsort(a, axis=-1, kind=None, order=None): return _wrapfunc(a, 'argsort', axis=axis, kind=kind, order=order) -def _argmax_dispatcher(a, axis=None, out=None): +def _argmax_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): return (a, out) @array_function_dispatch(_argmax_dispatcher) -def argmax(a, axis=None, out=None): +def argmax(a, axis=None, out=None, *, keepdims=np._NoValue): """ Returns the indices of the maximum values along an axis. @@ -1133,12 +1133,18 @@ def argmax(a, axis=None, out=None): out : array, optional If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. Returns ------- index_array : ndarray of ints Array of indices into the array. It has the same shape as `a.shape` - with the dimension along `axis` removed. + with the dimension along `axis` removed. If `keepdims` is set to True, + then the size of `axis` will be 1 with the resulting array having same + shape as `a.shape`. See Also -------- @@ -1191,16 +1197,23 @@ def argmax(a, axis=None, out=None): >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) array([4, 3]) + Setting `keepdims` to `True`, + + >>> x = np.arange(24).reshape((2, 3, 4)) + >>> res = np.argmax(x, axis=1, keepdims=True) + >>> res.shape + (2, 1, 4) """ - return _wrapfunc(a, 'argmax', axis=axis, out=out) + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmax', axis=axis, out=out, **kwds) -def _argmin_dispatcher(a, axis=None, out=None): +def _argmin_dispatcher(a, axis=None, out=None, *, keepdims=np._NoValue): return (a, out) @array_function_dispatch(_argmin_dispatcher) -def argmin(a, axis=None, out=None): +def argmin(a, axis=None, out=None, *, keepdims=np._NoValue): """ Returns the indices of the minimum values along an axis. @@ -1214,12 +1227,18 @@ def argmin(a, axis=None, out=None): out : array, optional If provided, the result will be inserted into this array. It should be of the appropriate shape and dtype. + keepdims : bool, optional + If this is set to True, the axes which are reduced are left + in the result as dimensions with size one. With this option, + the result will broadcast correctly against the array. Returns ------- index_array : ndarray of ints Array of indices into the array. It has the same shape as `a.shape` - with the dimension along `axis` removed. + with the dimension along `axis` removed. If `keepdims` is set to True, + then the size of `axis` will be 1 with the resulting array having same + shape as `a.shape`. See Also -------- @@ -1272,8 +1291,15 @@ def argmin(a, axis=None, out=None): >>> np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1).squeeze(axis=-1) array([2, 0]) + Setting `keepdims` to `True`, + + >>> x = np.arange(24).reshape((2, 3, 4)) + >>> res = np.argmin(x, axis=1, keepdims=True) + >>> res.shape + (2, 1, 4) """ - return _wrapfunc(a, 'argmin', axis=axis, out=out) + kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {} + return _wrapfunc(a, 'argmin', axis=axis, out=out, **kwds) def _searchsorted_dispatcher(a, v, side=None, sorter=None): -- cgit v1.2.1