Schedule guvectorize'd functions over dask.distributed #9495
Conversation
There was a problem hiding this comment.
@crusaderky, thanks for your contribution. Do you have a test case using Dask where current code would fail?
| class NEP13Array: | ||
| """https://numpy.org/neps/nep-0013-ufunc-overrides.html""" | ||
| def __init__(self, array): | ||
| self.array = array | ||
|
|
||
| def __array__(self): | ||
| return self.array | ||
|
|
||
| def tolist(self): | ||
| return self.array.tolist() | ||
|
|
||
| def __array_ufunc__(self, ufunc, method, *args, **kwargs): | ||
| if method != "__call__": | ||
| return NotImplemented | ||
|
|
||
| return NEP13Array(ufunc(*[np.asarray(x) for x in args], **kwargs)) |
There was a problem hiding this comment.
Is this change needed for the PR to work?
There was a problem hiding this comment.
No; I felt that it would make it easier to understand the test code if these two very similar, and yet subtly different stub classes were presented right next to each other. I can revert if you prefer.
import dask.array as da
import distributed
import numba
client = distributed.Client()
@numba.guvectorize(["f8,f8[:]"], "()->()")
def double(x, out):
out[:] = x * 2
double(da.arange(5)).compute()Output: |
|
@crusaderky, can you add a release notes file? https://numba.readthedocs.io/en/latest/developer/contributing.html#release-notes |
guilhermeleobas
added
4 - Waiting on author
Done |
guilhermeleobas
added
5 - Ready to merge
There was a problem hiding this comment.
Thanks, @crusaderky
Sorry for the delay on this last review.
Fix failure to pickle a dask array that embeds a guvectorized ufunc in its graph.