Merge pull request #8819 from mhvk/gufunc-axis-argument

ENH: Implement axes keyword argument for gufuncs.

3 files changed, 335 insertions, 29 deletions

diff --git a/doc/source/reference/ufuncs.rst b/doc/source/reference/ufuncs.rst
index 3711f660f..59d25a9ca 100644
--- a/doc/source/reference/ufuncs.rst
+++ b/doc/source/reference/ufuncs.rst
@@ -336,6 +336,20 @@ advanced usage and will not typically be used.
     of False indicate to leave the value in the output alone. This argument
     cannot be used for generalized ufuncs as those take non-scalar input.
 
+*axes*
+
+    .. versionadded:: 1.15
+
+    A list of tuples with indices of axes a generalized ufunc should operate
+    on. For instance, for a signature of ``(i,j),(j,k)->(i,k)`` appropriate
+    for matrix multiplication, the base elements are two-dimensional matrices
+    and these are taken to be stored in the two last axes of each argument.
+    The corresponding axes keyword would be ``[(-2, -1), (-2, -1), (-2, -1)]``.
+    For simplicity, for generalized ufuncs that operate on 1-dimensional arrays
+    (vectors), a single integer is accepted instead of a single-element tuple,
+    and for generalized ufuncs for which all outputs are scalars, the output
+    tuples can be omitted.
+
 *casting*
 
     .. versionadded:: 1.6
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index c67f60752..bf5a4ead3 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -555,7 +555,8 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
                     PyObject **out_extobj,
                     PyObject **out_typetup,
                     int *out_subok,
-                    PyArrayObject **out_wheremask)
+                    PyArrayObject **out_wheremask,
+                    PyObject **out_axes)
 {
     int i, nargs;
     int nin = ufunc->nin;
@@ -570,6 +571,9 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
 
     *out_extobj = NULL;
     *out_typetup = NULL;
+    if (out_axes != NULL) {
+        *out_axes = NULL;
+    }
     if (out_wheremask != NULL) {
         *out_wheremask = NULL;
     }
@@ -806,6 +810,13 @@ get_ufunc_arguments(PyUFuncObject *ufunc,
             }
 
             switch (str[0]) {
+                case 'a':
+                    /* possible axis argument for generalized ufunc */
+                    if (out_axes != NULL && strcmp(str, "axes") == 0) {
+                        *out_axes = value;
+                        bad_arg = 0;
+                    }
+                    break;
                 case 'c':
                     /* Provides a policy for allowed casting */
                     if (strcmp(str, "casting") == 0) {
@@ -995,6 +1006,10 @@ fail:
     *out_extobj = NULL;
     Py_XDECREF(*out_typetup);
     *out_typetup = NULL;
+    if (out_axes != NULL) {
+        Py_XDECREF(*out_axes);
+        *out_axes = NULL;
+    }
     if (out_wheremask != NULL) {
         Py_XDECREF(*out_wheremask);
         *out_wheremask = NULL;
@@ -1760,6 +1775,155 @@ make_arr_prep_args(npy_intp nin, PyObject *args, PyObject *kwds)
 }
 
 /*
+ * Check whether any of the outputs of a gufunc has core dimensions.
+ */
+static int
+_has_output_coredims(PyUFuncObject *ufunc) {
+    int i;
+    for (i = ufunc->nin; i < ufunc->nin + ufunc->nout; ++i) {
+        if (ufunc->core_num_dims[i] > 0) {
+            return 1;
+        }
+    }
+    return 0;
+}
+
+/*
+ * 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
+ * axes and then the core axes the gufunc operates on.
+ *
+ * Returns 0 on success, and -1 on failure
+ */
+static int
+_parse_axes_arg(PyUFuncObject *ufunc, PyObject *axes, PyArrayObject **op,
+                int broadcast_ndim, int **remap_axis) {
+    int nin = ufunc->nin;
+    int nout = ufunc->nout;
+    int nop = nin + nout;
+    int iop, list_size;
+
+    if (!PyList_Check(axes)) {
+        PyErr_SetString(PyExc_TypeError, "axes should be a list.");
+        return -1;
+    }
+    list_size = PyList_Size(axes);
+    if (list_size != nop) {
+        if (list_size != nin || _has_output_coredims(ufunc)) {
+            PyErr_Format(PyExc_ValueError,
+                         "axes should be a list with an entry for all "
+                         "%d inputs and outputs; entries for outputs can only "
+                         "be omitted if none of them has core axes.",
+                         nop);
+            return -1;
+        }
+        for (iop = nin; iop < nop; iop++) {
+            remap_axis[iop] = NULL;
+        }
+    }
+    for (iop = 0; iop < list_size; ++iop) {
+        int op_ndim, op_ncore, op_nbroadcast;
+        int have_seen_axis[NPY_MAXDIMS] = {0};
+        PyObject *op_axes_tuple, *axis_item;
+        int axis, op_axis;
+
+        op_ncore = ufunc->core_num_dims[iop];
+        if (op[iop] != NULL) {
+            op_ndim = PyArray_NDIM(op[iop]);
+            op_nbroadcast = op_ndim - op_ncore;
+        }
+        else {
+            op_nbroadcast = broadcast_ndim;
+            op_ndim = broadcast_ndim + op_ncore;
+        }
+        /*
+         * Get axes tuple for operand. If not a tuple already, make it one if
+         * there is only one axis (its content is checked later).
+         */
+        op_axes_tuple = PyList_GET_ITEM(axes, iop);
+        if (PyTuple_Check(op_axes_tuple)) {
+            if (PyTuple_Size(op_axes_tuple) != op_ncore) {
+                if (op_ncore == 1) {
+                    PyErr_Format(PyExc_ValueError,
+                                 "axes item %d should be a tuple with a "
+                                 "single element, or an integer", iop);
+                }
+                else {
+                    PyErr_Format(PyExc_ValueError,
+                                 "axes item %d should be a tuple with %d "
+                                 "elements", iop, op_ncore);
+                }
+                return -1;
+            }
+            Py_INCREF(op_axes_tuple);
+        }
+        else if (op_ncore == 1) {
+            op_axes_tuple = PyTuple_Pack(1, op_axes_tuple);
+            if (op_axes_tuple == NULL) {
+                return -1;
+            }
+        }
+        else {
+            PyErr_Format(PyExc_TypeError, "axes item %d should be a tuple",
+                         iop);
+            return -1;
+        }
+        /*
+         * Now create the remap, starting with the core dimensions, and then
+         * adding the remaining broadcast axes that are to be iterated over.
+         */
+        for (axis = op_nbroadcast; axis < op_ndim; axis++) {
+            axis_item = PyTuple_GET_ITEM(op_axes_tuple, axis - op_nbroadcast);
+            op_axis = PyArray_PyIntAsInt(axis_item);
+            if (error_converting(op_axis) ||
+                    (check_and_adjust_axis(&op_axis, op_ndim) < 0)) {
+                Py_DECREF(op_axes_tuple);
+                return -1;
+            }
+            if (have_seen_axis[op_axis]) {
+                PyErr_Format(PyExc_ValueError,
+                             "axes item %d has value %d repeated",
+                             iop, op_axis);
+                Py_DECREF(op_axes_tuple);
+                return -1;
+            }
+            have_seen_axis[op_axis] = 1;
+            remap_axis[iop][axis] = op_axis;
+        }
+        Py_DECREF(op_axes_tuple);
+        /*
+         * Fill the op_nbroadcast=op_ndim-op_ncore axes not yet set,
+         * using have_seen_axis to skip over entries set above.
+         */
+        for (axis = 0, op_axis = 0; axis < op_nbroadcast; axis++) {
+            while (have_seen_axis[op_axis]) {
+                op_axis++;
+            }
+            remap_axis[iop][axis] = op_axis++;
+        }
+        /*
+         * Check whether we are actually remapping anything. Here,
+         * op_axis can only equal axis if all broadcast axes were the same
+         * (i.e., the while loop above was never entered).
+         */
+        if (axis == op_axis) {
+            while (axis < op_ndim && remap_axis[iop][axis] == axis) {
+                axis++;
+            }
+        }
+        if (axis == op_ndim) {
+            remap_axis[iop] = NULL;
+        }
+    } /* end of for(iop) loop over operands */
+    return 0;
+}
+
+#define REMAP_AXIS(iop, axis) ((remap_axis != NULL && \
+                                remap_axis[iop] != NULL)? \
+                               remap_axis[iop][axis] : axis)
+
+/*
  * Validate the core dimensions of all the operands, and collect all of
  * the labelled core dimensions into 'core_dim_sizes'.
  *
@@ -1781,7 +1945,7 @@ make_arr_prep_args(npy_intp nin, PyObject *args, PyObject *kwds)
  */
 static int
 _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
-                   npy_intp* core_dim_sizes) {
+                   npy_intp* core_dim_sizes, int **remap_axis) {
     int i;
     int nin = ufunc->nin;
     int nout = ufunc->nout;
@@ -1796,27 +1960,14 @@ _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
             int dim_offset = ufunc->core_offsets[i];
             int num_dims = ufunc->core_num_dims[i];
             int core_start_dim = PyArray_NDIM(op[i]) - num_dims;
-
-            /* Check if operands have enough dimensions */
-            if (core_start_dim < 0) {
-                PyErr_Format(PyExc_ValueError,
-                        "%s: %s operand %d does not have enough "
-                        "dimensions (has %d, gufunc core with "
-                        "signature %s requires %d)",
-                        ufunc_get_name_cstr(ufunc), i < nin ? "Input" : "Output",
-                        i < nin ? i : i - nin, PyArray_NDIM(op[i]),
-                        ufunc->core_signature, num_dims);
-                return -1;
-            }
-
             /*
              * Make sure every core dimension exactly matches all other core
              * dimensions with the same label.
              */
             for (idim = 0; idim < num_dims; ++idim) {
                 int core_dim_index = ufunc->core_dim_ixs[dim_offset+idim];
-                npy_intp op_dim_size =
-                            PyArray_DIM(op[i], core_start_dim+idim);
+                npy_intp op_dim_size = PyArray_DIM(
+                    op[i], REMAP_AXIS(i, core_start_dim+idim));
 
                 if (core_dim_sizes[core_dim_index] == -1) {
                     core_dim_sizes[core_dim_index] = op_dim_size;
@@ -1950,7 +2101,9 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     /* The sizes of the core dimensions (# entries is ufunc->core_num_dim_ix) */
     npy_intp *core_dim_sizes = inner_dimensions + 1;
     int core_dim_ixs_size;
-
+    /* swapping around of axes */
+    int *remap_axis_memory = NULL;
+    int **remap_axis = NULL;
     /* The __array_prepare__ function to call for each output */
     PyObject *arr_prep[NPY_MAXARGS];
     /*
@@ -1962,8 +2115,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 and typetup contain borrowed references */
-    PyObject *extobj = NULL, *type_tup = NULL;
+    /* When provided, extobj, typetup, and axes contain borrowed references */
+    PyObject *extobj = NULL, *type_tup = NULL, *axes = NULL;
 
     if (ufunc == NULL) {
         PyErr_SetString(PyExc_ValueError, "function not supported");
@@ -1990,12 +2143,30 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     /* Get all the arguments */
     retval = get_ufunc_arguments(ufunc, args, kwds,
                 op, &order, &casting, &extobj,
-                &type_tup, &subok, NULL);
+                &type_tup, &subok, NULL, &axes);
     if (retval < 0) {
         goto fail;
     }
 
     /*
+     * Check that operands have the minimum dimensions required.
+     * (Just checks core; broadcast dimensions are tested by the iterator.)
+     */
+    for (i = 0; i < nop; i++) {
+        if (op[i] != NULL && PyArray_NDIM(op[i]) < ufunc->core_num_dims[i]) {
+            PyErr_Format(PyExc_ValueError,
+                         "%s: %s operand %d does not have enough "
+                         "dimensions (has %d, gufunc core with "
+                         "signature %s requires %d)",
+                         ufunc_get_name_cstr(ufunc),
+                         i < nin ? "Input" : "Output",
+                         i < nin ? i : i - nin, PyArray_NDIM(op[i]),
+                         ufunc->core_signature, ufunc->core_num_dims[i]);
+            goto fail;
+        }
+    }
+
+    /*
      * Figure out the number of iteration dimensions, which
      * is the broadcast result of all the input non-core
      * dimensions.
@@ -2026,8 +2197,27 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         goto fail;
     }
 
+    /* Possibly remap axes. */
+    if (axes) {
+        remap_axis = PyArray_malloc(sizeof(remap_axis[0]) * nop);
+        remap_axis_memory = PyArray_malloc(sizeof(remap_axis_memory[0]) *
+                                                  nop * NPY_MAXDIMS);
+        if (remap_axis == NULL || remap_axis_memory == NULL) {
+            PyErr_NoMemory();
+            goto fail;
+        }
+        for (i=0; i < nop; i++) {
+            remap_axis[i] = remap_axis_memory + i * NPY_MAXDIMS;
+        }
+        retval = _parse_axes_arg(ufunc, axes, op, broadcast_ndim,
+                                 remap_axis);
+        if(retval < 0) {
+            goto fail;
+        }
+    }
+
     /* Collect the lengths of the labelled core dimensions */
-    retval = _get_coredim_sizes(ufunc, op, core_dim_sizes);
+    retval = _get_coredim_sizes(ufunc, op, core_dim_sizes, remap_axis);
     if(retval < 0) {
         goto fail;
     }
@@ -2054,7 +2244,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         /* Broadcast all the unspecified dimensions normally */
         for (idim = 0; idim < broadcast_ndim; ++idim) {
             if (idim >= broadcast_ndim - n) {
-                op_axes_arrays[i][idim] = idim - (broadcast_ndim - n);
+                op_axes_arrays[i][idim] =
+                    REMAP_AXIS(i, idim - (broadcast_ndim - n));
             }
             else {
                 op_axes_arrays[i][idim] = -1;
@@ -2074,7 +2265,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
             for (idim = 0; idim < num_dims; ++idim) {
                 iter_shape[j] = core_dim_sizes[
                                         ufunc->core_dim_ixs[dim_offset + idim]];
-                op_axes_arrays[i][j] = n + idim;
+                op_axes_arrays[i][j] = REMAP_AXIS(i, n + idim);
                 ++j;
             }
         }
@@ -2220,11 +2411,12 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         for (j = 0; j < num_dims; ++j) {
             if (core_start_dim + j >= 0) {
                 /*
-                 * Force the stride to zero when the shape is 1, sot
+                 * Force the stride to zero when the shape is 1, so
                  * that the broadcasting works right.
                  */
-                if (shape[core_start_dim + j] != 1) {
-                    inner_strides[idim++] = strides[core_start_dim + j];
+                int remapped_axis = REMAP_AXIS(i, core_start_dim + j);
+                if (shape[remapped_axis] != 1) {
+                    inner_strides[idim++] = strides[remapped_axis];
                 } else {
                     inner_strides[idim++] = 0;
                 }
@@ -2375,7 +2567,8 @@ fail:
     }
     Py_XDECREF(type_tup);
     Py_XDECREF(arr_prep_args);
-
+    PyArray_free(remap_axis_memory);
+    PyArray_free(remap_axis);
     return retval;
 }
 
@@ -2450,7 +2643,7 @@ PyUFunc_GenericFunction(PyUFuncObject *ufunc,
     /* Get all the arguments */
     retval = get_ufunc_arguments(ufunc, args, kwds,
                 op, &order, &casting, &extobj,
-                &type_tup, &subok, &wheremask);
+                &type_tup, &subok, &wheremask, NULL);
     if (retval < 0) {
         goto fail;
     }
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 57e0ec272..7e1bfbdbe 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -597,6 +597,105 @@ class TestUfunc(object):
         umt.inner1d(a, b, out=c[..., 0])
         assert_array_equal(c[..., 0], np.sum(a*b, axis=-1), err_msg=msg)
 
+    def test_axes_argument(self):
+        # inner1d signature: '(i),(i)->()'
+        in1d = umt.inner1d
+        a = np.arange(27.).reshape((3, 3, 3))
+        b = np.arange(10., 19.).reshape((3, 1, 3))
+        # basic tests on inputs (outputs tested below with matrix_multiply).
+        c = in1d(a, b)
+        assert_array_equal(c, (a * b).sum(-1))
+        # default
+        c = in1d(a, b, axes=[(-1,), (-1,), ()])
+        assert_array_equal(c, (a * b).sum(-1))
+        # integers ok for single axis.
+        c = in1d(a, b, axes=[-1, -1, ()])
+        assert_array_equal(c, (a * b).sum(-1))
+        # mix fine
+        c = in1d(a, b, axes=[(-1,), -1, ()])
+        assert_array_equal(c, (a * b).sum(-1))
+        # can omit last axis.
+        c = in1d(a, b, axes=[-1, -1])
+        assert_array_equal(c, (a * b).sum(-1))
+        # can pass in other types of integer (with __index__ protocol)
+        c = in1d(a, b, axes=[np.int8(-1), np.array(-1, dtype=np.int32)])
+        assert_array_equal(c, (a * b).sum(-1))
+        # swap some axes
+        c = in1d(a, b, axes=[0, 0])
+        assert_array_equal(c, (a * b).sum(0))
+        c = in1d(a, b, axes=[0, 2])
+        assert_array_equal(c, (a.transpose(1, 2, 0) * b).sum(-1))
+        # Check errors for inproperly constructed axes arguments.
+        # should have list.
+        assert_raises(TypeError, in1d, a, b, axes=-1)
+        # needs enough elements
+        assert_raises(ValueError, in1d, a, b, axes=[-1])
+        # should pass in indices.
+        assert_raises(TypeError, in1d, a, b, axes=[-1.0, -1.0])
+        assert_raises(TypeError, in1d, a, b, axes=[(-1.0,), -1])
+        assert_raises(TypeError, in1d, a, b, axes=[None, 1])
+        # cannot pass an index unless there is only one dimension
+        # (output is wrong in this case)
+        assert_raises(TypeError, in1d, a, b, axes=[-1, -1, -1])
+        # or pass in generally the wrong number of axes
+        assert_raises(ValueError, in1d, a, b, axes=[-1, -1, (-1,)])
+        assert_raises(ValueError, in1d, a, b, axes=[-1, (-2, -1), ()])
+        # axes need to have same length.
+        assert_raises(ValueError, in1d, a, b, axes=[0, 1])
+
+        # matrix_multiply signature: '(m,n),(n,p)->(m,p)'
+        mm = umt.matrix_multiply
+        a = np.arange(12).reshape((2, 3, 2))
+        b = np.arange(8).reshape((2, 2, 2, 1)) + 1
+        # Sanity check.
+        c = mm(a, b)
+        assert_array_equal(c, np.matmul(a, b))
+        # Default axes.
+        c = mm(a, b, axes=[(-2, -1), (-2, -1), (-2, -1)])
+        assert_array_equal(c, np.matmul(a, b))
+        # Default with explicit axes.
+        c = mm(a, b, axes=[(1, 2), (2, 3), (2, 3)])
+        assert_array_equal(c, np.matmul(a, b))
+        # swap some axes.
+        c = mm(a, b, axes=[(0, -1), (1, 2), (-2, -1)])
+        assert_array_equal(c, np.matmul(a.transpose(1, 0, 2),
+                                        b.transpose(0, 3, 1, 2)))
+        # Default with output array.
+        c = np.empty((2, 2, 3, 1))
+        d = mm(a, b, out=c, axes=[(1, 2), (2, 3), (2, 3)])
+        assert_(c is d)
+        assert_array_equal(c, np.matmul(a, b))
+        # Transposed output array
+        c = np.empty((1, 2, 2, 3))
+        d = mm(a, b, out=c, axes=[(-2, -1), (-2, -1), (3, 0)])
+        assert_(c is d)
+        assert_array_equal(c, np.matmul(a, b).transpose(3, 0, 1, 2))
+        # Check errors for inproperly constructed axes arguments.
+        # wrong argument
+        assert_raises(TypeError, mm, a, b, axis=1)
+        # axes should be list
+        assert_raises(TypeError, mm, a, b, axes=1)
+        assert_raises(TypeError, mm, a, b, axes=((-2, -1), (-2, -1), (-2, -1)))
+        # list needs to have right length
+        assert_raises(ValueError, mm, a, b, axes=[])
+        assert_raises(ValueError, mm, a, b, axes=[(-2, -1)])
+        # list should contain tuples for multiple axes
+        assert_raises(TypeError, mm, a, b, axes=[-1, -1, -1])
+        assert_raises(TypeError, mm, a, b, axes=[(-2, -1), (-2, -1), -1])
+        assert_raises(TypeError,
+                      mm, a, b, axes=[[-2, -1], [-2, -1], [-2, -1]])
+        assert_raises(TypeError,
+                      mm, a, b, axes=[(-2, -1), (-2, -1), [-2, -1]])
+        assert_raises(TypeError, mm, a, b, axes=[(-2, -1), (-2, -1), None])
+        # tuples should not have duplicated values
+        assert_raises(ValueError, mm, a, b, axes=[(-2, -1), (-2, -1), (-2, -2)])
+        # arrays should have enough axes.
+        z = np.zeros((2, 2))
+        assert_raises(ValueError, mm, z, z[0])
+        assert_raises(ValueError, mm, z, z, out=z[:, 0])
+        assert_raises(ValueError, mm, z[1], z, axes=[0, 1])
+        assert_raises(ValueError, mm, z, z, out=z[0], axes=[0, 1])
+
     def test_innerwt(self):
         a = np.arange(6).reshape((2, 3))
         b = np.arange(10, 16).reshape((2, 3))