Repository navigation
Allow libraries that implement __array_ufunc__ to override DUFunc.__c… - #8995
Conversation
|
@jpivarski, Thank you for the effort. This PR is definitely in the right direction.
Yes, that would do. A minimal class to exercise the new code path in Dufunc would be sufficient.
I think we should keep this PR small and focus on In summary, I think all that's remaining for this PR are tests and a doc update. |
|
I added a test (we'll see what happens in CI), though I'm not sure what to say where for a doc update. I never saw any documentation that says that ufuncs made by Thanks! |
|
For docs, I think it's worth adding a note in https://github.com/numba/numba/blob/main/docs/source/reference/jit-compilation.rst#vectorized-functions-ufuncs-and-dufuncs to say what level of NEP13 is supported. |
|
Okay, I think that's it: the assertions are unittest-style and I've added a paragraph at the end of the jit-compilation.rst reference. |
|
Is this waiting on anything from me? As far as I'm aware, it's done. |
sklam
added
4 - Waiting on reviewer
sklam
added
5 - Ready to merge
ianna
left a comment
ianna
left a comment
There was a problem hiding this comment.
@jpivarski - I think the issue scikit-hep/awkward#3179 should be solved in Numba by adding similar functionality.
| return output | ||
| else: | ||
| return super().__call__(*args, **kws) | ||
|
|
There was a problem hiding this comment.
@jpivarski - this PR implemented:
>>> import numba as nb
>>> @nb.vectorize()
... def _square(x):
... return x * x
...
>>> _square
<numba._DUFunc '_square'>
we need to add similar functionality to a CUDAUFuncDispatcher that handles ufuncs on CUDA to allow the following:
>>> @nb.vectorize(
...,
... )
... def _square_cuda(x):
... return x * x
...
>>> _square_cuda
<numba.cuda.vectorizers.CUDAUFuncDispatcher object at 0x7259f82fa790>
|
Maybe this function: numba/numba/cuda/deviceufunc.py Lines 241 to 329 in 99f5cc6 is missing @gmarkall, I didn't see anything like that in the issues or PRs. Should @ianna open an issue? (Should it go to this codebase or the new one, in which Numba-CUDA is an external project?) |
This PR addresses a long-standing issue in which array-like objects that define an
__array_ufunc__method (NEP-18) can't be used with ufuncs created bynp.vectorizeunless a signature is provided. (I should have done this years ago!)For instance, in Awkward Array,
but
and even if we pass
nopython=True,Similarly, with the Sparse library,
but
and
However, with this PR, ufuncs created with
@nb.vectorizeand no type signature are recognized by Awkward Array:and Sparse:
It's because
__array_ufunc__methods are usually implemented by pre- and post-processing the result of applying the ufunc to some underlying NumPy array. If an__array_ufunc__method is implemented some other way, without Numba-compatible objects at any level, then@nb.vectorize-decorated functions wouldn't work with it, anyway.This is a draft PR because:
Some thought will be needed to know what kinds of tests would be appropriate. This mostly affects libraries that use Numba. Maybe a little class that implements
__array_ufunc__?It could be related to #6678, #4625, and/or #2979. For completeness, the PR in its current state only addresses ufunc's
__call__method; perhaps we wantaccumulate,at,outer,reduce, andreduceatas well. Also, it only addressesnp.vectorize; maybe something similar is needed fornp.guvectorize.I'm leaving it in this simple state so that I can get feedback about whether this is the right direction before proceeding to generalize it. (Also, which generalizations are important?)