Jax's cuda installation should align with TF's cuda installation

[davidshen84]

TF has this page that details its cuda and cudnn installation. Jax only briefly mentioned it here. If I want to use both Jax and TF on my system, I have to do trial and error a few times to find out which TF and Jax versions can work with my current cuda and cudnn installations on my system.

I have this simple docker file for TF 2.16 that was working a few months back.

dockerfile TF 2.16

from tensorflow/tensorflow:2.16.0-gpu-jupyter

run pip install --upgrade pip
copy requirements.txt .
run pip install --upgrade --requirement requirements.txt --find-link https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

entrypoint ["bash", "-c", "source /etc/bash.bashrc && jupyter lab --notebook-dir=/tf --ip 0.0.0.0 --no-browser --allow-root"]

requirements.txt

clu
flax
jax[cuda12_local]
jupyter_contrib_nbextensions
jupyterlab
jupyterlab-code-formatter
tensorflow_datasets
tensorflow_hub
tensorflow_text

I use jax[cuda12_local] because I want Jax to use the same cuda and cudnn installation as TF did, and it worked.

But when I upgraded to TF 2.18, I had to use jax[cuda12]. The result is that my docker image size increased from about 14GB to over 22GB.