jax, jaxlib compatability with s2fft

[MEH04]

Hello,

I am trying to install GPU-enabled jax and jaxlib onto a linux-based remote computer. I am following the instructions on this website: https://docs.jax.dev/en/latest/installation.html. However, when running pip install --upgrade "jax[cuda12]", I find a version incompatibility, where s2fft is incompatible with any jax version >=0.60. I have attempted to follow the instructions on the jax installation website for installing older versions, but fail to find any that are compatible with the GPU on this system, where my CUDA is 12.2, CUDNN is 8.9.3 and NCCL is 2.22.3.

I am wondering if there (a) plans to upgrade s2fft comptability to 0.6.0, or (b) any suggestions as to how I can install an older version jax and jaxlib onto my system.

Thank you.

[kmulderdas]

From the JAX changelogs it looks like your CUDNN version (8.9.3) is incompatible with the JAX version (0.6.0+) you are trying to install. Instructions for upgrading CUDNN are installation specific and can be found here.

For the GPU specific s2fft installation grabbing an older compatible JAX version (e.g. 0.5.3) should work via pip install --upgrade "jax[cuda12]"==0.5.3 with the listed CUDA/CUDNN versions.

[jasonltorchinsky]

I had encountered the same issue, and your jax installation instructions seemed to fix it! There didn't appear to be instructions from Jax on how to install an older version with CUDA. I think it would be worth including a warning about that in the s2fft installation instructions.