Issue 1:
Description

I am unable to install the latest version of JAX 0.4.12 through the command:

pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Version 0.4.11 is installed instead.

What jax/jaxlib version are you using?

No response

Which accelerator(s) are you using?

No response

Additional system info

No response

NVIDIA GPU info

+---------------------------------------------------------------------------------------+

| NVIDIA-SMI 530.30.02 Driver Version: 530.30.02 CUDA Version: 12.1 |

|-----------------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|=========================================+======================+======================|

| 0 NVIDIA RTX 6000 Ada Gener... Off| 00000000:41:00.0 Off | Off |

| 30% 44C P8 26W / 300W| 6MiB / 49140MiB | 0% Default |

| | | N/A |

+-----------------------------------------+----------------------+----------------------+

| 1 NVIDIA RTX 6000 Ada Gener... Off| 00000000:61:00.0 Off | Off |

| 30% 39C P8 22W / 300W| 226MiB / 49140MiB | 0% Default |

| | | N/A |

+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=======================================================================================|

| 0 N/A N/A 1925 G /usr/lib/xorg/Xorg 4MiB |

| 1 N/A N/A 1925 G /usr/lib/xorg/Xorg 135MiB |

| 1 N/A N/A 2061 G /usr/bin/gnome-shell 52MiB |

+---------------------------------------------------------------------------------------+

Issue 2:
Description

For CUDA 12.1 and cuDNN 8.8 No wheel is available. The suggested whl to use was CUDA 12 and cuDNN 8.9 but this does not work and I only have cuDNN 8.8.

Thanks,

What jax/jaxlib version are you using?

jaxlib==0.4.12

Which accelerator(s) are you using?

GPU

Additional system info

Python 3.8.10, Ubuntu 20.04.6 LTS

NVIDIA GPU info

Mon Jun 12 14:37:26 2023

+---------------------------------------------------------------------------------------+

| NVIDIA-SMI 530.41.03 Driver Version: 530.41.03 CUDA Version: 12.1 |

|-----------------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|=========================================+======================+======================|

| 0 NVIDIA RTX A6000 Off| 00000000:65:00.0 On | Off |

| 32% 61C P2 79W / 300W| 48625MiB / 49140MiB | 0% Default |

| | | N/A |

+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=======================================================================================|

| 0 N/A N/A 1189 G /usr/lib/xorg/Xorg 77MiB |

| 0 N/A N/A 1939 G /usr/lib/xorg/Xorg 161MiB |

| 0 N/A N/A 2074 G /usr/bin/gnome-shell 75MiB |

| 0 N/A N/A 6749 C python3 37240MiB |

| 0 N/A N/A 7046 C python3 10934MiB |

| 0 N/A N/A 11541 G /usr/lib/firefox/firefox 121MiB |

+---------------------------------------------------------------------------------------+