ENH: Adding keepdims to np.argmin,np.argmax

doc/release/upcoming_changes/19211.new_feature.rst

@@ -0,0 +1,9 @@
+`keepdims` optional argument added to `numpy.argmin`, `numpy.argmax`
+--------------------------------------------------------------------
+
+`keepdims` argument has been added to `numpy.argmin`, `numpy.argmax`.
+By default, it is `False` and hence no behaviour change should be expected
+in existing codes using `numpy.argmin`, `numpy.argmax`. If set to `True`,
+the axes which are reduced are left in the result as dimensions with size one.
+In simple words, the resulting array will be having same dimensions
+as the input array.

doc/source/reference/c-api/array.rst

@@ -2019,6 +2019,13 @@ Calculation
     Equivalent to :meth:`ndarray.argmax<numpy.ndarray.argmax>` (*self*, *axis*). Return the index of
     the largest element of *self* along *axis*.
 
+.. c:function:: PyObject* PyArray_ArgMaxKeepdims( \
+        PyArrayObject* self, int axis, PyArrayObject* out)
+
+    Equivalent to :meth:`ndarray.argmax<numpy.ndarray.argmax>` (*self*, *axis*). Return the index of
+    the largest element of *self* along *axis*. The number of dimensions of the result matches with the
+    that of the input object.
+
 .. c:function:: PyObject* PyArray_ArgMin( \
         PyArrayObject* self, int axis, PyArrayObject* out)
 

numpy/__init__.cython-30.pxd

@@ -644,6 +644,7 @@ cdef extern from "numpy/arrayobject.h":
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
     object PyArray_ArgMax (ndarray, int, ndarray)
+    object PyArray_ArgMaxKeepdims (ndarray, int, ndarray)
     object PyArray_ArgMin (ndarray, int, ndarray)
     object PyArray_Reshape (ndarray, object)
     object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER)

numpy/__init__.pxd

@@ -602,6 +602,7 @@ cdef extern from "numpy/arrayobject.h":
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
     object PyArray_ArgMax (ndarray, int, ndarray)
+    object PyArray_ArgMaxKeepdims (ndarray, int, ndarray)
     object PyArray_ArgMin (ndarray, int, ndarray)
     object PyArray_Reshape (ndarray, object)
     object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER)

numpy/__init__.pyi

@@ -1299,37 +1299,43 @@ class _ArrayOrScalarCommon:
         self,
         axis: None = ...,
         out: None = ...,
+        keepdims: L[False] = ...,
     ) -> intp: ...
     @overload
     def argmax(
         self,
         axis: _ShapeLike = ...,
         out: None = ...,
+        keepdims: bool = ...,
     ) -> Any: ...
     @overload
     def argmax(
         self,
         axis: Optional[_ShapeLike] = ...,
         out: _NdArraySubClass = ...,
+        keepdims: bool = ...,
     ) -> _NdArraySubClass: ...
 
     @overload
     def argmin(
         self,
         axis: None = ...,
         out: None = ...,
+        keepdims: L[False] = ...,
     ) -> intp: ...
     @overload
     def argmin(
         self,
         axis: _ShapeLike = ...,
-         out: None = ...,
+        out: None = ...,
+        keepdims: bool = ...,
     ) -> Any: ...
     @overload
     def argmin(
         self,
         axis: Optional[_ShapeLike] = ...,
         out: _NdArraySubClass = ...,
+        keepdims: bool = ...,
     ) -> _NdArraySubClass: ...
 
     def argsort(

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):
-    return (a, out)
+def _argmax_dispatcher(a, axis=None, out=None, keepdims=None):
+    return (a, out, keepdims)
 
 
 @array_function_dispatch(_argmax_dispatcher)
-def argmax(a, axis=None, out=None):
+def argmax(a, axis=None, out=None, keepdims=False):
     """
     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)
+    return _wrapfunc(a, 'argmax', axis=axis, out=out, 
+                        keepdims=keepdims)
 
 
-def _argmin_dispatcher(a, axis=None, out=None):
-    return (a, out)
+def _argmin_dispatcher(a, axis=None, out=None, keepdims=None):
+    return (a, out, keepdims)
 
 
 @array_function_dispatch(_argmin_dispatcher)
-def argmin(a, axis=None, out=None):
+def argmin(a, axis=None, out=None, keepdims=False):
     """
     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,24 @@ 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)
+    res = _wrapfunc(a, 'argmin', axis=axis, out=out)
+
+    if keepdims:
+        if axis is None:
+            new_shape = (1,)*a.ndim
+        else:
+            new_shape = list(a.shape)
+            new_shape[axis] = 1
+        return res.reshape(new_shape)
+
+    return res
 
 
 def _searchsorted_dispatcher(a, v, side=None, sorter=None):

numpy/core/fromnumeric.pyi

@@ -130,25 +130,29 @@ def argmax(
     a: ArrayLike,
     axis: None = ...,
     out: Optional[ndarray] = ...,
+    keepdims: Literal[False] = ...,
 ) -> intp: ...
 @overload
 def argmax(
     a: ArrayLike,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
+    keepdims: bool = ...,
 ) -> Any: ...
 
 @overload
 def argmin(
     a: ArrayLike,
     axis: None = ...,
     out: Optional[ndarray] = ...,
+    keepdims: Literal[False] = ...,
 ) -> intp: ...
 @overload
 def argmin(
     a: ArrayLike,
     axis: Optional[int] = ...,
     out: Optional[ndarray] = ...,
+    keepdims: bool = ...,
 ) -> Any: ...
 
 @overload

numpy/core/src/multiarray/calculation.c

@@ -17,6 +17,8 @@
 #include "calculation.h"
 #include "array_assign.h"
 
+#include "alloc.h"
+
 static double
 power_of_ten(int n)
 {
@@ -150,6 +152,41 @@ PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out)
     return NULL;
 }
 
+NPY_NO_EXPORT PyObject*
+PyArray_ArgMaxKeepdims(PyArrayObject *op, int axis, PyArrayObject* out) {
+    PyArrayObject* ret = (PyArrayObject*)PyArray_ArgMax(op, axis, NULL);
+    PyArray_Dims newdims;
+    newdims.len = PyArray_NDIM(op);
+    newdims.ptr = npy_alloc_cache_dim(newdims.len);
+    npy_intp* currdim = PyArray_DIMS(op);
+    for( int dim = 0; dim < newdims.len; dim++ ) {
+        if( dim == axis || axis == NPY_MAXDIMS ) {
+            newdims.ptr[dim] = 1;
+        } else {
+            newdims.ptr[dim] = currdim[dim];
+        }
+    }
+
+    // If ret is same as out, then should we reshape or raise an error?
+    ret = (PyArrayObject*)PyArray_Newshape(ret, &newdims, NPY_CORDER);
+    if (out != NULL) {
+        if ((PyArray_NDIM(out) != PyArray_NDIM(ret)) ||
+                !PyArray_CompareLists(PyArray_DIMS(out), PyArray_DIMS(ret),
+                                      PyArray_NDIM(out))) {
+            PyErr_SetString(PyExc_ValueError,
+                    "output array dimensions do not match with that of result from np.argmax.");
+            Py_XDECREF(ret);
+            return NULL;
+        }
+        PyArray_CopyInto(out, ret);
+        Py_DECREF(ret);
+        ret = out;
+        Py_INCREF(ret);
+    }
+    npy_free_cache_dim_obj(newdims);
+    return (PyObject*)ret;
+}
+
 /*NUMPY_API
  * ArgMin
  */

numpy/core/src/multiarray/calculation.h

@@ -4,6 +4,9 @@
 NPY_NO_EXPORT PyObject*
 PyArray_ArgMax(PyArrayObject* self, int axis, PyArrayObject *out);
 
+NPY_NO_EXPORT PyObject*
+PyArray_ArgMaxKeepdims(PyArrayObject* self, int axis, PyArrayObject *out);
+
 NPY_NO_EXPORT PyObject*
 PyArray_ArgMin(PyArrayObject* self, int axis, PyArrayObject *out);
 

numpy/core/src/multiarray/methods.c

@@ -284,16 +284,23 @@ array_argmax(PyArrayObject *self,
 {
     int axis = NPY_MAXDIMS;
     PyArrayObject *out = NULL;
+    npy_bool keepdims = NPY_FALSE;
     NPY_PREPARE_ARGPARSER;
 
     if (npy_parse_arguments("argmax", args, len_args, kwnames,
             "|axis", &PyArray_AxisConverter, &axis,
             "|out", &PyArray_OutputConverter, &out,
+            "|keepdims", &PyArray_BoolConverter, &keepdims,
             NULL, NULL, NULL) < 0) {
         return NULL;
     }
 
-    PyObject *ret = PyArray_ArgMax(self, axis, out);
+    PyObject *ret;
+    if( keepdims ) {
+        ret = PyArray_ArgMaxKeepdims(self, axis, out);
+    } else {
+        ret = PyArray_ArgMax(self, axis, out);
+    }
 
     /* this matches the unpacking behavior of ufuncs */
     if (out == NULL) {

numpy/core/src/umath/ufunc_object.c

@@ -4001,6 +4001,7 @@ static PyObject *
 PyUFunc_GenericReduction(PyUFuncObject *ufunc,
         PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames, int operation)
 {
+
     int i, naxes=0, ndim;
     int axes[NPY_MAXDIMS];
 

numpy/core/tests/test_multiarray.py

@@ -4333,7 +4333,29 @@ def test_object_argmax_with_NULLs(self):
         assert_equal(a.argmax(), 3)
         a[1] = 30
         assert_equal(a.argmax(), 1)
+
+    def test_np_argmax_keepdims(self):
 
+        sizes = [(3,), (3, 2), (2, 3),
+                 (3, 3), (2, 3, 4), (4, 3, 2)]
+        for size in sizes:
+            arr = np.random.normal(size=size)
+            for axis in range(len(size)):
+                res = np.argmax(arr, axis=axis, keepdims=True)
+                assert_(res.ndim == arr.ndim)
+                assert_(res.shape[axis] == 1)
+                outarray = np.empty(res.shape, dtype=res.dtype)
+                np.argmax(arr, axis=axis, out=outarray, 
+                          keepdims=True)
+                assert_equal(res, outarray)
+
+            # Testing for axis=None, keepdims=True
+            res = np.argmax(arr, axis=None, keepdims=True)
+            assert_(res.ndim == arr.ndim)
+            assert_(res.shape == (1,)*arr.ndim)
+            outarray = np.empty(res.shape, dtype=res.dtype)
+            np.argmax(arr, axis=None, out=outarray, keepdims=True)
+            assert_equal(res, outarray)
 
 class TestArgmin:
 
@@ -4489,7 +4511,23 @@ def test_object_argmin_with_NULLs(self):
         assert_equal(a.argmin(), 3)
         a[1] = 10
         assert_equal(a.argmin(), 1)
-
+
+    # Avoid until np.argmin is ready to support keepdims
+    # def test_np_argmin_keepdims(self):
+
+    #     sizes = [(3,), (3, 2), (2, 3),
+    #              (3, 3), (2, 3, 4), (4, 3, 2)]
+    #     for size in sizes:
+    #         arr = np.random.normal(size=size)
+    #         for axis in range(len(size)):
+    #             res = np.argmin(arr, axis=axis, keepdims=True)
+    #             assert_(res.ndim == arr.ndim)
+    #             assert_(res.shape[axis] == 1)
+
+    #         # Testing for axis=None, keepdims=True
+    #         res = np.argmin(arr, axis=None, keepdims=True)
+    #         assert_(res.ndim == arr.ndim)
+    #         assert_(res.shape == (1,)*arr.ndim)
 
 class TestMinMax: