Add Windows WSL2 instructions

[randyau]

Notes for compiling jax w/ CUDA support in Windows WSL environment

[joeltgray]

Any chance of getting one for AMD GPUs using ROCm and Windows WSL?

[Phildo]

followed instructions (including installing jax from source, which takes ~1hr! would recommend adding a warning about that), and am still getting the "warning": WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.), with the frontend failing to connect to the server.

Any ideas what I could be missing?

Also a note: I think jax fixed the "config error" you warned about and provided a workaround for, so I don't think that's necessary any longer.

[matthew-lowe]

followed instructions (including installing jax from source, which takes ~1hr! would recommend adding a warning about that), and am still getting the "warning": WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.), with the frontend failing to connect to the server.

I had the same issue and after lots of fiddling I think I got it to work properly w/ the GPU.

I used Ubuntu 20.04 WSL and CUDA 11.3 (so it would match the version for pytorch exactly). CUDA install was standard as described on the official page (https://developer.nvidia.com/cuda-11.3.0-download-archive).

I downloaded the latest cuDNN version from the archive that matched the CUDA version (https://developer.nvidia.com/rdp/cudnn-archive). Specifically, I got the 'Library for Linux (x86_64)', and the 'Runtime Library', 'Development Library' and 'Samples' for 'Ubuntu 20.04 x86_64 (Deb)'. Not sure if they're all needed though. Installed the .deb packages normally and the tar file as described here (https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html).

I don't think you need to compile jaxlib from source if you follow the standard install instructions for 'pip installation: GPU (CUDA)'. I couldn't even get it to compile properly but the install worked fine.

From here I just followed the included instructions. After that the warning disappeared and everything seemed to be okay. Also make sure your drivers on Windows are up to date. Good luck!