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