Request: Document compatible versions of Flax, Distrax, and CUDA libraries for older JAX releases

[Revivekirin]

Hello,

I’m currently working on an RL project with the JAX framework as a beginner. While the tutorial command

pip install jax[cuda12]

is very helpful for setting up the latest JAX environment on CUDA GPUs, I faced some difficulties when trying to inherit and run a project that depends on older JAX versions.

For example, when installing JAX 0.6.2 with:

pip install --upgrade \
  jax==0.6.2 \
  jaxlib==0.6.2 \
  jax-cuda12-plugin==0.6.2 \
  -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

I found that I also needed to install additional libraries like Flax and Distrax, but it was not clear which versions were compatible with JAX 0.6.2.

Moreover, I had to manually install nvidia-cudnn-cu12. The latest version (9.12) did not work, and I had to specifically download version 9.8.0.97, which took extra time to figure out.

Would it be possible to add documentation that explicitly lists compatible versions of Flax, Distrax, and CUDA/cuDNN for older JAX releases?
This would save users a lot of time when setting up legacy projects that don’t run on the latest JAX release.

Thank you!

[jakevdp]

For compatibility of downstream packages like flax and distrax, I think it would be best to inquire directly with those projects (there are a lot of projects built on JAX, and we don't have the bandwidth or expertise to maintain a complete catalog of version compatibility).

For CUDA compatibility of older JAX versions, I find it easiest to look at JAX installation docs for the particular version. For example, here's the CUDA compatibility info for JAX v0.6.2: https://github.com/jax-ml/jax/blob/jax-v0.6.2/docs/installation.md?plain=1#L73

If you do pip install jax[cuda12]==0.6.2 then you should get all the nvidia dependencies automatically. I think that should pin the required older versions, but please let me know if that's not the case.

[Revivekirin]

Thank you for pointing this out! I hadn’t noticed the section about cuDNN compatibility before — that clears things up

[Revivekirin]

We noticed that when experimenting in an Ubuntu 24.04 environment with Miniconda, installing JAX 0.6.2 did not automatically pull in compatible versions of Flax and Distrax. We had to install these libraries manually, but we could not find clear documentation regarding which versions are guaranteed to be compatible with JAX.

Additionally, if we install the latest versions of Flax or Distrax, the dependency resolver automatically upgrades JAX to 0.7.1, which introduces compatibility issues for our workflow. This makes version pinning difficult without prior knowledge of the correct combinations.

It would be helpful if the documentation (similar to the cuDNN compatibility notes) could also include guidance on Flax and Distrax compatibility with JAX releases, so that users can set up environments more reliably.

[jakevdp]

installing JAX 0.6.2 did not automatically pull in compatible versions of Flax and Distrax.

I would not expect compatible versions of downstream packages to be pulled in automatically. You'll need to check with the particular downstream projects regarding JAX version compatibility.

It would be helpful if the documentation could also include guidance on Flax and Distrax compatibility with JAX releases

As I mentioned above, the JAX team does not have the bandwidth nor the expertise to advise about compatibility/requirements of downstream packages. I'd suggest checking with the Flax and Distrax maintainers for guidance on that.