Skip to content

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

Merged
merged 8 commits into from
Jul 22, 2024

Conversation

WarrenWeckesser
Copy link
Member

@WarrenWeckesser WarrenWeckesser commented Jul 11, 2024

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 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,

In [1]: import numpy as np

In [2]: from experiment import conv1d_full, euclidean_pdist

In [3]: type(conv1d_full), type(euclidean_pdist)
Out[3]: (numpy.ufunc, numpy.ufunc)

In [4]: x = np.array([1, 2, 3.5])

In [5]: y = np.array([-1, 0, 4, 1, 0.5, 2])

conv1d_full is a gufunkified version of np.convolve with mode='full':

In [6]: conv1d_full(x, y)
Out[6]: array([-1.  , -2.  ,  0.5 ,  9.  , 16.5 ,  6.5 ,  5.75,  7.  ])

In [7]: np.convolve(x, y, mode='full')
Out[7]: array([-1.  , -2.  ,  0.5 ,  9.  , 16.5 ,  6.5 ,  5.75,  7.  ])

Basic broadcasting:

In [8]: conv1d_full([x, -x, 2*x + 1], y)
Out[8]: 
array([[ -1.  ,  -2.  ,   0.5 ,   9.  ,  16.5 ,   6.5 ,   5.75,   7.  ],
       [  1.  ,   2.  ,  -0.5 ,  -9.  , -16.5 ,  -6.5 ,  -5.75,  -7.  ],
       [ -3.  ,  -5.  ,   4.  ,  23.  ,  38.5 ,  16.5 ,  14.  ,  16.  ]])

For an array of points with shape (m, n), euclidean_pdist computes the m*(m-1)/2 pairwise distances. If out is not provided, it creates an output with the correct shape:

In [9]: pts = np.array([[0, 0, 0], [0, 0, 1], [1, 2, 3], [-1, 0, 2], [2, 2, 2]])

In [10]: euclidean_pdist(pts)
Out[10]: 
array([1.        , 3.74165739, 2.23606798, 3.46410162, 3.        ,
       1.41421356, 3.        , 3.        , 1.41421356, 3.60555128])

If out is given but has the wrong shape, an exception is raised:

In [11]: out = np.zeros(5)

In [12]: euclidean_pdist(pts, out=out)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[12], line 1
----> 1 euclidean_pdist(pts, out=out)

ValueError: euclidean_pdist: the core dimension p of the out parameter does not equal m*(m - 1)/2, where m is the first core dimension of the input x; got m=5, so p must be 10, but got p=5).

This implementation of euclidean_pdist requires at least one point:

In [13]: z = np.zeros((0, 4))

In [14]: euclidean_pdist(z)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[14], line 1
----> 1 euclidean_pdist(z)

ValueError: euclidean_pdist: the core dimension m of the input parameter x must be at least 1; got 0.

@WarrenWeckesser WarrenWeckesser added 01 - Enhancement component: numpy._core component: numpy.ufunc 63 - C API Changes or additions to the C API. Mailing list should usually be notified. labels Jul 11, 2024
@WarrenWeckesser WarrenWeckesser marked this pull request as draft July 11, 2024 01:07
@WarrenWeckesser
Copy link
Member Author

WarrenWeckesser commented Jul 11, 2024

Marked as draft. I still need to ensure that gufuncs with possible missing dimensions can be handled correctly. Flexible dimensions work. The PR could use a few a more tests to cover such things, so still marked as draft.

@WarrenWeckesser
Copy link
Member Author

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 ufuncobject.h and 7 lines in ufunc_object.c. Everything else is there to add a couple ufuncs to _umath_tests.c.src that use the new feature, with corresponding tests added to test_ufunc.py.

@WarrenWeckesser
Copy link
Member Author

We could use this feature to coalesce some of the redundant linalg gufuncs. For example, we have:

gufunc   signature    use when
------   ----------   --------
svd_m    (m,n)->(m)   n >= m
svd_n    (m,n)->(n)   n <= m

We can combine these into one gufunc with effective signature (m,n)->(min(m,n)). In the actual implementation, the gufunc signature would be (m,n)->(p), and in the process_core_dims_func function, the p dimension would be set to min(m, n).

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.
@WarrenWeckesser WarrenWeckesser force-pushed the user-process-core-dims branch from 84fac1f to a3f61e1 Compare July 12, 2024 19:27
Copy link
Member

@seberg seberg left a 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;
}
}

Copy link
Member

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.

Copy link
Member Author

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.

Copy link
Member

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.

Copy link
Contributor

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.

Copy link
Member Author

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!

@WarrenWeckesser
Copy link
Member Author

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 process_core_dims_func in the ufunc docstring.

I'll simplify the tests; I wanted to exercise the feature with nontrivial examples, but I might have gone a little overboard.

@WarrenWeckesser
Copy link
Member Author

WarrenWeckesser commented Jul 13, 2024

By the way, this feature allows cross to be implemented as a gufunc that handles both the 2-d and 3-d cases. The differences between this and np.cross are that crossing a 2-d with a 3-d is not allowed (at least in my example implementation, because that never should have been allowed in np.cross!), and, because a core dim signature of (p) gives an array of length 0 when p is 0, the shape signature in the 2-d case is (2),(2)->(1), so crossing two length-2 arrays gives a length 1 array, not a scalar.

E.g.

In [3]: from experiment import cross

In [4]: cross
Out[4]: <ufunc 'cross'>

In [5]: cross.signature
Out[5]: '(m),(m)->(p)'

In [6]: cross([1.5, 2, -1], [2.5, 0, 1])  # 3-d
Out[6]: array([ 2., -4., -5.])

In [7]: cross([8.5, 10], [-1.0, 2.0])  # 2-d
Out[7]: array([27.])

In [8]: cross([[1.5, 0], [2, 3], [8.5, 10]], [-1.0, 2.0])  # Broadcasting 2-d case
Out[8]: 
array([[ 3.],
       [ 7.],
       [27.]])

In [9]: cross([1, 2, 3, 4], [2, 5, -1, 2])  # Invalid shape
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], line 1
----> 1 cross([1, 2, 3, 4], [2, 5, -1, 2])

ValueError: cross: core dimension of the input must be 2 or 3; got 4.

I'm not saying we should do this in numpy, but it does demonstrate the flexibility that this feature provides.

Copy link
Contributor

@mhvk mhvk left a 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:

  1. Definitely needs to have some additional text for the construct-your-own-ufunc section; referring to _umath_tests.c.src seems only fair!
  2. Should this extension become an addendum to NEP 20? In principle, the convoluted matmul signature could now be removed...
  3. 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;
}
}

Copy link
Contributor

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.

@seberg
Copy link
Member

seberg commented Jul 15, 2024

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!

one could still use the ? for that perhaps.

Am I missing things: I think the ? syntax is still very much required and thus the matmul signature cannot be replaced by this (which I think is fine)?
To some degree, I think the best thing would be to allow any token that isn't currently defined to just be the name of that dimension.

For example, rather than (m),(n)->(p), write:

(m),(n)->(m + n - 3) 

And m + n - 3 could just be considered the name of that dimension for the sake of the first processing step.
To not interfere with other syntax, it could also be `m + n - 3` (with e.g. backticks or similar to indicate this is not a single character name and thus typically a "functional dependency on other axes").

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.)

@mhvk
Copy link
Contributor

mhvk commented Jul 15, 2024

Am I missing things: I think the ? syntax is still very much required and thus the matmul signature cannot be replaced by this (which I think is fine)?

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.

To some degree, I think the best thing would be to allow any token that isn't currently defined to just be the name of that dimension.

I like that idea (even without the backticks, though that may make parsing easier)!

@WarrenWeckesser
Copy link
Member Author

WarrenWeckesser commented Jul 15, 2024

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 m and n can't both be 0, so it will test the two motivating use-cases: adding constraints to input sizes and computing output sizes.

@WarrenWeckesser
Copy link
Member Author

WarrenWeckesser commented Jul 15, 2024

@mvhk wrote

Another item that can be implemented with this is short-circuiting all_equal - a long-standing issue (#8513).

I don't see how this helps with all_equal. You can already implement all_equal as a gufunc with signature (n)->(). That's what I did in ufunclab, although I called it all_same for some reason.

Argh, nevermind, I didn't read the issue well enough. The all_equal that is discussed has shape signature (n),(n)->(). And the question in #8513 is, can it be a gufunc that broadcasts its core dimensions?

So now I agree that, yes, as implemented here, the process_core_dims_func can be used to implement all_equal, and have it broadcast. The shape signature would be (m),(n)->(), and the hook function would check that one of three conditions holds: m == n, m == 1, or n == 1 (i.e. regular broadcast rules). The one quirk is that the length-1 argument can't be a scalar; it must have a core dimension with length 1. So to compare a vector x to the value 0, you would write all_equal(x, [0]). Fixed below...

@WarrenWeckesser
Copy link
Member Author

The one quirk is that the length-1 argument can't be a scalar; it must have a core dimension with length 1. So to compare a vector x to the value 0, you would write all_equal(x, [0]).

But if the signature is (m?),(n?)->(), scalars work. Here it is in action:

In [1]: from experiment import all_equal

In [2]: all_equal([2, 2, 2], [2, 3, 4])
Out[2]: np.False_

In [3]: all_equal([2, 2, 2], [2, 2, 2])
Out[3]: np.True_

In [4]: all_equal([2, 2, 2], [3])
Out[4]: np.False_

In [5]: all_equal([2, 2, 2], [2])
Out[5]: np.True_

In [6]: all_equal([2, 2, 2], 2)
Out[6]: np.True_

In [7]: all_equal([2, 2, 2], [[2],[3]])
Out[7]: array([ True, False])

In [8]: all_equal([[2, 2, 2, 2], [6, 7, 8, 9]], [[[2]],[[3]],[[4]]])
Out[8]: 
array([[ True, False],
       [False, False],
       [False, False]])

@WarrenWeckesser
Copy link
Member Author

WarrenWeckesser commented Jul 16, 2024

I updated the tests (now just the one test gufunc conv1d_full is used), added comments about PyUFunc_ProcessCoreDimsFunc in ufuncfobject.h, and updated generalized-ufuncs.rst. (I didn't run the full CI on the last commit, since it only tweaked the .rst file.)

I'm still thinking about the ideas for handling the signature. Some cases to consider:

  • A gufunc that implements cross to match the current np.cross as closely as possible (including--the horror!--the ability to cross a 2-d vector with a 3-d vector) with input core dimensions m and n. In Python syntax, the expression for the output core dimension would be something like 3 if m == 3 or n == 3 else 1. (The hook function also imposes the constraint that m and n can only be 2 or 3.)
  • A gufunc that doubles the dimensions of a square 2-d array. The effective shape signature is (m, m)->(2*m, 2*m). Currently this would be written (m, m)->(p, p) and the hook function would ensure that p is 2*m. This just shows that if we say that an output core dimension token can be an arbitrary string, we still have to check for equal tokens to ensure that only one output core dimension is defined.

By the way, I have two follow-up PRs planned:

I can add either of those to this PR if is preferable.

@seberg
Copy link
Member

seberg commented Jul 16, 2024

we still have to check for equal tokens to ensure that only one output core dimension is defined.

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.
You could do p=... so that you can re-use the p (and use its alphabetical sorting). But, I am not really convinced that duplicating the same (non-trivial) entry happens enough to worry about it. Plus, you can still write p=... and p, you would just get it passed twice ;).

I can add either of those to this PR if is preferable.

I don't think it is necessary, we know it has it's uses, and it is a very simple API addition after all.

Copy link
Contributor

@mhvk mhvk left a 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;
}
Copy link
Contributor

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
Copy link
Member

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.

@seberg
Copy link
Member

seberg commented Jul 22, 2024

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!

@WarrenWeckesser
Copy link
Member Author

Thanks @mhvk and @seberg! (I've been off the grid a lot for the past week because of, believe it or not, tornadoes in upstate NY last week. Power was out for a few days (hooray for backup generators!), and I'm still waiting for internet service to be restored at home.)

@seberg seberg merged commit db93110 into numpy:main Jul 22, 2024
68 checks passed
@WarrenWeckesser WarrenWeckesser deleted the user-process-core-dims branch July 23, 2024 16:26
WarrenWeckesser added a commit to WarrenWeckesser/numpy that referenced this pull request Jul 23, 2024
[skip actions] [skip azp] [skip cirrus]
@WarrenWeckesser
Copy link
Member Author

Release note: #27017

mattip added a commit that referenced this pull request Jul 23, 2024
WarrenWeckesser added a commit to WarrenWeckesser/numpy that referenced this pull request Jul 27, 2024
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
ArvidJB pushed a commit to ArvidJB/numpy that referenced this pull request Nov 1, 2024
)

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>
ArvidJB pushed a commit to ArvidJB/numpy that referenced this pull request Nov 1, 2024
[skip actions] [skip azp] [skip cirrus]
ArvidJB pushed a commit to ArvidJB/numpy that referenced this pull request Nov 1, 2024
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
01 - Enhancement 63 - C API Changes or additions to the C API. Mailing list should usually be notified. component: numpy._core component: numpy.ufunc
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants