JAX and CUDA compatibility when installing JAX into WSL2

[ToshiyukiBandai]

Hi all,

I'm installing jax.0.4.28 into Windows Linux Subsystem 2 (WSL2). I understand the support for WSL2 is experimental, but I want to share some experiences to hopefully get some remedies.

After installing jax via the commad as in the installation document

pip install -U "jax[cuda12]"

I test running JAX and got the following error

This error disappears after a trick (#18027)

python3 -m pip install tensorflow[and-cuda]

I don't know why the error disappeared because by running nvidia-smi I can see the version of CUDA is 12.3, not 12.5, which was required at first.

Does anyone know what is going on?

[ToshiyukiBandai]

For record, I just needed to install the latest version of NVIDIA driver (https://docs.nvidia.com/cuda/wsl-user-guide/index.html). Make sure to install the driver into the Windows system only, not the WSL2 Linux system. Unfortunately, it seems some bug in WSL2 showing "segmentation fault" if you check the driver version by nvidia-smi (microsoft/WSL#11277). I got the message on WSL2 on WIndows 11 but not WSL2 on Windows10.

[hawkinsp]

Yes, that is a warning saying that your CUDA driver is older than your ptxas version. Upgrading your driver will work.

python3 -m pip install tensorflow[and-cuda] works also because that will downgrade ptxas, since tensorflow depends on a specific older version. In general I'd recommend upgrading rather than downgrading!