Jax 0.6.0 now requires cudnn 9.8

[adam-hartshorne]

Pip installing Torch2Jax from github overwrites the existing version of nvidia-cudnn-cu12 package to 9.7. From JAX 0.6, its requires 9.8.x

[rdyro]

Good catch, I just tried this in a new venv:

$ uv pip install -U "jax[cuda]==0.6.0" torch

and I'm getting an incompatibility between JAX and PyTorch 😢

Using Python 3.12.7 environment at: venv
  × No solution found when resolving dependencies:
  ╰─▶ Because torch>=2.2.0,<=2.3.1 depends on nvidia-cudnn-cu12{platform_machine
      == 'x86_64' and sys_platform == 'linux'}==8.9.2.26 and
      nvidia-cudnn-cu12{platform_machine == 'x86_64' and sys_platform ==
      'linux'}==9.1.0.70, we can conclude that nvidia-cudnn-cu12!=8.9.2.26,
      torch>=2.2.0,<=2.6.0, nvidia-cudnn-cu12{platform_machine == 'x86_64' and
      sys_platform == 'linux'}!=9.1.0.70 are incompatible. (1)

      Because only the following versions of torch are available:
          torch==1.0.0
          torch==1.0.1
          torch==1.1.0
          torch==1.2.0
          torch==1.3.0
          torch==1.3.1
          torch==1.4.0
          torch==1.5.0
          torch==1.5.1
          torch==1.6.0
          torch==1.7.0
          torch==1.7.1
          torch==1.8.0
          torch==1.8.1
          torch==1.9.0
          torch==1.9.1
          torch==1.10.0
          torch==1.10.1
          torch==1.10.2
          torch==1.11.0
          torch==1.12.0
          torch==1.12.1
          torch==1.13.0
          torch==1.13.1
          torch==2.0.0
          torch==2.0.1
          torch==2.1.0
          torch==2.1.1
          torch==2.1.2
          torch==2.2.0
          torch==2.2.1
          torch==2.2.2
          torch==2.3.0
          torch==2.3.1
          torch==2.4.0
          torch==2.4.1
          torch==2.5.0
          torch==2.5.1
          torch==2.6.0
          torch==2.7.0
      and torch<=2.1.2 has no wheels with a matching Python ABI tag (e.g.,
      `cp312`), we can conclude that torch<=2.1.2 cannot be used.
      And because we know from (1) that nvidia-cudnn-cu12!=8.9.2.26,
      torch>=2.2.0,<=2.6.0, nvidia-cudnn-cu12{platform_machine == 'x86_64' and
      sys_platform == 'linux'}!=9.1.0.70 are incompatible, we can conclude that
      torch<2.7.0 depends on one of:
          nvidia-cudnn-cu12==8.9.2.26
          nvidia-cudnn-cu12==9.1.0.70

      And because torch==2.7.0 depends on nvidia-cudnn-cu12{platform_machine
      == 'x86_64' and sys_platform == 'linux'}==9.5.1.17
      and jax-cuda12-plugin[with-cuda]==0.6.0 depends on
      nvidia-cudnn-cu12>=9.8, we can conclude that all versions of torch and
      jax-cuda12-plugin[with-cuda]==0.6.0 are incompatible.
      And because only jax-cuda12-plugin[with-cuda]<=0.6.0 is available and
      jax[cuda]==0.6.0 depends on jax-cuda12-plugin[with-cuda]==0.6.0, we can
      conclude that jax[cuda]==0.6.0 and all versions of torch are incompatible.
      And because you require jax[cuda]==0.6.0 and torch, we can conclude that
      your requirements are unsatisfiable.

      hint: You require CPython 3.12 (`cp312`), but we only found wheels for
      `torch` (v2.1.2) with the following Python ABI tags: `cp38`, `cp39`,
      `cp310`, `cp311`

Hopefully pytorch can release new wheels soon. torch2jax itself doesn't pin any jax or torch versions, so as long as torch and jax can be installed in a compatible way, it should hopefully work!

[adam-hartshorne]

I have managed to get it working. You have to install torch first, then delete all the nvidia runtimes it installs, then install jax.

[rdyro]

Huh, I suppose that makes sense. If NVIDIA maintains good backwards compatibility, that could work. Thanks for investigating this!