fix: restore numpyro to accelerate jax

Author: mmcky

.github/workflows/cache.yml

@@ -64,6 +64,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade "jax[cuda]==0.4.2" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          pip install --upgrade "numpyro[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
           nvidia-smi
       - name: Build HTML
         shell: bash -l {0}