fix: JAX version pinning

Author: mmcky

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
         run: |
           pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
           pip install pyro-ppl
-          pip install --upgrade "jax[cuda12-local]==0.6.2"
+          pip install --upgrade "jax[cuda12-local]==0.8.0"
           pip install numpyro pyro-ppl
           python scripts/test-jax-install.py
       - name: Check nvidia Drivers