jax 0.5.3 compatibility issues

[michaelxu01]

When installing Jax 0.5.3, the default version for nvidia-cudnn-cu12 is 9.11, which gives the following error at model compilation. Downgrading to nvidia-cudnn-cu12<=9.10 seems to fix this issue.

This is in addition to the issue https://github.com/jax-ml/jax/pull/28897 and needing to downgrade to nvidia-cublas-cu12<12.9

This is in Python 3.12.11.

INFO:phaser.utils.num:JIT-compiling kernel 'run_model'...
2025-07-28 20:58:26.880340: W external/xla/xla/service/gpu/autotuning/conv_algorithm_picker.cc:864] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.

[michaelxu01]

Edit for additional troubleshooting information - previous solution of disabling or adjusting memory preallocation doesn't seem to work in this case https://github.com/jax-ml/jax/issues/12461

[hexane360]

Probably the easiest solution here is to update JAX version (e.g. to 0.7.0). I've done some testing on it, seems like everything is working.
IIRC this was a bug in a specific version of the nvidia-cublas-cu12 PyPI package that later got fixed. Maybe a fresh environment would also fix?

[michaelxu01]

Good to know, I'll give 0.7.0 a shot and see how it goes
I did try a fresh environment for the above issue, and it didn't seem to resolve it

[michaelxu01]

Seems like V100s are no longer supported? jax-ml/jax#30663

[michaelxu01]

Jax 0.7.0 with the following packages work me. The cublas issue from before is resolved, but the cudnn version needs to be downgraded to <9.11 if running on V100s.

nvidia-cublas-cu12==12.9.1.4
nvidia-cuda-cupti-cu12==12.9.79
nvidia-cuda-nvcc-cu12==12.9.86
nvidia-cuda-nvrtc-cu12==12.9.86
nvidia-cuda-runtime-cu12==12.9.79
nvidia-cudnn-cu12==9.10.2.21
nvidia-cufft-cu12==11.4.1.4
nvidia-cusolver-cu12==11.7.5.82
nvidia-cusparse-cu12==12.5.10.65
nvidia-nccl-cu12==2.27.6
nvidia-nvjitlink-cu12==12.9.86
nvidia-nvshmem-cu12==3.3.9

I did however have to run unset LD_LIBRARY_PATH based on: jax-ml/jax#28929

EDIT: python 3.13.5