gpu - avoid cpu/gpu double implementations of same code

[lastephey]

In a perfect world it be ideal to avoid

if cpu:
    do something
if gpu:
    do something else

and instead have a nicer way (maybe a decorator) of making this more portable.

[dmargala]

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>

[lastephey]

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.

[lastephey]

Thanks for your help. Closing in the spirit of Closember.