MAINT: Install numpyro with GPU

Author: mmcky

.github/workflows/cache.yml

@@ -64,6 +64,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade "jax[cuda]" -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}

.github/workflows/ci.yml

@@ -61,6 +61,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade "jax[cuda]" -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: Install latex dependencies
         shell: bash -l {0}

.github/workflows/publish.yml

@@ -67,6 +67,7 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade "jax[cuda]" -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: Display Conda Environment Versions
         shell: bash -l {0}