GPU Support through CuPy about autograd HOT 18 OPEN

@mattjj I'm out of grad school now! 😄

What more needs to be done to support CuPy? Could we chat about this sometime?

from autograd.

@matjj I started a branch on my fork! https://github.com/ericmjl/autograd/tree/cupy

Still familiarizing myself with both CuPy and autograd's codebase. However, one small win I have now is that I can import autograd.cupy as cp and no errors show up. The next thing I need help with is error messages that show up when I try to instantiate an array.

cc: @samuela, @MInner, @iaroslav-ai, @oeway, @kastnerkyle, @outlace, @tscohen I definitely can't do this on my own, and I'd love to get your input on the branch, especially w.r.t. testing and finding problems that show up!

from autograd.

@JarvisDevon, if you're interested in what I've done thus far, you can look at https://github.com/ericmjl/autograd-cupy.

Otherwise, I'd suggest going to jax: it provides autodiff on the NumPy API with JIT compilation to GPU/TPU!

from autograd.

Any updates? I and a few others are very interested, and @tscohen 's comment makes me think it could be pretty straightforward.

from autograd.

Hey Taco! Thanks for the suggestion. We had noticed CuPy, and I think a student is currently looking into the necessary adaptations. This github issue is probably a good place to track any progress on that front. If you have any further thoughts on getting autograd+CuPy working, let us know!

from autograd.

Please support OpenCL if you're planning to add GPU support.

from autograd.

No updates, but pull requests are welcome :)

from autograd.

+1
looking forward to this feature.

from autograd.

+1

from autograd.

Did anybody try using numpy linked with nvblas - a drop-in replacement for BLAS that uses cuBLAS under the hood? If all numpy ops are automatically GPU-accelerated this way, autograd should get accelerated too, right?

from autograd.

It's not too hard to wrap numerical libraries with autograd. The trouble has been that there isn't really a good "GPU-backed numpy" to wrap. That is, the challenge is finding a well-developed GPU-backed numerical library for expressing functions to be differentiated; if we had that, wrapping it in autograd to add autodiff support wouldn't be much work.

As an example, a few weeks ago I made a cupy wrapper in a branch, including a gpu-accelerated MNIST net example. Take a look at the autograd/cupy directory in that branch to see the wrapping code. Basically, it requires three files: cupy_wrapper.py to wrap the cupy functions in primitives, cupy_extra.py to set up a cupy array boxed data type, and cupy_grads.py to define the VJPs. However, cupy is a bit rough around the edges, and I haven't pushed on developing that wrapper further.

I haven't looked at nvblas, but because I haven't heard of it being used, I suspect it is under-developed. I'd like to hear more about it, though!

from autograd.

@mattjj I'm also not personally familiar with nvblas but it appears to be an Nvidia project, so it should be actively maintained.

from autograd.

Any updates guys?

from autograd.

Please check my branch: https://github.com/ericmjl/autograd/tree/cupy

from autograd.

Also, some progress has been made.

from autograd.

Thanks, will check it out.

from autograd.

Any progress on this? The branch at https://github.com/ericmjl/autograd/tree/cupy seems to be gone too?

from autograd.

https://github.com/google/jax

from autograd.