ucl-bug / jwave

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs
GNU Lesser General Public License v3.0
137 stars 21 forks source link

Support `conda` feedstock jax wheels #71

Open astanziola opened 2 years ago

astanziola commented 2 years ago

There are community-driven wheels for jax and jaxlib with GPU support via conda-forge

Relevant links:

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

astanziola commented 2 years ago

@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

djps commented 2 years ago

I tried this but had conflicts between the version of cuda I have and the version it wanted to install