Jaxlib cuda releases only for Cuda 10.2 and cuda 11.2

[yashk2810]

Hello everyone,

I have been thinking about releasing GPU jaxlib wheels only for cuda 10.2 and 11.2. This makes it in line with Pytorch and TF.

What does everyone think?

[jakevdp]

I believe Colab currently uses cuda110 wheels for jaxlib, so a first order of business before making this change would be to ensure Colab is compatible with updated versions.

Edit: it appears that Colab also uses pytorch's cuda 10.2 wheels, so it may be as easy as updating the jaxlib wheels in the Colab container, and not having to touch cuda (which is a more involved change)

[yashk2810]

Is there a test I can run by installing the current 11.2 jaxlib in colab and making sure jax runs?

AFAIK, cuda is backwards compatible for 11.x versions.

[jakevdp]

I just poked around the Colab container a bit... the problem there is that there is no cupy-cuda111 release that supports Python 2, and at least for the time being Python 2 seems to be part of the Colab container image (though I believe its use is deprecated). I suspect maintaining their planned deprecation schedule will be higher priority for the Colab team than having the newest version of JAX, so the effect of this change may be that jaxlib will be frozen at version 0.1.70 in Colab for the foreseeable future. I'll try a container build with cuda111 to see whether I can find a way around this.

[yashk2810]

I'll try a container build with cuda111 to see whether I can find a way around this.

Thank you!

The freeze sounds good to me.

Or we could also build a custom wheel for colab for 11.0.

[jakevdp]

In my opinion we should do everything possible to avoid a situation where we are unable update jax in Colab – it's an important platform for JAX users & the main way we let people interact with our documentation.

[yashk2810]

Sure, we can do that (but not sustainably). Colab will need to update its Cuda version at some point.

We can build a wheel for colab only for 11.0 using the current pipeline (not the new one). The command in colab containers can be updated to point to the GCS directory right?

[yashk2810]

But I am also of the opinion that we should push colab to update to Cuda 11.1+ (maybe its in their plan?)

[yashk2810]

Now the official release for jaxlibs will be for Cuda 10.2 and Cuda 11.1 (similar to Pytorch)