FIX: pin jax=0.6.2 to all workflows (#498)

Author: mmcky

.github/workflows/cache.yml

@@ -24,7 +24,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-          pip install --upgrade "jax[cuda12-local]"
+          pip install --upgrade "jax[cuda12-local]==0.6.2"
           pip install numpyro  
           python scripts/test-jax-install.py
       - name: Check nvidia drivers

.github/workflows/publish.yml

@@ -23,7 +23,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-          pip install --upgrade "jax[cuda12-local]"
+          pip install --upgrade "jax[cuda12-local]==0.6.2"
           pip install numpyro  
           python scripts/test-jax-install.py
       - name: Check nvidia drivers