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 <sebastian@sipsolutions.net>

* 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 <wieser.eric@gmail.com>

* Apply suggestions from code review

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

1 files changed, 34 insertions, 8 deletions

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):