diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h index a24a0d83774f..eb9a3c63d8c9 100644 --- a/numpy/core/include/numpy/ufuncobject.h +++ b/numpy/core/include/numpy/ufuncobject.h @@ -227,6 +227,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/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c index 385d59f88426..b8b38aa1269e 100644 --- a/numpy/core/src/umath/ufunc_object.c +++ b/numpy/core/src/umath/ufunc_object.c @@ -578,6 +578,28 @@ _is_alnum_underscore(char ch) return _is_alpha_underscore(ch) || (ch >= '0' && ch <= '9'); } +/* + * 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 */ @@ -645,10 +667,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; } @@ -677,8 +702,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) { @@ -689,6 +721,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; @@ -733,6 +766,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; @@ -1914,7 +1950,8 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc, * to a core dimension, it won't be visited. */ for (i = 0; i < ufunc->core_num_dim_ix; ++i) { - core_dim_sizes[i] = 1; + npy_intp frozen_size = ufunc->core_dim_szs[i]; + core_dim_sizes[i] = frozen_size == -1 ? 1 : frozen_size; } for (i = 0; i < nop; ++i) { if (op[i] != NULL) { @@ -1949,9 +1986,10 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *ufunc, } for (; idim < num_dims; ++idim) { int core_dim_index = ufunc->core_dim_ixs[dim_offset + idim]; + npy_intp frozen_size = ufunc->core_dim_szs[core_dim_index]; npy_intp op_dim_size = PyArray_SHAPE(op[i])[core_start_dim + idim]; - if (core_dim_sizes[core_dim_index] == 1) { + if (frozen_size == -1 && core_dim_sizes[core_dim_index] == 1) { core_dim_sizes[core_dim_index] = op_dim_size; } else if ((i >= nin || op_dim_size != 1) && core_dim_sizes[core_dim_index] != op_dim_size) { @@ -4370,6 +4408,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); @@ -4716,6 +4755,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); diff --git a/numpy/core/src/umath/umath_tests.c.src b/numpy/core/src/umath/umath_tests.c.src index 46d06288d37c..6eea812471f8 100644 --- a/numpy/core/src/umath/umath_tests.c.src +++ b/numpy/core/src/umath/umath_tests.c.src @@ -177,6 +177,43 @@ 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 di = dimensions[0]; + npy_intp i; + 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**/ + /* The following lines were generated using a slightly modified version of code_generators/generate_umath.py and adding these lines to defdict: @@ -207,6 +244,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 void addUfuncs(PyObject *dictionary) { @@ -234,6 +274,14 @@ addUfuncs(PyObject *dictionary) { 0, matrix_multiply_signature); PyDict_SetItemString(dictionary, "matrix_multiply", 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); + PyDict_SetItemString(dictionary, "cross1d", f); + Py_DECREF(f); + } /* diff --git a/numpy/core/src/umath/umathmodule.c b/numpy/core/src/umath/umathmodule.c index 57b2bb2397d9..6220ad934a84 100644 --- a/numpy/core/src/umath/umathmodule.c +++ b/numpy/core/src/umath/umathmodule.c @@ -121,6 +121,7 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS self->core_num_dim_ix = 0; self->core_num_dims = NULL; self->core_dim_ixs = NULL; + self->core_dim_szs = NULL; self->core_offsets = NULL; self->core_signature = NULL; self->op_flags = PyArray_malloc(sizeof(npy_uint32)*self->nargs);