API: Require reduce promoters to start with None to match

But additionally, we add an (implicit) fallback promoter for
the reduction case that fills in the first dtype with the second
one.

Author: seberg

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;
         }