ENH: test, document, implement flexible signature.

Goal is allow signatures like (m?,n),(n,p?)->(m?,p?) for matmul.

5 files changed, 289 insertions, 79 deletions

diff --git a/doc/source/reference/c-api.generalized-ufuncs.rst b/doc/source/reference/c-api.generalized-ufuncs.rst
index dd8cf6558..59aa11d00 100644
--- a/doc/source/reference/c-api.generalized-ufuncs.rst
+++ b/doc/source/reference/c-api.generalized-ufuncs.rst
@@ -127,34 +127,49 @@ The formal syntax of signatures is as follows::
     <Output arguments>     ::= <Argument list>
     <Argument list>        ::= nil | <Argument> | <Argument> "," <Argument list>
     <Argument>             ::= "(" <Core dimension list> ")"
-    <Core dimension list>  ::= nil | <Core dimension name> |
-                               <Core dimension name> "," <Core dimension list>
-    <Core dimension name>  ::= valid Python variable name
-
+    <Core dimension list>  ::= nil | <Core dimension> |
+                               <Core dimension> "," <Core dimension list>
+    <Core dimension>       ::= <Dimension name> <Dimension modifier>
+    <Dimension name>       ::= valid Python variable name | valid integer
+    <Dimension modifier>   ::= nil | "?"
 
 Notes:
 
 #. All quotes are for clarity.
-#. Core dimensions that share the same name must have the exact same size.
+#. Unmodified core dimensions that share the same name must have the same size.
    Each dimension name typically corresponds to one level of looping in the
    elementary function's implementation.
 #. White spaces are ignored.
+#. An integer as a dimension name freezes that dimension to the value.
+#. If the name is suffixed with the "?" modifier, the dimension is a core
+   dimension only if it exists on all inputs and outputs that share it;
+   otherwise it is ignored (and replaced by a dimension of size 1 for the
+   elementary function).
 
 Here are some examples of signatures:
 
-+-------------+------------------------+-----------------------------------+
-| add         | ``(),()->()``          |                                   |
-+-------------+------------------------+-----------------------------------+
-| inner1d     | ``(i),(i)->()``        |                                   |
-+-------------+------------------------+-----------------------------------+
-| sum1d       | ``(i)->()``            |                                   |
-+-------------+------------------------+-----------------------------------+
-| dot2d       | ``(m,n),(n,p)->(m,p)`` | matrix multiplication             |
-+-------------+------------------------+-----------------------------------+
-| outer_inner | ``(i,t),(j,t)->(i,j)`` | inner over the last dimension,    |
-|             |                        | outer over the second to last,    |
-|             |                        | and loop/broadcast over the rest. |
-+-------------+------------------------+-----------------------------------+
++-------------+----------------------------+-----------------------------------+
+| add         | ``(),()->()``              |                                   |
++-------------+----------------------------+-----------------------------------+
+| sum1d       | ``(i)->()``                |                                   |
++-------------+----------------------------+-----------------------------------+
+| inner1d     | ``(i),(i)->()``            |                                   |
++-------------+----------------------------+-----------------------------------+
+| matmat      | ``(m,n),(n,p)->(m,p)``     | matrix multiplication             |
++-------------+----------------------------+-----------------------------------+
+| vecmat      | ``(n),(n,p)->(p)``         | vector-matrix multiplication      |
++-------------+----------------------------+-----------------------------------+
+| matvec      | ``(m,n),(n)->(m)``         | matrix-vector multiplication      |
++-------------+----------------------------+-----------------------------------+
+| matmul      | ``(m?,n),(n,p?)->(m?,p?)`` | combination of the four above     |
++-------------+----------------------------+-----------------------------------+
+| cross1d     | ``(3),(3)->(3)``           | cross product where the last      |
+|             |                            | dimension must be 3               |
++-------------+----------------------------+-----------------------------------+
+| outer_inner | ``(i,t),(j,t)->(i,j)``     | inner over the last dimension,    |
+|             |                            | outer over the second to last,    |
+|             |                            | and loop/broadcast over the rest. |
++-------------+----------------------------+-----------------------------------+
 
 C-API for implementing Elementary Functions
 -------------------------------------------
diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 7fe5e9f82..24848c3d0 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -213,9 +213,15 @@ typedef struct _tagPyUFuncObject {
         /* New in version 1 and above */
 
         /*
-         * sizes of frozen core dimensions, or -1 if unset
+         * sizes of `frozen` core dimensions, -1 if unset (not frozen)
          */
-         npy_intp *core_dim_szs;
+        npy_intp *core_dim_szs;
+
+        /*
+         * for each core_num_dim_ix, 1 for flexible (signature has ?) 0 otherwise
+         */
+        int *core_dim_flexible;
+
 } PyUFuncObject;
 
 #include "arrayobject.h"
diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src
index 1d00fa76c..0a152b03b 100644
--- a/numpy/core/src/umath/_umath_tests.c.src
+++ b/numpy/core/src/umath/_umath_tests.c.src
@@ -452,9 +452,10 @@ static PyObject *
 UMath_Tests_test_signature(PyObject *NPY_UNUSED(dummy), PyObject *args)
 {
     int nin, nout, i;
-    PyObject *signature, *sig_str;
-    PyUFuncObject *f = NULL;
-    PyObject *core_num_dims = NULL, *core_dim_ixs = NULL;
+    PyObject *signature=NULL, *sig_str=NULL;
+    PyUFuncObject *f=NULL;
+    PyObject *core_num_dims=NULL, *core_dim_ixs=NULL, *core_dim_szs=NULL;
+    PyObject *core_dim_flexible=NULL;
     int core_enabled;
     int core_num_ixs = 0;
 
@@ -517,13 +518,44 @@ UMath_Tests_test_signature(PyObject *NPY_UNUSED(dummy), PyObject *args)
         Py_INCREF(Py_None);
         core_dim_ixs = Py_None;
     }
+    if (f->core_dim_szs != NULL) {
+        core_dim_szs = PyTuple_New(f->core_num_dim_ix);
+        if (core_num_dims == NULL) {
+            goto fail;
+        }
+        for (i = 0; i < f->core_num_dim_ix; i++) {
+            PyObject * val = PyLong_FromLong(f->core_dim_szs[i]);
+            PyTuple_SET_ITEM(core_dim_szs, i, val);
+        }
+    }
+    else {
+        Py_INCREF(Py_None);
+        core_dim_szs = Py_None;
+    }
+    if (f->core_num_dim_ix != NULL) {
+        core_dim_flexible = PyTuple_New(f->core_num_dim_ix);
+        if (core_dim_flexible == NULL) {
+            goto fail;
+        }
+        for (i = 0; i < f->core_num_dim_ix; i++) {
+            PyObject * val = PyLong_FromLong(f->core_dim_flexible[i]);
+            PyTuple_SET_ITEM(core_dim_flexible, i, val);
+        }
+    }
+    else {
+        Py_INCREF(Py_None);
+        core_dim_flexible = Py_None;
+    }
     Py_DECREF(f);
-    return Py_BuildValue("iOO", core_enabled, core_num_dims, core_dim_ixs);
+    return Py_BuildValue("iOOOO", core_enabled, core_num_dims,
+                         core_dim_ixs, core_dim_szs, core_dim_flexible);
 
 fail:
     Py_XDECREF(f);
     Py_XDECREF(core_num_dims);
     Py_XDECREF(core_dim_ixs);
+    Py_XDECREF(core_dim_szs);
+    Py_XDECREF(core_dim_flexible);
     return NULL;
 }
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index e7031557e..1f939c2c0 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -371,6 +371,9 @@ _get_end_of_name(const char* str, int offset)
     while (_is_alnum_underscore(str[ret])) {
         ret++;
     }
+    if (str[ret] == '?') {
+        ret ++;
+    }
     return ret;
 }
 
@@ -432,12 +435,17 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
     /* The next two items will be shrunk later */
     ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len);
     ufunc->core_dim_szs = PyArray_malloc(sizeof(npy_intp) * len);
+    ufunc->core_dim_flexible = PyArray_malloc(sizeof(int) * len);
 
     if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL ||
-        ufunc->core_offsets == NULL || ufunc->core_dim_szs == NULL) {
+        ufunc->core_offsets == NULL || ufunc->core_dim_szs == NULL ||
+        ufunc->core_dim_flexible == NULL) {
         PyErr_NoMemory();
         goto fail;
     }
+    for (i = 0; i < len; i++) {
+        ufunc->core_dim_flexible[i] = -1;
+    }
 
     i = _next_non_white_space(signature, 0);
     while (signature[i] != '\0') {
@@ -471,7 +479,7 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
 
             if (!_is_alpha_underscore(signature[i]) &&
                 (frozen_size = _get_size(signature + i)) < 0) {
-                parse_error = "expect dimension name or frozen size";
+                parse_error = "expect dimension name or non-zero frozen size";
                 goto fail;
             }
             while (j < ufunc->core_num_dim_ix) {
@@ -486,9 +494,23 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
                 ufunc->core_num_dim_ix++;
             }
             ufunc->core_dim_ixs[cur_core_dim] = j;
+            i = _get_end_of_name(signature, i);
+            if (i > 0 && signature[i - 1] == '?') {
+                if (ufunc->core_dim_flexible[j] == 0) {
+                    parse_error = "? cannot be used, name already seen without ?";
+                    goto fail;
+                }
+                ufunc->core_dim_flexible[j] = 1;
+            }
+            else {
+                if (ufunc->core_dim_flexible[j] == 1) {
+                    parse_error = "? must be used, name already seen with ?";
+                    goto fail;
+                }
+                ufunc->core_dim_flexible[j] = 0;
+            }
             cur_core_dim++;
             nd++;
-            i = _get_end_of_name(signature, i);
             i = _next_non_white_space(signature, i);
             if (signature[i] != ',' && signature[i] != ')') {
                 parse_error = "expect ',' or ')'";
@@ -527,8 +549,10 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
     }
     ufunc->core_dim_ixs = PyArray_realloc(ufunc->core_dim_ixs,
             sizeof(int)*cur_core_dim);
-    // ufunc->core_dim_szs = PyArray_realloc(ufunc->core_dim_szs,
-            // sizeof(npy_intp)*ufunc->core_num_dim_ix);
+    ufunc->core_dim_szs = PyArray_realloc(ufunc->core_dim_szs,
+            sizeof(npy_intp)*ufunc->core_num_dim_ix);
+    ufunc->core_dim_flexible = PyArray_realloc(ufunc->core_dim_flexible,
+            sizeof(int)*(ufunc->core_num_dim_ix));
 
     /* check for trivial core-signature, e.g. "(),()->()" */
     if (cur_core_dim == 0) {
@@ -2103,6 +2127,7 @@ _parse_axis_arg(PyUFuncObject *ufunc, int core_num_dims[], PyObject *axis,
  */
 static int
 _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
+                   int *core_num_dims, int * flexible_activated,
                    npy_intp* core_dim_sizes, int **remap_axis) {
     int i;
     int nin = ufunc->nin;
@@ -2110,34 +2135,55 @@ _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
     int nop = nin + nout;
 
     for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
-        core_dim_sizes[i] = -1;
+        /* support fixed-size dim names */
+        core_dim_sizes[i] = ufunc->core_dim_szs[i];
     }
     for (i = 0; i < nop; ++i) {
         if (op[i] != NULL) {
             int idim;
             int dim_offset = ufunc->core_offsets[i];
-            int num_dims = ufunc->core_num_dims[i];
+            int num_dims = core_num_dims[i];
             int core_start_dim = PyArray_NDIM(op[i]) - num_dims;
+            int dim_delta = 0;
+
+            /* 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) {
+            for (idim = 0; idim < ufunc->core_num_dims[i]; ++idim) {
                 int core_dim_index = ufunc->core_dim_ixs[dim_offset+idim];
-                npy_intp op_dim_size = PyArray_DIM(
-                    op[i], REMAP_AXIS(i, core_start_dim+idim));
-
+                npy_intp op_dim_size;
+                if (flexible_activated[core_dim_index] > 0) {
+                    op_dim_size = 0;
+                    dim_delta ++;
+                }
+                else {
+                    op_dim_size = PyArray_DIM(op[i],
+                             REMAP_AXIS(i, core_start_dim + idim - dim_delta));
+                }
                 if (core_dim_sizes[core_dim_index] == -1) {
                     core_dim_sizes[core_dim_index] = op_dim_size;
                 }
                 else if (op_dim_size != core_dim_sizes[core_dim_index]) {
                     PyErr_Format(PyExc_ValueError,
-                            "%s: %s operand %d has a mismatch in its "
+                            "1 %s: %s operand %d has a mismatch in its "
                             "core dimension %d, with gufunc "
                             "signature %s (size %zd is different "
                             "from %zd)",
                             ufunc_get_name_cstr(ufunc), i < nin ? "Input" : "Output",
-                            i < nin ? i : i - nin, idim,
+                            i < nin ? i : i - nin, idim - dim_delta,
                             ufunc->core_signature, op_dim_size,
                             core_dim_sizes[core_dim_index]);
                     return -1;
@@ -2163,7 +2209,7 @@ _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op,
         int out_op;
         for (out_op = nin; out_op < nop; ++out_op) {
             int first_idx = ufunc->core_offsets[out_op];
-            int last_idx = first_idx + ufunc->core_num_dims[out_op];
+            int last_idx = first_idx + core_num_dims[out_op];
             for (i = first_idx; i < last_idx; ++i) {
                 if (ufunc->core_dim_ixs[i] == missing_core_dim) {
                     break;
@@ -2240,6 +2286,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     int *core_num_dims;
     int op_axes_arrays[NPY_MAXARGS][NPY_MAXDIMS];
     int *op_axes[NPY_MAXARGS];
+    int n_dims_used[NPY_MAXARGS];
+    int flexible_activated[NPY_MAXARGS];
 
     npy_uint32 op_flags[NPY_MAXARGS];
     npy_intp iter_shape[NPY_MAXARGS];
@@ -2293,6 +2341,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     for (i = 0; i < nop; ++i) {
         dtypes[i] = NULL;
         arr_prep[i] = NULL;
+        n_dims_used[i] = ufunc->core_num_dims[i];
+    }
+    for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
+        flexible_activated[i] = (ufunc->core_dim_flexible[i] > 0) ? -1 : 0;
     }
 
     NPY_UF_DBG_PRINT("Getting arguments\n");
@@ -2344,19 +2396,104 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
      * (Just checks core; broadcast dimensions are tested by the iterator.)
      */
     for (i = 0; i < nop; i++) {
-        if (op[i] != NULL && PyArray_NDIM(op[i]) < core_num_dims[i]) {
-            PyErr_Format(PyExc_ValueError,
+        /* TODO: store _n_required after parsing, ahead-of-time */
+        int n_required = ufunc->core_num_dims[i];
+        if (n_required > 0) {
+            /* can be less if dim is flexible */
+            for (j = ufunc->core_offsets[i];
+                     j < ufunc->core_offsets[i] + core_num_dims[i]; j++) {
+                int dim_index = ufunc->core_dim_ixs[j];
+                n_required -= ufunc->core_dim_flexible[dim_index];
+            }
+        }
+        if (op[i] != NULL) {
+            n_dims_used[i] = PyArray_NDIM(op[i]);
+            if (n_dims_used[i] < n_required) {
+                PyErr_Format(PyExc_ValueError,
                          "%s: %s operand %d does not have enough "
                          "dimensions (has %d, gufunc core with "
                          "signature %s requires %d)",
                          ufunc_name,
                          i < nin ? "Input" : "Output",
-                         i < nin ? i : i - nin,
-                         PyArray_NDIM(op[i]),
-                         ufunc->core_signature,
-                         core_num_dims[i]);
-            retval = -1;
-            goto fail;
+                         i < nin ? i : i - nin, PyArray_NDIM(op[i]),
+                         ufunc->core_signature, n_required);
+                retval = -1;
+                goto fail;
+            }
+            if (n_dims_used[i] < core_num_dims[i]) {
+                int flexible_used = core_num_dims[i] - n_dims_used[i];
+                /* one (or more?) flexible signature names have been used. Mark them */
+                for (j = ufunc->core_offsets[i];
+                     j < ufunc->core_offsets[i] + core_num_dims[i];
+                     j++) {
+                    int dim_index = ufunc->core_dim_ixs[j];
+                    if (flexible_activated[dim_index] < 0) {
+                        flexible_activated[dim_index] = 1;
+                    }
+                    flexible_used -= flexible_activated[dim_index];
+                    if (flexible_used <= 0) {
+                        break;
+                    }
+                }
+                if (flexible_used > 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, n_required + flexible_used);
+                    goto fail;
+                }
+            }
+            else {
+                /* never use more than the signature requires */
+                n_dims_used[i] = core_num_dims[i];
+                /* make sure any flexible dimensions are used in the same way */
+                for (j = ufunc->core_offsets[i];
+                     j < ufunc->core_offsets[i] + core_num_dims[i];
+                     j++) {
+                    if (j <= ufunc->core_num_dim_ix) {
+                        int dim_index = ufunc->core_dim_ixs[j];
+                        if (flexible_activated[dim_index] < 0) {
+                            flexible_activated[dim_index] = 0;
+                        }
+                        else if (flexible_activated[dim_index] == 1) {
+                            PyErr_Format(PyExc_ValueError,
+                                "%s: %s operand %d with flexible dimension "
+                                "index %d signature %s previouosly marked "
+                                "flexible but now is not",
+                                 ufunc_get_name_cstr(ufunc),
+                                 i < nin ? "Input" : "Output",
+                                 i < nin ? i : i - nin, dim_index,
+                                 ufunc->core_signature);
+                            goto fail;
+                        }
+                    }
+                }
+            }
+        }
+        else {
+            /* NULL output provided. Use flexible_used to determine ndims_used */
+            n_dims_used[i] = core_num_dims[i];
+            for (j = ufunc->core_offsets[i];
+                 j < ufunc->core_offsets[i] + core_num_dims[i]; j++) {
+                if (j <= ufunc->core_num_dim_ix) {
+                    int dim_index = ufunc->core_dim_ixs[j];
+                    if (flexible_activated[dim_index] < 0) {
+                        PyErr_Format(PyExc_ValueError,
+                            "%s: %s operand %d with flexible dimension index "
+                            " %d signature %s cannot be uniquely determined",
+                            ufunc_get_name_cstr(ufunc),
+                            i < nin ? "Input" : "Output",
+                            i < nin ? i : i - nin, dim_index,
+                            ufunc->core_signature);
+                        goto fail;
+                    }
+                    n_dims_used[i] -= flexible_activated[dim_index];
+                }
+            }
         }
     }
 
@@ -2367,7 +2504,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
      */
     broadcast_ndim = 0;
     for (i = 0; i < nin; ++i) {
-        int n = PyArray_NDIM(op[i]) - core_num_dims[i];
+        int n = PyArray_NDIM(op[i]) - n_dims_used[i];
         if (n > broadcast_ndim) {
             broadcast_ndim = n;
         }
@@ -2381,7 +2518,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
      */
     iter_ndim = broadcast_ndim;
     for (i = nin; i < nop; ++i) {
-        iter_ndim += core_num_dims[i];
+        iter_ndim += n_dims_used[i];
     }
     if (iter_ndim > NPY_MAXDIMS) {
         PyErr_Format(PyExc_ValueError,
@@ -2415,20 +2552,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
             goto fail;
         }
     }
-    /*
-     * Validate the core dimensions of all the operands,
-     * and collect all of the labeled core dimension sizes
-     * into the array 'core_dim_sizes'. Initialize them to
-     * 1, for example in the case where the operand broadcasts
-     * to a core dimension, it won't be visited.
-     */
-    for (i = 0; i < ufunc->core_num_dim_ix; ++i) {
-        npy_intp frozen_size = ufunc->core_dim_szs[i];
-        core_dim_sizes[i] = frozen_size == -1 ? 1 : frozen_size;
-    }
 
     /* Collect the lengths of the labelled core dimensions */
-    retval = _get_coredim_sizes(ufunc, op, core_dim_sizes, remap_axis);
+    retval = _get_coredim_sizes(ufunc, op, n_dims_used, flexible_activated,
+                                core_dim_sizes, remap_axis);
     if(retval < 0) {
         goto fail;
     }
@@ -2444,11 +2571,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         int n;
 
         if (op[i]) {
-            /*
-             * Note that n may be negative if broadcasting
-             * extends into the core dimensions.
-             */
-            n = PyArray_NDIM(op[i]) - core_num_dims[i];
+            n = PyArray_NDIM(op[i]) - n_dims_used[i];
         }
         else {
             n = broadcast_ndim;
@@ -2473,16 +2596,27 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
         if (i >= nin) {
             int dim_offset = ufunc->core_offsets[i];
             int num_dims = core_num_dims[i];
+            int has_flexible=0;
             /*
              * Fill in 'iter_shape' and 'op_axes' for the core dimensions
              * of this output. Here, we have to be careful: if keepdims
-             * was used, then this axis is not a real core dimension,
-             * but is being added back for broadcasting, so its size is 1.
+             * was used or if it is a flexible dimension, then this axis
+             * is not a real core dimension, but is being added back for broadcasting,
+             * so its size is 1.
              */
             for (idim = 0; idim < num_dims; ++idim) {
-                iter_shape[j] = keepdims ? 1 : core_dim_sizes[
-                                        ufunc->core_dim_ixs[dim_offset + idim]];
-                op_axes_arrays[i][j] = REMAP_AXIS(i, n + idim);
+                if (keepdims) {
+                    iter_shape[j] = 1;
+                } else {
+                    int core_index = ufunc->core_dim_ixs[dim_offset + idim];
+                    if (flexible_activated[core_index]) {
+                        /* skip it */
+                        has_flexible ++;
+                        continue;
+                    }
+                    iter_shape[j] = core_dim_sizes[core_index];
+                }
+                op_axes_arrays[i][j] = REMAP_AXIS(i, n + idim - has_flexible);
                 ++j;
             }
         }
@@ -2613,6 +2747,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc,
     idim = nop;
     for (i = 0; i < nop; ++i) {
         int num_dims = ufunc->core_num_dims[i];
+        /* Could be negative if flexible dims are used, or if keepdims */
         int core_start_dim = PyArray_NDIM(op[i]) - num_dims;
         /*
          * Need to use the arrays in the iterator, not op, because
diff --git a/numpy/core/tests/test_ufunc.py b/numpy/core/tests/test_ufunc.py
index 0e564e305..c77e5f0b8 100644
--- a/numpy/core/tests/test_ufunc.py
+++ b/numpy/core/tests/test_ufunc.py
@@ -288,17 +288,45 @@ class TestUfunc(object):
 
     def test_signature(self):
         # the arguments to test_signature are: nin, nout, core_signature
-        # pass
-        enabled, num_dims, ixs = umt.test_signature(2, 1, "(i),(i)->()")
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, "(i),(i)->()")
         assert_equal(enabled, 1)
         assert_equal(num_dims, (1,  1,  0))
         assert_equal(ixs, (0, 0))
+        assert_equal(szs, (-1,))
 
         # empty core signature; treat as plain ufunc (with trivial core)
-        enabled, num_dims, ixs = umt.test_signature(2, 1, "(),()->()")
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, "(),()->()")
         assert_equal(enabled, 0)
         assert_equal(num_dims, (0,  0,  0))
         assert_equal(ixs, ())
+        assert_equal(szs, ())
+
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, "(n,k),(k,m)->(n,m)")
+        assert_equal(enabled, 1)
+        assert_equal(num_dims, (2, 2, 2))
+        assert_equal(ixs, (0, 1, 1, 2, 0, 2))
+        assert_equal(szs, (-1, -1, -1))
+
+        # more complicated names for variables
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, "(i1,i2),(J_1)->(_kAB)")
+        assert_equal(enabled, 1)
+        assert_equal(num_dims, (2, 1, 1))
+        assert_equal(ixs, (0, 1, 2, 3))
+        assert_equal(szs, (-1, -1, -1, -1))
+
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, u"(i1, i12),   (J_1)->(i12, i2)")
+        assert_equal(enabled, 1)
+        assert_equal(num_dims, (2, 1, 2))
+        assert_equal(ixs, (0, 1, 2, 1, 3))
+        assert_equal(szs, (-1, -1, -1, -1))
+        assert_equal(flex, (0, 0, 0, 0))
+
+        enabled, num_dims, ixs, szs, flex = umt.test_signature(2, 1, "(n?,k),(k,m?)->(n?,m?)")
+        assert_equal(enabled, 1)
+        assert_equal(num_dims, (2, 2, 2))
+        assert_equal(ixs, (0, 1, 1, 2, 0, 2))
+        assert_equal(szs, (-1, -1, -1))
+        assert_equal(flex, (1, 0, 1))
 
         # in the following calls, a ValueError should be raised because
         # of error in core signature
@@ -336,12 +364,6 @@ class TestUfunc(object):
         except ValueError:
             pass
 
-        # more complicated names for variables
-        enabled, num_dims, ixs = umt.test_signature(2, 1, "(i1,i2),(J_1)->(_kAB)")
-        assert_equal(enabled, 1)
-        assert_equal(num_dims, (2, 1, 1))
-        assert_equal(ixs, (0, 1, 2, 3))
-
     def test_get_signature(self):
         assert_equal(umt.inner1d.signature, "(i),(i)->()")