Closed lastephey closed 1 year ago
It looks like there is support for using the numpy api directly on cupy objects which is achieved through the __array_function__
protocol: https://numpy.org/neps/nep-0018-array-function-protocol.html
For example, consider the following code:
import cupy as cp
import numpy as np
x = cp.random.normal(size=(1000, 1000))
w, v = np.linalg.eigh(x)
print(w.device)
Outputs:
<CUDA Device 0>
I like that this could help us prevent repeated code. I dislike that it makes it harder to read/debug the code and immediately see if we're on the cpu or gpu, but that's maybe unavoidable.
In a perfect world it be ideal to avoid
and instead have a nicer way (maybe a decorator) of making this more portable.