Support `conda` feedstock jax wheels
astanziola opened this issue · comments
Antonio Stanziola commented
There are community-driven wheels for jax
and jaxlib
with GPU support via conda-forge
Relevant links:
- https://github.com/conda-forge/jax-feedstock
- https://github.com/conda-forge/jaxlib-feedstock
- google/jax#189
I managed to get a working installation of jax
on with only conda
installed (no cuda or cudnn prior installation) using
$ conda create -n condajax python=3.9
$ conda activate condajax
$ conda install jaxlib jax --channel conda-forge
It is interesting to see if this works under windows, to avoid using the WSL or building jax
from scratch
Antonio Stanziola commented
@djps FYI, not sure if you want to give it a try.
I don't have a Windows with CUDA machine at hand now, but will test this out as soon as I can
David Sinden commented
I tried this but had conflicts between the version of cuda I have and the version it wanted to install