-
-
Notifications
You must be signed in to change notification settings - Fork 12k
Description
I recently sat down with @ericmjl to see if we could make autograd work with cupy by making them both understand protocols like __array_ufunc__
The goal here being to improve all projects to produce and consume numpy protocols rather than explicitly import each other in pair-wise plugin mechanisms.
One issue that we ran into is that autograd needs to produce new arrays, like ones (the gradient of sum) and random (not sure why yet). It would be nice to be able to say "produce an array of ones with a particular shape and dtype, but using the module that created this particular array object". This would allow autograd to produce dask arrays of ones, cupy arrays of ones, etc..
The numpy.ones_like function almost does this except for the following two issues:
- It takes the shape from the given array (which makes perfect sense given its original objective)
- There is no protocol for it, so
np.ones_like(my_duck_array, ...)produces a numpy array
Also, to be clear, ones here is an example of a larger problem of how to improve dispatch for a wider set of numpy functions.