Fix JAX CUDA installation to use cuda12_pip instead of cuda12-local

Changes jax[cuda12-local] to jax[cuda12_pip] to avoid cuDNN version
compatibility issues. The cuda12_pip variant includes compatible
CUDA and cuDNN libraries bundled with JAX, preventing runtime errors
from mismatched local CUDA installations.

Fixes cuDNN 9.10.0 backward-compatibility error.

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