JAX CUDA Enabled Installation on HPC Cluster

[Naeemkh]

Summarize the working versions of Python, CUDA, JAX, and cuDNN required for successful installation of CUDA-enabled JAX on an HPC cluster (x86 architecture).

[Naeemkh]

This combination works for me:

module load python
mamba create -n jaxtest python=3.12 -y
mamba activate jaxtest
module load cuda/12.4.1-fasrc01
module load cudnn/9.1.1.17_cuda12-fasrc01
pip install jax==0.4.29  jaxlib==0.4.29+cuda12.cudnn91  -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
python -c "import jax; print(jax.devices())"

[mmshad]

This combination worked for me:

module load python
mamba create -n jaxtest python=3.13 -y
mamba activate jaxtest
module load cuda/12.9.1-fasrc01
module load cudnn/9.10.2.21_cuda12-fasrc01
pip install -U "jax[cuda12]"

Test it

python -c "import jax; print(jax.devices())"
[CudaDevice(id=0)]