MAINT: remove jax, pyro, torch, and gpu related software installs + GPU admonition (#453)

* MAINT: remove installs of gpu related software

* Update lectures requiring GPU access

* add arviz to environment

* fix installs at the lecture level for non-gpu packages

* add pyro-ppl package to GPU packages

* switch to hide-output for collab compat testing

* ensure latest collab docker runtime is used

* add arviz to environment

* ensure latest arviz in collab

* update collab build software

Author: mmcky

.github/workflows/ci.yml

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

.github/workflows/collab.yml

@@ -4,7 +4,7 @@ jobs:
   execution-checks:
     runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=ubuntu24-gpu-x64/disk=large"
     container:
-      image: docker://us-docker.pkg.dev/colab-images/public/runtime
+      image: docker://us-docker.pkg.dev/colab-images/public/runtime:latest
       options: --gpus all
     steps:
       - uses: actions/checkout@v4
@@ -32,7 +32,7 @@ jobs:
       - name: Install Build Software
         shell: bash -l {0}
         run: |
-          pip install jupyter-book==0.15.1 docutils==0.17.1 quantecon-book-theme==0.7.2 sphinx-tojupyter==0.3.0 sphinxext-rediraffe==0.2.7 sphinx-reredirects sphinx-exercise==0.4.1 sphinxcontrib-youtube==1.1.0 sphinx-togglebutton==0.3.1 arviz==0.13.0 sphinx-proof
+          pip install jupyter-book==1.0.3 quantecon-book-theme==0.8.2 sphinx-tojupyter==0.3.0 sphinxext-rediraffe==0.2.7 sphinxcontrib-youtube==1.3.0 sphinx-togglebutton==0.3.2 arviz sphinx-proof sphinx-exercise sphinx-reredirects
       # Build of HTML (Execution Testing)
       - name: Build HTML
         shell: bash -l {0}

environment.yml

@@ -16,6 +16,3 @@ dependencies:
     - ghp-import==1.1.0
     - sphinxcontrib-youtube==1.3.0 #Version 1.3.0 is required as quantecon-book-theme is only compatible with sphinx<=5
     - sphinx-togglebutton==0.3.2
-    # Docker Requirements
-    - pytz
-

lectures/_admonition/gpu.md

@@ -0,0 +1,9 @@
+```{admonition} GPU
+:class: warning
+
+This lecture was built using a machine with the latest CUDA and CUDANN frameworks installed with access to a GPU.
+
+To run this lecture on [Google Colab](https://colab.research.google.com/), click on the "play" icon top right, select Colab, and set the runtime environment to include a GPU.
+
+To run this lecture on your own machine, you need to install the software listed following this notice.
+```

lectures/ar1_bayes.md

@@ -13,15 +13,25 @@ kernelspec:
 
 # Posterior Distributions for  AR(1) Parameters
 
-We'll begin with some Python imports.
+```{include} _admonition/gpu.md
+```
 
+```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install numpyro jax
+```
+
+In addition to what's included in base Anaconda, we need to install the following packages
 
 ```{code-cell} ipython3
 :tags: [hide-output]
 
-!pip install arviz pymc numpyro jax
+!pip install arviz pymc
 ```
 
+We'll begin with some Python imports.
+
 ```{code-cell} ipython3
 
 import arviz as az

lectures/back_prop.md

@@ -13,10 +13,21 @@ kernelspec:
 
 # Introduction to Artificial Neural Networks
 
+```{include} _admonition/gpu.md
+```
+
+```{code-cell} ipython3
+:tags: [skip-execution]
+
+!pip install --upgrade jax
+```
+
+In addition to what's included in base Anaconda, we need to install the following packages
+
 ```{code-cell} ipython3
 :tags: [hide-output]
 
-!pip install --upgrade jax jaxlib kaleido
+!pip install kaleido
 !conda install -y -c plotly plotly plotly-orca retrying
 ```
 

lectures/bayes_nonconj.md

@@ -13,6 +13,15 @@ kernelspec:
 
 # Non-Conjugate Priors
 
+```{include} _admonition/gpu.md
+```
+
+```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install numpyro pyro-ppl torch jax
+```
+
 This lecture is a sequel to the {doc}`quantecon lecture <prob_meaning>`.
 
 That lecture offers a Bayesian interpretation of probability in a setting in which the likelihood function and the prior distribution
@@ -43,13 +52,6 @@ The two Python modules are
 
 As usual, we begin by importing some Python code.
 
-```{code-cell} ipython3
-:tags: [hide-output]
-
-# install dependencies
-!pip install numpyro pyro-ppl torch jax
-```
-
 ```{code-cell} ipython3
 import numpy as np
 import seaborn as sns

lectures/mix_model.md

@@ -14,11 +14,12 @@ kernelspec:
 (likelihood-ratio-process)=
 # Incorrect Models
 
-In addition to what's in Anaconda, this lecture will need the following libraries:
+```{include} _admonition/gpu.md
+```
+
 ```{code-cell} ipython
----
-tags: [hide-output]
----
+:tags: [hide-output]
+
 !pip install numpyro jax
 ```