Merge pull request numpy#26090 from seberg/minimal-reduce-fix

API: Require reduce promoters to start with None to match

Author: ngoldbaum

numpy/_core/src/umath/dispatching.c

@@ -274,21 +274,20 @@ resolve_implementation_info(PyUFuncObject *ufunc,
                     /* Unspecified out always matches (see below for inputs) */
                     continue;
                 }
+                assert(i == 0);
                 /*
-                 * This is a reduce-like operation, which always have the form
-                 * `(res_DType, op_DType, res_DType)`.  If the first and last
-                 * dtype of the loops match, this should be reduce-compatible.
+                 * This is a reduce-like operation, we enforce that these
+                 * register with None as the first DType.  If a reduction
+                 * uses the same DType, we will do that promotion.
+                 * A `(res_DType, op_DType, res_DType)` pattern can make sense
+                 * in other context as well and could be confusing.
                  */
-                if (PyTuple_GET_ITEM(curr_dtypes, 0)
-                        == PyTuple_GET_ITEM(curr_dtypes, 2)) {
+                if (PyTuple_GET_ITEM(curr_dtypes, 0) == Py_None) {
                     continue;
                 }
-                /*
-                 * This should be a reduce, but doesn't follow the reduce
-                 * pattern.  So (for now?) consider this not a match.
-                 */
+                /* Otherwise, this is not considered a match */
                 matches = NPY_FALSE;
-                continue;
+                break;
             }
 
             if (resolver_dtype == (PyArray_DTypeMeta *)Py_None) {
@@ -488,7 +487,7 @@ resolve_implementation_info(PyUFuncObject *ufunc,
  * those defined by the `signature` unmodified).
  */
 static PyObject *
-call_promoter_and_recurse(PyUFuncObject *ufunc, PyObject *promoter,
+call_promoter_and_recurse(PyUFuncObject *ufunc, PyObject *info,
         PyArray_DTypeMeta *op_dtypes[], PyArray_DTypeMeta *signature[],
         PyArrayObject *const operands[])
 {
@@ -498,37 +497,51 @@ call_promoter_and_recurse(PyUFuncObject *ufunc, PyObject *promoter,
     int promoter_result;
     PyArray_DTypeMeta *new_op_dtypes[NPY_MAXARGS];
 
-    if (PyCapsule_CheckExact(promoter)) {
-        /* We could also go the other way and wrap up the python function... */
-        PyArrayMethod_PromoterFunction *promoter_function = PyCapsule_GetPointer(
-                promoter, "numpy._ufunc_promoter");
-        if (promoter_function == NULL) {
+    if (info != NULL) {
+        PyObject *promoter = PyTuple_GET_ITEM(info, 1);
+        if (PyCapsule_CheckExact(promoter)) {
+            /* We could also go the other way and wrap up the python function... */
+            PyArrayMethod_PromoterFunction *promoter_function = PyCapsule_GetPointer(
+                    promoter, "numpy._ufunc_promoter");
+            if (promoter_function == NULL) {
+                return NULL;
+            }
+            promoter_result = promoter_function((PyObject *)ufunc,
+                    op_dtypes, signature, new_op_dtypes);
+        }
+        else {
+            PyErr_SetString(PyExc_NotImplementedError,
+                    "Calling python functions for promotion is not implemented.");
             return NULL;
         }
-        promoter_result = promoter_function((PyObject *)ufunc,
-                op_dtypes, signature, new_op_dtypes);
-    }
-    else {
-        PyErr_SetString(PyExc_NotImplementedError,
-                "Calling python functions for promotion is not implemented.");
-        return NULL;
-    }
-    if (promoter_result < 0) {
-        return NULL;
-    }
-    /*
-     * If none of the dtypes changes, we would recurse infinitely, abort.
-     * (Of course it is nevertheless possible to recurse infinitely.)
-     */
-    int dtypes_changed = 0;
-    for (int i = 0; i < nargs; i++) {
-        if (new_op_dtypes[i] != op_dtypes[i]) {
-            dtypes_changed = 1;
-            break;
+        if (promoter_result < 0) {
+            return NULL;
+        }
+        /*
+        * If none of the dtypes changes, we would recurse infinitely, abort.
+        * (Of course it is nevertheless possible to recurse infinitely.)
+        *
+        * TODO: We could allow users to signal this directly and also move
+        *       the call to be (almost immediate).  That would call it
+        *       unnecessarily sometimes, but may allow additional flexibility.
+        */
+        int dtypes_changed = 0;
+        for (int i = 0; i < nargs; i++) {
+            if (new_op_dtypes[i] != op_dtypes[i]) {
+                dtypes_changed = 1;
+                break;
+            }
+        }
+        if (!dtypes_changed) {
+            goto finish;
         }
     }
-    if (!dtypes_changed) {
-        goto finish;
+    else {
+        /* Reduction special path */
+        new_op_dtypes[0] = NPY_DT_NewRef(op_dtypes[1]);
+        new_op_dtypes[1] = NPY_DT_NewRef(op_dtypes[1]);
+        Py_XINCREF(op_dtypes[2]);
+        new_op_dtypes[2] = op_dtypes[2];
     }
 
     /*
@@ -788,13 +801,13 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
 
     /*
      * At this point `info` is NULL if there is no matching loop, or it is
-     * a promoter that needs to be used/called:
+     * a promoter that needs to be used/called.
+     * TODO: It may be nice to find a better reduce-solution, but this way
+     *       it is a True fallback (not registered so lowest priority)
      */
-    if (info != NULL) {
-        PyObject *promoter = PyTuple_GET_ITEM(info, 1);
-
+    if (info != NULL || op_dtypes[0] == NULL) {
         info = call_promoter_and_recurse(ufunc,
-                promoter, op_dtypes, signature, ops);
+                info, op_dtypes, signature, ops);
         if (info == NULL && PyErr_Occurred()) {
             return NULL;
         }

numpy/_core/src/umath/reduction.c

@@ -21,6 +21,7 @@
 #include "array_coercion.h"
 #include "array_method.h"
 #include "ctors.h"
+#include "refcount.h"
 
 #include "numpy/ufuncobject.h"
 #include "lowlevel_strided_loops.h"
@@ -438,7 +439,7 @@ PyUFunc_ReduceWrapper(PyArrayMethod_Context *context,
     Py_INCREF(result);
 
     if (initial_buf != NULL && PyDataType_REFCHK(PyArray_DESCR(result))) {
-        PyArray_Item_XDECREF(initial_buf, PyArray_DESCR(result));
+        PyArray_ClearBuffer(PyArray_DESCR(result), initial_buf, 0, 1, 1);
     }
     PyMem_FREE(initial_buf);
     NPY_AUXDATA_FREE(auxdata);
@@ -450,7 +451,7 @@ PyUFunc_ReduceWrapper(PyArrayMethod_Context *context,
 
 fail:
     if (initial_buf != NULL && PyDataType_REFCHK(op_dtypes[0])) {
-        PyArray_Item_XDECREF(initial_buf, op_dtypes[0]);
+        PyArray_ClearBuffer(op_dtypes[0], initial_buf, 0, 1, 1);
     }
     PyMem_FREE(initial_buf);
     NPY_AUXDATA_FREE(auxdata);

numpy/_core/tests/test_stringdtype.py

@@ -746,6 +746,25 @@ def test_add_promoter(string_list):
         assert_array_equal(arr + op, rresult)
 
 
+def test_add_promoter_reduce():
+    # Exact TypeError could change, but ensure StringDtype doesn't match
+    with pytest.raises(TypeError, match="the resolved dtypes are not"):
+        np.add.reduce(np.array(["a", "b"], dtype="U"))
+
+    # On the other hand, using `dtype=T` in the *ufunc* should work.
+    np.add.reduce(np.array(["a", "b"], dtype="U"), dtype=np.dtypes.StringDType)
+
+
+def test_multiply_reduce():
+    # At the time of writing (NumPy 2.0) this is very limited (and rather
+    # ridiculous anyway).  But it works and actually makes some sense...
+    # (NumPy does not allow non-scalar initial values)
+    repeats = np.array([2, 3, 4])
+    val = "school-🚌"
+    res = np.multiply.reduce(repeats, initial=val, dtype=np.dtypes.StringDType)
+    assert res == val * np.prod(repeats)
+
+
 @pytest.mark.parametrize("use_out", [True, False])
 @pytest.mark.parametrize("other", [2, [2, 1, 3, 4, 1, 3]])
 @pytest.mark.parametrize(