-
-
Notifications
You must be signed in to change notification settings - Fork 10.8k
ENH: Provide a hook for gufuncs to process core dimensions. #26908
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
Conversation
Marked as draft. |
1a80540
to
3eb1f94
Compare
The two failing CI jobs are also failing in other PRs and are not related to this change. The significant changes here are just 7 lines in |
We could use this feature to coalesce some of the redundant linalg gufuncs. For example, we have:
We can combine these into one gufunc with effective signature |
The field `process_core_dims_func` is added to the PyUFuncObject. This is a pointer to a function typedef'd as typedef int (PyUFunc_ProcessCoreDimsFunc)( struct _tagPyUFuncObject *ufunc, npy_intp *core_dim_sizes); The author of a gufunc can set the field with a function that they implement. The function will be called when the gufunc is called. (The actual call is in the internal function _get_coredim_sizes.) The user-defined function can set an exception and return an error status if any of the core dimensions in `core_dim_sizes` do not satisfy the assumptions of the gufunc. The user-defined function can also *set* the value of core dimensions that are passed in as -1, meaning the correspond out parameter was not given. This allows calculations such pairwise distances (which generates m*(m-1)/2 output values for an input with shape (m, n)) and full convolution (generates m + n - 1 output values from two inputs with shapes m and n) to be implemented as gufuncs with automatic allocation of the output with the correct shape. The output shape is computed and set in the user-defined function.
84fac1f
to
a3f61e1
Compare
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.
I like this addition and I think we should just do it, still relying on the named axes seems fine and not really limiting (as discussed in emails also).
This is also naturally defined on the ufunc (no DType specific component).
The Python exposed .signature
is not ideal after this. As discussed, one can hack around that, but that is awkward outside of NumPy (you cannot free the signature, or you encode private NumPy ABI for one thing).
But I am happy to not worry about that here, I think it might be best to solve by giving ufuncs a __dict__
or a .signature
property; if you like allow mutating only as _set_signature()
).
I do find the tests a bit much, TBH. I can be convinced to ignore that but they almost entice to use them directly from the umath tests module :).
After all, all we need is a do-nothing loop with and with a process-core-dims function and some tests that check the output shape (and error when out=
is passed).
Otherwise, it feels almost like we should just make convolve/correlate a gufunc here :)
return -1; | ||
} | ||
} | ||
|
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.
So one thing we do not do here is sanity check the output. I.e. if the core_dim_size
is already set but the user mutates it, we could be in trouble.
To be clear, I think that is OK, but it should be documented.
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.
It is easy enough to pass a copy to process_core_dims_func
and verify on return that nothing was changed that shouldn't be changed. I'll add that.
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.
I don't think it matters, just occured to me that we should document that you may only validate (raise an error) or fill in -1
.
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.
I wavered a bit about whether or not to test, but in the end, making a mistake in the processing here is no different from making a mistake in the actual function that gets called - one can get a core dump regardless. So, agreed with @seberg that a check after the call is not needed.
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.
OK, I won't add the check!
I'm fine with not hacking the signature. In the API docs (not yet written), I'll strongly recommend that gufunc authors document the formula for any shapes set in I'll simplify the tests; I wanted to exercise the feature with nontrivial examples, but I might have gone a little overboard. |
By the way, this feature allows E.g.
I'm not saying we should do this in numpy, but it does demonstrate the flexibility that this feature provides. |
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.
I like this very much. I'm also all for the extra added tests -- when I first started to do anything with ufuncs, the test suite was very helpful to understand how things work, and I think this adds to that.
Another item that can be implemented with this is short-circuiting all_equal
- a long-standing issue (#8513). Indeed, I think this nicely flexible approach allows one to implement most of the things noted in #8811 as well - one could implement the broadcasting option by simply assigning a new variable and checking in the function whether that is either equal to another or 1.
Some larger suggestions regarding how to document this:
- Definitely needs to have some additional text for the construct-your-own-ufunc section; referring to
_umath_tests.c.src
seems only fair! - Should this extension become an addendum to NEP 20? In principle, the convoluted
matmul
signature could now be removed... - For the signature, one option would be to allow placing a comment at the end, say something like
# p = m + n - 1
. Alternatively, one might have a single character that indicates that a parameter is unusual - one could still use the?
for that perhaps.
Note that one could do items 2 and 3 could be done later - don't want to stop something nice now!
return -1; | ||
} | ||
} | ||
|
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.
I wavered a bit about whether or not to test, but in the end, making a mistake in the processing here is no different from making a mistake in the actual function that gets called - one can get a core dump regardless. So, agreed with @seberg that a check after the call is not needed.
Nice thought on adding a note to NEP 20. To some degree, I think it is independent, but NEP 20 is the main document you find on gufunc signatures, and if you are interested in that, there is a good chance you want to know about this flexibility as well. Plus, if we allow things like the backtick idea below, the signature parsing is actually modified!
Am I missing things: I think the For example, rather than
And We wouldn't actually automatically evaluate the formula, the user has to ensure that the formula is indeed correct. (At least I think I prefer that, we just consider it "documentation" for which it would be nice if it can be evaluated.) |
Hmm, I thought one could do it with a check on the numbers, but you are right that that wouldn't work for 1-D arrays.
I like that idea (even without the backticks, though that may make parsing easier)! |
Re: tests: how about I remove the "wl1_pdist" gufunc and tests, but keep "conv1d_full"? The tests for wl1_pdist don't add anything new. The signature has flexible dimensions, but that doesn't make any difference in the hook function as far as testing goes. I'm adding a check to conv1d_full that will require that |
@mvhk wrote
Argh, nevermind, I didn't read the issue well enough. The So now I agree that, yes, as implemented here, the |
But if the signature is
|
[skip actions] [skip azp] [skip cirrus]
I updated the tests (now just the one test gufunc I'm still thinking about the ideas for handling the signature. Some cases to consider:
By the way, I have two follow-up PRs planned:
I can add either of those to this PR if is preferable. |
Yes, I don't think it matters much. I think the parsing is currently in C, but we can just call back into Python so that all the C-side has to do is parse a tuple of ints.
I don't think it is necessary, we know it has it's uses, and it is a very simple API addition after all. |
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.
I think this looks great now! We can worry about niceties later.
return -1; | ||
} | ||
return 0; | ||
} |
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.
What a really nice & clear addition to the documentation!
@@ -191,6 +224,12 @@ typedef struct _tagPyUFuncObject { | |||
/* A PyListObject of `(tuple of DTypes, ArrayMethod/Promoter)` */ | |||
PyObject *_loops; | |||
#endif | |||
#if NPY_FEATURE_VERSION >= NPY_2_1_API_VERSION |
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.
Just noticed, but just noting to not forget: NPY_2_1_API_VERSION
isn't defined yet I think. But don't worry about it, I can just do it (maybe even later today) since evertying else should be good.
Pushed a commit to introduce 2.1 as a specific C-API version (hash is unchanged, because we don't hash the headers where the change is). Looked through, and this looks all good to me, so planning to merge it later today if nobody comments, thanks! |
[skip actions] [skip azp] [skip cirrus]
Release note: #27017 |
DOC: Release note for feature added in gh-26908.
The new feature added in numpygh-26908 is used to reduce the number of gufuncs in np.linalg._umath_linalg. For example, before this change, the gufuncs `svd_m` and `svd_n` used the same code, but two distinct gufuncs were necessary because, for input with shape (m, n), the output has shape min(m, n). `svd_m` returned shape (m,), and `svd_n` returned shape (n,). The wrapper code had to determine which gufunc to call. This can be handled now with a single gufunc by defining a function that sets the output dimension to min(m, n). In this PR, the following pairs of gufuncs are coalesced into a one gufunc: lstsq_m & lstsq_n => lstsq qr_r_raw_m & qr_r_raw_n => qr_r_raw svd_m & svd_n => svd svd_m_f & svd_n_f => svd_f svd_m_s & svd_n_s => svd_s
) The field `process_core_dims_func` is added to the PyUFuncObject. This is a pointer to a function typedef'd as typedef int (PyUFunc_ProcessCoreDimsFunc)( struct _tagPyUFuncObject *ufunc, npy_intp *core_dim_sizes); The author of a gufunc can set the field with a function that they implement. The function will be called when the gufunc is called. (The actual call is in the internal function _get_coredim_sizes.) The user-defined function can set an exception and return an error status if any of the core dimensions in `core_dim_sizes` do not satisfy the assumptions of the gufunc. The user-defined function can also *set* the value of core dimensions that are passed in as -1, meaning the correspond out parameter was not given. This allows calculations such pairwise distances (which generates m*(m-1)/2 output values for an input with shape (m, n)) and full convolution (generates m + n - 1 output values from two inputs with shapes m and n) to be implemented as gufuncs with automatic allocation of the output with the correct shape. The output shape is computed and set in the user-defined function. * MAINT: Update 2.1 C-API version and use it --------- Co-authored-by: Sebastian Berg <sebastianb@nvidia.com>
[skip actions] [skip azp] [skip cirrus]
The new feature added in numpygh-26908 is used to reduce the number of gufuncs in np.linalg._umath_linalg. For example, before this change, the gufuncs `svd_m` and `svd_n` used the same code, but two distinct gufuncs were necessary because, for input with shape (m, n), the output has shape min(m, n). `svd_m` returned shape (m,), and `svd_n` returned shape (n,). The wrapper code had to determine which gufunc to call. This can be handled now with a single gufunc by defining a function that sets the output dimension to min(m, n). In this PR, the following pairs of gufuncs are coalesced into a one gufunc: lstsq_m & lstsq_n => lstsq qr_r_raw_m & qr_r_raw_n => qr_r_raw svd_m & svd_n => svd svd_m_f & svd_n_f => svd_f svd_m_s & svd_n_s => svd_s
The field
process_core_dims_func
is added to the PyUFuncObject. This is a pointer to a function typedef'd asThe author of a gufunc can set the field with a function that they implement. The function will be called when the gufunc is called. (The actual call is in the internal function _get_coredim_sizes.) The user-defined function can set an exception and return an error status if any of the core dimensions in
core_dim_sizes
do not satisfy the assumptions of the gufunc. The user-defined function can also set the value of core dimensions that are passed in as -1, meaning the corresponding out parameter was not given. This allows calculations such pairwise distances (which generates m*(m-1)/2 output values for an input with shape (m, n)) and full convolution (generates m + n - 1 output values from two inputs with shapes m and n) to be implemented as gufuncs with automatic allocation of the output with the correct shape. The output shape is computed and set in the user-defined function.The two examples mentioned above are implemented at https://github.com/WarrenWeckesser/experiments/tree/main/python/numpy/gufunc-process_core_dims_func. For example,
conv1d_full
is a gufunkified version ofnp.convolve
withmode='full'
:Basic broadcasting:
For an array of points with shape (m, n),
euclidean_pdist
computes the m*(m-1)/2 pairwise distances. Ifout
is not provided, it creates an output with the correct shape:If
out
is given but has the wrong shape, an exception is raised:This implementation of
euclidean_pdist
requires at least one point: