Repository: google/jax

Issue:
Installation Bug: No Matching distribution found for jaxlib==0.4.12+cuda12.cudnn88; extra == "cuda12_local" (from jax[cuda12_local])
### 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 |
+---------------------------------------------------------------------------------------+



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -22,9 +22,9 @@
 _current_jaxlib_version = '0.4.12'
 # The following should be updated with each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.4.12'
-_available_cuda11_cudnn_versions = ['82', '86']
+_available_cuda11_cudnn_versions = ['86']
 _default_cuda11_cudnn_version = '86'
-_default_cuda12_cudnn_version = '88'
+_default_cuda12_cudnn_version = '89'
 _libtpu_version = '0.1.dev20230608'
 
 _dct = {}


Test patch: