Added reshape approach to C implementation

numpy · seberg · Jul 7, 2021 · Jun 10, 2021 · Jun 10, 2021 · Jun 10, 2021
commit 09433d6202f0151c382ebc9d3e997ed5a721afcd
diff --git a/doc/source/reference/c-api/array.rst b/doc/source/reference/c-api/array.rst
@@ -2014,12 +2014,18 @@ Calculation
 
 
 .. c:function:: PyObject* PyArray_ArgMax( \
-        PyArrayObject* self, int axis, PyArrayObject* out, \
-        int keepdims)
+        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*.
 
+.. 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)
 

diff --git a/numpy/__init__.cython-30.pxd b/numpy/__init__.cython-30.pxd
@@ -643,7 +643,8 @@ cdef extern from "numpy/arrayobject.h":
     int PyArray_Sort (ndarray, int, NPY_SORTKIND)
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
-    object PyArray_ArgMax (ndarray, int, ndarray, int)
+    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)

diff --git a/numpy/__init__.pxd b/numpy/__init__.pxd
@@ -601,7 +601,8 @@ cdef extern from "numpy/arrayobject.h":
     int PyArray_Sort (ndarray, int, NPY_SORTKIND)
     object PyArray_ArgSort (ndarray, int, NPY_SORTKIND)
     object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *)
-    object PyArray_ArgMax (ndarray, int, ndarray, int)
+    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)

diff --git a/numpy/core/src/multiarray/calculation.c b/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)
 {
@@ -38,8 +40,7 @@ power_of_ten(int n)
  * ArgMax
  */
 NPY_NO_EXPORT PyObject *
-PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out,
-               int keepdims)
+PyArray_ArgMax(PyArrayObject *op, int axis, PyArrayObject *out)
 {
     PyArrayObject *ap = NULL, *rp = NULL;
     PyArray_ArgFunc* arg_func;
@@ -151,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
  */

diff --git a/numpy/core/src/multiarray/calculation.h b/numpy/core/src/multiarray/calculation.h
@@ -2,8 +2,10 @@
 #define _NPY_CALCULATION_H_
 
 NPY_NO_EXPORT PyObject*
-PyArray_ArgMax(PyArrayObject* self, int axis, PyArrayObject *out, 
-               int keepdims);
+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);

diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
@@ -284,7 +284,7 @@ array_argmax(PyArrayObject *self,
 {
     int axis = NPY_MAXDIMS;
     PyArrayObject *out = NULL;
-    int keepdims = -1;
+    npy_bool keepdims = NPY_FALSE;
     NPY_PREPARE_ARGPARSER;
 
     if (npy_parse_arguments("argmax", args, len_args, kwnames,
@@ -295,7 +295,12 @@ array_argmax(PyArrayObject *self,
         return NULL;
     }
 
-    PyObject *ret = PyArray_ArgMax(self, axis, out, keepdims);
+    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) {

diff --git a/numpy/core/src/umath/ufunc_object.c b/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];
 

diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -4344,11 +4344,18 @@ def test_np_argmax_keepdims(self):
                 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.argmin(arr, 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:
 
@@ -4505,21 +4512,22 @@ def test_object_argmin_with_NULLs(self):
         a[1] = 10
         assert_equal(a.argmin(), 1)
 
-    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)
87A2
+    # 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)
+    #         # 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: