ENH: Add frozen dimensions to gufunc signatures

3 files changed, 127 insertions, 6 deletions

diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 4b1b3d325..a191e5766 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -209,6 +209,11 @@ typedef struct _tagPyUFuncObject {
          * set by nditer object.
          */
         npy_uint32 iter_flags;
+
+        /*
+         * sizes of frozen core dimensions, or -1 if unset
+         */
+         npy_intp *core_dim_szs;
 } 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 fcbdbe330..1d00fa76c 100644
--- a/numpy/core/src/umath/_umath_tests.c.src
+++ b/numpy/core/src/umath/_umath_tests.c.src
@@ -195,6 +195,41 @@ static void
 
 /**end repeat**/
 
+char *cross1d_signature = "(3),(3)->(3)";
+
+/**begin repeat
+
+   #TYPE=LONG,DOUBLE#
+   #typ=npy_long, npy_double#
+*/
+
+/*
+ *  This implements the cross product:
+ *        out[n, 0] = in1[n, 1]*in2[n, 2] - in1[n, 2]*in2[n, 1]
+ *        out[n, 1] = in1[n, 2]*in2[n, 0] - in1[n, 0]*in2[n, 2]
+ *        out[n, 2] = in1[n, 0]*in2[n, 1] - in1[n, 1]*in2[n, 0]
+ */
+static void
+@TYPE@_cross1d(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED(func))
+{
+    INIT_OUTER_LOOP_3
+    npy_intp is1=steps[0], is2=steps[1], os = steps[2];
+    BEGIN_OUTER_LOOP_3
+        char *ip1=args[0], *ip2=args[1], *op=args[2];
+
+        *(@typ@ *)op = *(@typ@ *)(ip1 + is1) * *(@typ@ *)(ip2 + 2*is2) -
+                       *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)(ip2 + is2);
+        op += os;
+        *(@typ@ *)op = *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)ip2 -
+                       *(@typ@ *)ip1 * *(@typ@ *)(ip2 + 2*is2);
+        op += os;
+        *(@typ@ *)op = *(@typ@ *)ip1 * *(@typ@ *)(ip2 + is2) -
+                       *(@typ@ *)(ip1 + is1) * *(@typ@ *)ip2;
+    END_OUTER_LOOP
+}
+
+/**end repeat**/
+
 char *euclidean_pdist_signature = "(n,d)->(p)";
 
 /**begin repeat
@@ -285,6 +320,26 @@ static void
 
 /**end repeat**/
 
+/*  The following lines were generated using a slightly modified
+    version of code_generators/generate_umath.py and adding these
+    lines to defdict:
+
+defdict = {
+'inner1d' :
+    Ufunc(2, 1, None_,
+        r'''inner on the last dimension and broadcast on the rest \n"
+        "     \"(i),(i)->()\" \n''',
+          TD('ld'),
+          ),
+'innerwt' :
+    Ufunc(3, 1, None_,
+          r'''inner1d with a weight argument \n"
+          "     \"(i),(i),(i)->()\" \n''',
+          TD('ld'),
+          ),
+}
+
+*/
 
 static PyUFuncGenericFunction inner1d_functions[] = { LONG_inner1d, DOUBLE_inner1d };
 static void * inner1d_data[] = { (void *)NULL, (void *)NULL };
@@ -295,7 +350,9 @@ static char innerwt_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_LONG, NPY
 static PyUFuncGenericFunction matrix_multiply_functions[] = { LONG_matrix_multiply, FLOAT_matrix_multiply, DOUBLE_matrix_multiply };
 static void *matrix_multiply_data[] = { (void *)NULL, (void *)NULL, (void *)NULL };
 static char matrix_multiply_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG,  NPY_FLOAT, NPY_FLOAT, NPY_FLOAT,  NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
-
+static PyUFuncGenericFunction cross1d_functions[] = { LONG_cross1d, DOUBLE_cross1d };
+static void * cross1d_data[] = { (void *)NULL, (void *)NULL };
+static char cross1d_signatures[] = { NPY_LONG, NPY_LONG, NPY_LONG, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE };
 static PyUFuncGenericFunction euclidean_pdist_functions[] =
                             { FLOAT_euclidean_pdist, DOUBLE_euclidean_pdist };
 static void *eucldiean_pdist_data[] = { (void *)NULL, (void *)NULL };
@@ -376,6 +433,16 @@ addUfuncs(PyObject *dictionary) {
     }
     PyDict_SetItemString(dictionary, "inner1d_no_doc", f);
     Py_DECREF(f);
+    f = PyUFunc_FromFuncAndDataAndSignature(cross1d_functions, cross1d_data,
+                    cross1d_signatures, 2, 2, 1, PyUFunc_None, "cross1d",
+                    "cross product on the last dimension and broadcast on the rest \n"\
+                    "     \"(3),(3)->(3)\" \n",
+                    0, cross1d_signature);
+    if (f == NULL) {
+        return -1;
+    }
+    PyDict_SetItemString(dictionary, "cross1d", f);
+    Py_DECREF(f);
 
     return 0;
 }
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index a3fd72839..47a4d876b 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -340,6 +340,28 @@ _is_alnum_underscore(char ch)
 }
 
 /*
+ * Convert a string into a number
+ */
+ static npy_intp
+ _get_size(const char* str)
+ {
+    char *stop;
+    #if defined(_MSC_VER)
+        #define strtoll _strtoi64
+    #endif
+    npy_intp size = (npy_intp)strtoll(str, &stop, 10);
+    #if defined(_MSC_VER)
+        #undef strtoll
+    #endif
+
+    if (stop == str || _is_alpha_underscore(*stop)) {
+        /* not a well formed number */
+         return -1;
+    }
+    return size;
+ }
+
+/*
  * Return the ending position of a variable name
  */
 static int
@@ -406,10 +428,13 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
     ufunc->core_enabled = 1;
     ufunc->core_num_dim_ix = 0;
     ufunc->core_num_dims = PyArray_malloc(sizeof(int) * ufunc->nargs);
-    ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len); /* shrink this later */
     ufunc->core_offsets = PyArray_malloc(sizeof(int) * ufunc->nargs);
-    if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL
-        || ufunc->core_offsets == NULL) {
+    /* 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);
+
+    if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL ||
+        ufunc->core_offsets == NULL || ufunc->core_dim_szs == NULL) {
         PyErr_NoMemory();
         goto fail;
     }
@@ -438,8 +463,15 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
         while (signature[i] != ')') {
             /* loop over core dimensions */
             int j = 0;
-            if (!_is_alpha_underscore(signature[i])) {
-                parse_error = "expect dimension name";
+            npy_intp frozen_size = -1;
+            if (signature[i] == '\0') {
+                parse_error = "unexpected end of signature string";
+                goto fail;
+            }
+
+            if (!_is_alpha_underscore(signature[i]) &&
+                (frozen_size = _get_size(signature + i)) < 0) {
+                parse_error = "expect dimension name or frozen size";
                 goto fail;
             }
             while (j < ufunc->core_num_dim_ix) {
@@ -450,6 +482,7 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
             }
             if (j >= ufunc->core_num_dim_ix) {
                 var_names[j] = signature+i;
+                ufunc->core_dim_szs[j] = frozen_size;
                 ufunc->core_num_dim_ix++;
             }
             ufunc->core_dim_ixs[cur_core_dim] = j;
@@ -494,6 +527,9 @@ _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);
+
     /* check for trivial core-signature, e.g. "(),()->()" */
     if (cur_core_dim == 0) {
         ufunc->core_enabled = 0;
@@ -2379,6 +2415,17 @@ 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);
@@ -4651,6 +4698,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
     ufunc->core_dim_ixs = NULL;
     ufunc->core_offsets = NULL;
     ufunc->core_signature = NULL;
+    ufunc->core_dim_szs = NULL;
     if (signature != NULL) {
         if (_parse_signature(ufunc, signature) != 0) {
             Py_DECREF(ufunc);
@@ -4997,6 +5045,7 @@ ufunc_dealloc(PyUFuncObject *ufunc)
 {
     PyArray_free(ufunc->core_num_dims);
     PyArray_free(ufunc->core_dim_ixs);
+    PyArray_free(ufunc->core_dim_szs);
     PyArray_free(ufunc->core_offsets);
     PyArray_free(ufunc->core_signature);
     PyArray_free(ufunc->ptr);