ENH: add "axis" argument to generalized ufuncs.

It is only allowed for gufuncs with a single, shared
core dimension.

1 files changed, 116 insertions, 11 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 631999bc1..1091cdf6e 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -566,11 +566,12 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
                     NPY_ORDER *out_order,
                     NPY_CASTING *out_casting,
                     PyObject **out_extobj,
-                    PyObject **out_typetup,
-                    int *out_subok,
-                    PyArrayObject **out_wheremask,
-                    PyObject **out_axes,
-                    int *out_keepdims)
+                    PyObject **out_typetup,  /* type: Tuple[np.dtype] */
+                    int *out_subok,  /* bool */
+                    PyArrayObject **out_wheremask, /* PyArray of bool */
+                    PyObject **out_axes,  /* type: List[Tuple[T]] */
+                    PyObject **out_axis,  /* type: T */
+                    int *out_keepdims)  /* bool */
 {
     int i, nargs;
     int nin = ufunc->nin;
@@ -588,6 +589,9 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
     if (out_axes != NULL) {
         *out_axes = NULL;
     }
+    if (out_axis != NULL) {
+        *out_axis = NULL;
+    }
     if (out_wheremask != NULL) {
         *out_wheremask = NULL;
     }
@@ -827,9 +831,23 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
                 case 'a':
                     /* possible axes argument for generalized ufunc */
                     if (out_axes != NULL && strcmp(str, "axes") == 0) {
+                        if (out_axis != NULL && *out_axis != NULL) {
+                            PyErr_SetString(PyExc_TypeError,
+                                "cannot specify both 'axis' and 'axes'");
+                            goto fail;
+                        }
                         Py_INCREF(value);
                         *out_axes = value;
-
+                        bad_arg = 0;
+                    }
+                    else if (out_axis != NULL && strcmp(str, "axis") == 0) {
+                        if (out_axes != NULL && *out_axes != NULL) {
+                            PyErr_SetString(PyExc_TypeError,
+                                "cannot specify both 'axis' and 'axes'");
+                            goto fail;
+                        }
+                        Py_INCREF(value);
+                        *out_axis = value;
                         bad_arg = 0;
                     }
                     break;
@@ -1045,6 +1063,10 @@ fail:
         Py_XDECREF(*out_axes);
         *out_axes = NULL;
     }
+    if (out_axis != NULL) {
+        Py_XDECREF(*out_axis);
+        *out_axis = NULL;
+    }
     return -1;
 }
 
@@ -1891,6 +1913,27 @@ _has_output_coredims(PyUFuncObject *ufunc) {
 }
 
 /*
+ * Check whether the gufunc can be used with axis, i.e., that there is only
+ * a single, shared core dimension (which means that operands either have
+ * that dimension, or have no core dimensions).  Returns 0 if all is fine,
+ * and sets an error and returns -1 if not.
+ */
+static int
+_check_axis_support(PyUFuncObject *ufunc) {
+    if (ufunc->core_num_dim_ix != 1) {
+        PyErr_Format(PyExc_TypeError,
+                     "%s: axis can only be used with a single shared core "
+                     "dimension, not with the %d distinct ones implied by "
+                     "signature %s.",
+                     ufunc_get_name_cstr(ufunc),
+                     ufunc->core_num_dim_ix,
+                     ufunc->core_signature);
+        return -1;
+    }
+    return 0;
+}
+
+/*
  * Check whether the gufunc can be used with keepdims, i.e., that all its
  * input arguments have the same number of core dimension, and all output
  * arguments have no core dimensions. Returns 0 if all is fine, and sets
@@ -1905,7 +1948,7 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
         if (ufunc->core_num_dims[i] != (i < nin ? input_core_dims : 0)) {
             PyErr_Format(PyExc_TypeError,
                 "%s does not support keepdims: its signature %s requires "
-                "that %s %d has %d core dimensions, but keepdims can only "
+                "%s %d to have %d core dimensions, but keepdims can only "
                 "be used when all inputs have the same number of core "
                 "dimensions and all outputs have no core dimensions.",
                 ufunc_get_name_cstr(ufunc),
@@ -1920,6 +1963,42 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
 }
 
 /*
+ * Translate axis to axes list of the form [(axis,), ...], with an
+ * empty tuple for operands without core dimensions.
+ * Returns an axes tuple or NULL on failure.
+ */
+static PyObject*
+_build_axes_tuple_from_axis(PyObject *axis, int core_num_dims[], int nop) {
+    int i;
+    PyObject *axes = NULL, *axis_tuple = NULL, *tuple;
+    PyObject *empty_tuple = PyTuple_New(0);  /* cannot realistically fail */
+
+    axes = PyList_New(nop);
+    if (axes == NULL) {
+        return NULL;
+    }
+    axis_tuple = PyTuple_Pack(1, axis);
+    if (axis_tuple == NULL) {
+        goto fail;
+    }
+
+    for (i = 0; i < nop; i++) {
+        tuple = core_num_dims[i] == 1 ? axis_tuple : empty_tuple;
+        Py_INCREF(tuple);
+        PyList_SET_ITEM(axes, i, tuple);
+    }
+    Py_DECREF(axis_tuple);
+    Py_DECREF(empty_tuple);
+    return axes;
+
+fail:
+    Py_XDECREF(axis_tuple);
+    Py_XDECREF(axes);
+    Py_DECREF(empty_tuple);
+    return NULL;
+}
+
+/*
  * Interpret a possible axes keyword argument, using it to fill the remap_axis
  * array which maps default to actual axes for each operand, indexed as
  * as remap_axis[iop][iaxis]. The default axis order has first all broadcast
@@ -2245,7 +2324,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     NPY_ORDER order = NPY_KEEPORDER;
     /* Use the default assignment casting rule */
     NPY_CASTING casting = NPY_DEFAULT_ASSIGN_CASTING;
-    PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL;
+    /* other possible keyword arguments */
+    PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL, *axis = NULL;
     int keepdims = -1;
 
     if (ufunc == NULL) {
@@ -2270,10 +2350,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
 
     NPY_UF_DBG_PRINT("Getting arguments\n");
 
-    /* Get all the arguments */
+    /*
+     * Get all the arguments.
+     */
     retval = get_ufunc_arguments(ufunc, args, kwds,
                 op, &order, &casting, &extobj,
-                &type_tup, &subok, NULL, &axes, &keepdims);
+                &type_tup, &subok, NULL, &axes, &axis, &keepdims);
     if (retval < 0) {
         goto fail;
     }
@@ -2288,6 +2370,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
             goto fail;
         }
     }
+    if (axis != NULL) {
+        retval = _check_axis_support(ufunc);
+        if (retval < 0) {
+            goto fail;
+        }
+    }
     /*
      * If keepdims is set and true, signal all dimensions will be the same.
      */
@@ -2355,6 +2443,19 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         goto fail;
     }
 
+    /*
+     * Translate axis to axes list of the form [(axis,), ...], with an
+     * empty tuple for operands without core dimensions.
+     */
+    if (axis) {
+        assert(axes == NULL);  /* parser prevents passing in both axis & axes */
+        axes = _build_axes_tuple_from_axis(axis, core_num_dims, nop);
+        if (axes == NULL) {
+            retval = -1;
+            goto fail;
+        }
+    }
+
     /* Possibly remap axes. */
     if (axes) {
         remap_axis = PyArray_malloc(sizeof(remap_axis[0]) * nop);
@@ -2714,8 +2815,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     Py_XDECREF(type_tup);
     Py_XDECREF(extobj);
     Py_XDECREF(axes);
+    Py_XDECREF(axis);
     Py_XDECREF(full_args.in);
     Py_XDECREF(full_args.out);
+    Py_XDECREF(axes);
 
     NPY_UF_DBG_PRINT("Returning Success\n");
 
@@ -2735,8 +2838,10 @@ fail:
     Py_XDECREF(type_tup);
     Py_XDECREF(extobj);
     Py_XDECREF(axes);
+    Py_XDECREF(axis);
     Py_XDECREF(full_args.in);
     Py_XDECREF(full_args.out);
+    Py_XDECREF(axes);
     PyArray_free(remap_axis_memory);
     PyArray_free(remap_axis);
     return retval;
@@ -2812,7 +2917,7 @@ PyUFunc_GenericFunction(PyUFuncObject *ufunc,
     /* Get all the arguments */
     retval = get_ufunc_arguments(ufunc, args, kwds,
                 op, &order, &casting, &extobj,
-                &type_tup, &subok, &wheremask, NULL, NULL);
+                &type_tup, &subok, &wheremask, NULL, NULL, NULL);
     if (retval < 0) {
         goto fail;
     }