Compatible with CUDA 11.2?

[joshdorrington]

I'm trying to set up neuralgcm on my local cluster. I think there might be a fundamental version incompatibility, but I wanted to check if I'm missing anything obvious first.

Running nvidia-smi in terminal:

NVIDIA-SMI 460.106.00 Driver Version: 460.106.00 CUDA Version: 11.2 |

I tried finding a suitable wheel from here: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

But there seem to be none that support jaxlib>=4.27, CUDA<=11.2 and Python>=3.10

Am I out of luck until I can get the sysadmin to upgrade our CUDA to 12?

[shoyer]

Unfortunately I don't know much about CUDA versions and JAX.

This would be a good question for posting on the JAX GitHub tracker: https://github.com/google/jax

[joshdorrington]

OK thanks Stephan,
Just to clarify, these seemed to be the requirements of neuralgcm specifically (jaxlib>=4.27 and Python>=3.10) which is why I asked here. I'll take my question over to JAX if you think they will have a better idea.

Cheers,
Josh