Update JAX to use cuda13-local for CUDA 13 support

Changes all workflows from cuda12-local to cuda13-local to match
the CUDA 13 installation in the new AMI (ami-0edec81935264b6d3).

Author: mmcky

.github/workflows/cache.yml

@@ -23,7 +23,7 @@ jobs:
       - name: Install JAX and Numpyro
         shell: bash -l {0}
         run: |
-          pip install --upgrade "jax[cuda12-local]"
+          pip install --upgrade "jax[cuda13-local]"
           pip install numpyro
           python scripts/test-jax-install.py
       - name: Check nvidia drivers

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ jobs:
       - name: Install JAX and Numpyro
         shell: bash -l {0}
         run: |
-          pip install --upgrade "jax[cuda12-local]"
+          pip install --upgrade "jax[cuda13-local]"
           pip install numpyro
           python scripts/test-jax-install.py
       - name: Check nvidia Drivers

.github/workflows/publish.yml

@@ -22,7 +22,7 @@ jobs:
       - name: Install JAX and Numpyro
         shell: bash -l {0}
         run: |
-          pip install --upgrade "jax[cuda12-local]"
+          pip install --upgrade "jax[cuda13-local]"
           pip install numpyro
           python scripts/test-jax-install.py
       - name: Check nvidia drivers