revert to jax docs install

Author: mmcky

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Install jax (and install checks for GPU)
         shell: bash -l {0}
         run: |
-          pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          pip install -U "jax[cuda12]"
           python --version
           python scripts/test-jax-install.py
           nvidia-smi