ENH: add "axis" argument to generalized ufuncs.

It is only allowed for gufuncs with a single, shared core dimension.
numpy · mhvk · Apr 30, 2018 · Apr 30, 2018 · May 2, 2018 · May 27, 2018
commit 7a09080a87fa9d2a2b9b0f5ca91abe2701d4f971
diff --git a/doc/scipy-sphinx-theme b/doc/scipy-sphinx-theme
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
@@ -565,11 +565,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;
@@ -826,8 +827,21 @@ 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_RuntimeError,
+                                "cannot specify both 'axis' and 'axes'");
+                            goto fail;
+                        }
                         *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_RuntimeError,
+                                "cannot specify both 'axis' and 'axes'");
+                            goto fail;
+                        }
+                        *out_axis = value;
                         bad_arg = 0;
                     }
                     break;
@@ -1884,6 +1898,27 @@ _has_output_coredims(PyUFuncObject *ufunc) {
     return 0;
 }
 
+/*
+ * 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
@@ -1899,7 +1934,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),
@@ -1913,6 +1948,42 @@ _check_keepdims_support(PyUFuncObject *ufunc) {
     return 0;
 }
 
+/*
+ * 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
@@ -2239,8 +2310,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     NPY_ORDER order = NPY_KEEPORDER;
     /* Use the default assignment casting rule */
     NPY_CASTING casting = NPY_DEFAULT_ASSIGN_CASTING;
-    /* When provided, extobj, typetup, and axes contain borrowed references */
-    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) {
@@ -2265,10 +2336,17 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
 
     NPY_UF_DBG_PRINT("Getting arguments\n");
 
-    /* Get all the arguments */
+    /*
+     * Get all the arguments.
+     *
+     * Here, when provided, extobj, typetup, axes, and axis contain borrowed
+     * references; axes, however, may be set internally (from axis), and thus
+     * is XDECREF'd at the end. Hence, if passed in, we get our own reference.
+     */
     retval = get_ufunc_arguments(ufunc, args, kwds,
                 op, &order, &casting, &extobj,
-                &type_tup, &subok, NULL, &axes, &keepdims);
+                &type_tup, &subok, NULL, &axes, &axis, &keepdims);
+    Py_XINCREF(axes);
     if (retval < 0) {
         goto fail;
     }
@@ -2283,6 +2361,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.
      */
@@ -2350,6 +2434,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);
@@ -2709,6 +2806,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     Py_XDECREF(type_tup);
     Py_XDECREF(full_args.in);
     Py_XDECREF(full_args.out);
+    Py_XDECREF(axes);
 
     NPY_UF_DBG_PRINT("Returning Success\n");
 
@@ -2728,6 +2826,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     Py_XDECREF(type_tup);
     Py_XDECREF(full_args.in);
     Py_XDECREF(full_args.out);
+    Py_XDECREF(axes);
     PyArray_free(remap_axis_memory);
     PyArray_free(remap_axis);
     return retval;
@@ -2804,7 +2903,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;
     }

diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
@@ -718,6 +718,36 @@ def test_axes_argument(self):
         assert_raises(ValueError, mm, z[1], z, axes=[0, 1])
         assert_raises(ValueError, mm, z, z, out=z[0], axes=[0, 1])
 
+    def test_axis_argument(self):
+        # inner1d signature: '(i),(i)->()'
+        inner1d = umt.inner1d
+        a = np.arange(27.).reshape((3, 3, 3))
+        b = np.arange(10., 19.).reshape((3, 1, 3))
+        c = inner1d(a, b)
+        assert_array_equal(c, (a * b).sum(-1))
+        c = inner1d(a, b, axis=-1)
+        assert_array_equal(c, (a * b).sum(-1))
+        out = np.zeros_like(c)
+        d = inner1d(a, b, axis=-1, out=out)
+        assert_(d is out)
+        assert_array_equal(d, c)
+        c = inner1d(a, b, axis=0)
+        assert_array_equal(c, (a * b).sum(0))
+        # Sanity check on innerwt.
+        a = np.arange(6).reshape((2, 3))
+        b = np.arange(10, 16).reshape((2, 3))
+        w = np.arange(20, 26).reshape((2, 3))
+        assert_array_equal(umt.innerwt(a, b, w, axis=0),
+                           np.sum(a * b * w, axis=0))
+        # Check errors.
+        # Cannot pass in both axis and axes.
+        assert_raises(RuntimeError, inner1d, a, b, axis=0, axes=[0, 0])
+        # Not an integer.
+        assert_raises(TypeError, inner1d, a, b, axis=[0])
+        # more than 1 core dimensions.
+        mm = umt.matrix_multiply
+        assert_raises(TypeError, mm, a, b, axis=1)
+
     def test_keepdims_argument(self):
         # inner1d signature: '(i),(i)->()'
         inner1d = umt.inner1d
@@ -733,7 +763,15 @@ def test_keepdims_argument(self):
         d = inner1d(a, b, keepdims=True, out=out)
         assert_(d is out)
         assert_array_equal(d, c)
-        # Now combined with axes.
+        # Now combined with axis and axes.
+        c = inner1d(a, b, axis=-1, keepdims=False)
+        assert_array_equal(c, (a * b).sum(-1, keepdims=False))
+        c = inner1d(a, b, axis=-1, keepdims=True)
+        assert_array_equal(c, (a * b).sum(-1, keepdims=True))
+        c = inner1d(a, b, axis=0, keepdims=False)
+        assert_array_equal(c, (a * b).sum(0, keepdims=False))
+        c = inner1d(a, b, axis=0, keepdims=True)
+        assert_array_equal(c, (a * b).sum(0, keepdims=True))
         c = inner1d(a, b, axes=[(-1,), (-1,), ()], keepdims=False)
         assert_array_equal(c, (a * b).sum(-1))
         c = inner1d(a, b, axes=[(-1,), (-1,), (-1,)], keepdims=True)
@@ -777,10 +815,12 @@ def test_keepdims_argument(self):
         w = np.arange(20, 26).reshape((2, 3))
         assert_array_equal(umt.innerwt(a, b, w, keepdims=True),
                            np.sum(a * b * w, axis=-1, keepdims=True))
+        assert_array_equal(umt.innerwt(a, b, w, axis=0, keepdims=True),
+                           np.sum(a * b * w, axis=0, keepdims=True))
         # Check errors.
         # Not a boolean
         assert_raises(TypeError, inner1d, a, b, keepdims='true')
-        # 1 core dimension only.
+        # More than 1 core dimension, and core output dimensions.
         mm = umt.matrix_multiply
         assert_raises(TypeError, mm, a, b, keepdims=True)
         assert_raises(TypeError, mm, a, b, keepdims=False)