Use cuda12-local for JAX installation across all workflows

Changes all workflows to use jax[cuda12-local] to leverage the CUDA
and cuDNN libraries pre-installed in the new AMI (ami-0edec81935264b6d3).
This is faster than cuda12_pip and uses the system libraries.

Also removes version pin (==0.6.2) from cache.yml and publish.yml
to make all workflows consistent.

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_pip]==0.6.2"
+          pip install --upgrade "jax[cuda12-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_pip]"
+          pip install --upgrade "jax[cuda12-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_pip]==0.6.2"
+          pip install --upgrade "jax[cuda12-local]"
           pip install numpyro
           python scripts/test-jax-install.py
       - name: Check nvidia drivers