JAX 0.4.14 jaxlib cudnn error - non-docker installation

Hi,

I am currently installing the latest Alphafold (from main branch, so 2.3.2 plus new features, new openmm=7.7.0 , no openmm patch, python=3.10 etc..) on a HPC cluster without Docker. I am confused about the jaxlib versions, as in the requirements. txt file it says jax=0.4.14 , but in the dockerfile it is still at 0.3.25.

CUDA=11.2.2 (or 11.8.0 or 11.1.0 , tried all of them)
If I generate a conda env with:
(from the docker file and requirements.txt)

```
channels:
- pytorch
- conda-forge
- defaults
- anaconda
- bioconda
dependencies:
- python==3.10
- pip
- openmm==7.7.0
- cudnn # Change version if not compatible with current system
- cudatoolkit
- pdbfixer
- hmmer==3.4
- hhsuite==3.3.0
- kalign2==2.04
- pip:
- absl-py==1.0.0
- biopython==1.79
- chex==0.0.7
- dm-haiku==0.0.10
- immutabledict==2.0.0
- ml-collections==0.1.0
- numpy==1.24.3
- scipy==1.11.1
- tensorflow-cpu==2.13.0
- jax==0.4.14
- pandas==2.0.3
- dm-tree==0.1.8
```

and then do

`pip3 install --upgrade --no-cache-dir jax==0.3.25 jaxlib==0.3.25+cuda11.cudnn805 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`

If I do :
`python -c "import jax; print(f'Jax backend: {jax.default_backend()}')` , I get `Jax backend: gpu` and if I check for jax.devices(), I get
`jax.devices output [gpu(id=0)] `

alphafold runs ( 5 min slower than the old version on the cluster for a 60AA protein), but I get the Warning during runtime:
`Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.`

As described in issue #88 , tensorflow is installed as tensorflow-cpu on purpose, if I understand correctly.

When I install the same with
` pip install --upgrade --no-cache-dir jax==0.4.14 jaxlib==0.4.14+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`

I dont get the Warning anymore, but alphafold fails after the "features" step during runtime with :
Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR

Anyone any idea?