-
-
Notifications
You must be signed in to change notification settings - Fork 11.1k
WIP: ENH: Expand gufunc signature to allow flexible dimension specs #11132
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
e9a5ab5
cbbe0f6
dbbeb58
18505d7
67e752e
50d59fd
d479a5b
f3d1622
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -339,6 +339,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 | ||
*/ | ||
|
@@ -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]) && | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think here and below it might make more sense to separate out more:
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed by removing frozen size code |
||
(frozen_size = _get_size(signature + i)) < 0) { | ||
parse_error = "expect dimension name or frozen size"; | ||
goto fail; | ||
} | ||
while (j < ufunc->core_num_dim_ix) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This
(Could also be a Though perhaps it is better to just split the frozen and normal paths altogether (see above) There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed by removing code |
||
|
@@ -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; | ||
|
@@ -2087,10 +2123,14 @@ _get_coredim_sizes(PyUFuncObject *ufunc, PyArrayObject **op, | |
*/ | ||
for (idim = 0; 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_DIM( | ||
op[i], REMAP_AXIS(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 (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]) { | ||
|
@@ -2364,6 +2404,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); | ||
|
@@ -4671,6 +4722,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; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should also initialize There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed |
||
if (signature != NULL) { | ||
if (_parse_signature(ufunc, signature) != 0) { | ||
Py_DECREF(ufunc); | ||
|
@@ -5017,6 +5069,7 @@ ufunc_dealloc(PyUFuncObject *ufunc) | |
{ | ||
PyArray_free(ufunc->core_num_dims); | ||
PyArray_free(ufunc->core_dim_ixs); | ||
PyArray_free(ufunc->core_dim_szs); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should also free There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fixed |
||
PyArray_free(ufunc->core_offsets); | ||
PyArray_free(ufunc->core_signature); | ||
PyArray_free(ufunc->ptr); | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
next three
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
removing szs makes it two again, so "fixed"