Point to the conda-forge jaxlib wheels in the JAX readme.

[hawkinsp]

Fixes #189

[hawkinsp]

@ngam @ericmjl @xhochy

Any comments?

[ngam]

Any comments?

Yes. Two comments:

1. I am more than happy to help but I ran into a super challenging issue that I've been struggling to resolve (we are yet to release jaxlib 0.3.15 on our end; I am waiting for conda-specific global updates to filter through and I will give it another go soon). This is all to say, I would add a stronger disclaimer so people don't spam this repo with issues...
2. I believe we have really solid cuda builds on conda-forge now (and we work quite hard to make sure that our jax builds can easily coexist with tensorflow). I would preemptively add instructions for cuda installation (that's not covered by the readme; I can potentially fix it, but the bots will delete that in no time :D). Something like this, which is also slowly getting out of date: https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch

I will suggest some edits to address both points in a bit (or can do so in a later PR)

[ngam]

Suggested change

	### Conda installation
	There is a community-supported Conda build of `jaxlib`. See the
	[jaxlib conda-forge](https://github.com/conda-forge/jaxlib-feedstock#installing-jaxlib)
	project for installation instructions.
	### Conda installation
	There is a community-supported Conda builds of `jaxlib` and `jax`, which are driven independently of the JAX team. See the
	[jaxlib](https://github.com/conda-forge/jaxlib-feedstock#installing-jaxlib) and [jax](https://github.com/conda-forge/jax-feedstock#installing-jax)
	repositories for installation instructions. Note that, the conda-forge project redistributes CUDA and CuDNN; as such, there is no need to install them beforehand. If you would like to install the CUDA-enabled builds of `jaxlib`, follow the instructions in the [Tips & tricks](https://conda-forge.org/docs/user/tipsandtricks.html#installing-cuda-enabled-packages-like-tensorflow-and-pytorch) section of the conda-forge website.

[ngam]

(I don't know if this is the correct formatting... it looks weird; feel free to discard now and I can submit a PR later if needed)

[hawkinsp]

I incorporated a number of your suggestions. By the way, I found out that it is actually not sufficient to let conda install CUDA for you. JAX needs ptxas, which isn't distributed. So you must install at least the ptxas binary yourself. Everything else necessary is in the conda builds.

[ngam]

I thought we get ptxas under conda-forge too. Could you give me a quick test to investigate this? There was a discussion about this previously in conda-forge. I will try to locate it and see the updates on that.

[ngam]

Okay, now I remember. I was even part of the discussion:

• ptxas executable missing conda-forge/tensorflow-feedstock#170
• ptxas executable conda-forge/cudatoolkit-feedstock#72

[ngam]

Thanks for catching the ptxas bit! Crucial!

[ngam]

No need for a test as we likely won't even catch it as we do have ptxas in the containers we use for building but not in the distributed cudatoolkit --- in the future, we could discuss building an "optimized" container on top of the conda-forge binaries if you'd like, which will be on top of the nvidia containers

[hawkinsp]

The ptxas issue is fixable with a little bit of engineering work on our side (there's a library version of the same thing: libnvptxcompiler_static.a) or with the agreement of NVidia to allow redistribution of ptxas. I'll raise it with them.

[gmarkall]

Could this PTX Compiler Python binding be useful as an alternative to calling ptxas: https://github.com/rapidsai/ptxcompiler ?

[hawkinsp]

I can do you one better:

https://pypi.org/project/nvidia-cuda-nvcc-cu11/