Trying to run graphcast from cds input data but stuck on JAX version issue

[Mehzooz]

I am trying to run graphcast using cds input data. I have installed the model and requirements following the isntructions on this page. I am getting this error message:

AttributeError: module 'jax' has no attribute 'linear_util'

I saw a similar issue[12] and I tried to fix it using the response from that. The solution was to change the jax version on requirements.txt like this:

jax[cuda11_pip]==0.4.23
git+https://github.com/deepmind/graphcast.git

But then I am getting this error:

ERROR: Ignored the following yanked versions: 0.4.32
ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.23+cuda11.cudnn86; extra == "cuda11-pip" (from jax[cuda11-pip]) (from versions: 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29, 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35)
ERROR: No matching distribution found for jaxlib==0.4.23+cuda11.cudnn86; extra == "cuda11-pip"

How do I fix this?

[owainkenwayucl]

For annoying reasons, Google de-list old versions of JaX from PyPI. You need to use their Google storage API mirror:

pip install jax[cuda11_pip]==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html