Update jax box at top of lecture, update CUDA to 11.8 (#319)

* Fix jax upgrade with GPU

* check using cuda

* try cudnn82

* revert back to default

* check cudnn installation status

* try docker/cuda=11.8

* remove failing cmd

* set jax==0.4.2

* adjust pip install and moved the jax[cpu] into the gpu warning box

---------

Co-authored-by: Humphrey Yang <u6474961@anu.edu.au>
Co-authored-by: mmcky <mamckay@gmail.com>

Author: 3 people

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
     needs: deploy-runner
     runs-on: [self-hosted, cml-gpu]
     container:
-      image: docker://nvidia/cuda:11.2.1-devel-ubuntu20.04
+      image: docker://nvidia/cuda:11.8.0-devel-ubuntu22.04
       options: --gpus all
     steps:
       - uses: actions/checkout@v3

lectures/kesten_processes.md

@@ -19,14 +19,14 @@ kernelspec:
 
 # Kesten Processes and Firm Dynamics
 
-```{admonition} GPU in use
+```{admonition} GPU
 :class: warning
 
-This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming. 
+This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming.
 
 Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.
 
-Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax) for installing JAX with GPU support.
+Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax) for installing JAX with GPU support. If you would like to install jax running on the `cpu` only you can use `pip install jax[cpu]`
 ```
 
 ```{index} single: Linear State Space Models
@@ -44,9 +44,6 @@ tags: [hide-output]
 ---
 !pip install quantecon
 !pip install --upgrade yfinance
-# If your machine has CUDA support, please follow the guide in GPU Warning. 
-# Otherwise, run the line below:
-!pip install --upgrade "jax[CPU]"
 ```
 
 ## Overview
@@ -686,7 +683,7 @@ s_init = 1.0      # initial condition for each firm
 :class: dropdown
 ```
 
-Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html). 
+Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html).
 
 First let's import the necessary modules and check the backend for JAX
 
@@ -731,10 +728,10 @@ def generate_draws(μ_a=-0.5,
         exp_a = jnp.exp(a_random[t, :])
         exp_b = jnp.exp(b_random[t, :])
         exp_e = jnp.exp(e_random[t, :])
-        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar, 
+        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar,
                              exp_e,
                              exp_a * s[:, t] + exp_b))
-    
+
     return s[:, -1]
 
 %time data = generate_draws().block_until_ready()
@@ -761,7 +758,7 @@ plt.show()
 
 The plot produces a straight line, consistent with a Pareto tail.
 
-It is possible to further speed up our code by replacing the `for` loop with [`lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) 
+It is possible to further speed up our code by replacing the `for` loop with [`lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)
 to reduce the loop overhead in the compilation of the jitted function
 
 ```{code-cell} ipython3
@@ -779,10 +776,10 @@ def generate_draws_lax(μ_a=-0.5,
                        M=1_000_000,
                        s_init=1.0,
                        seed=123):
-  
+
     key = random.PRNGKey(seed)
     keys = random.split(key, 3)
-    
+
     # Generate random draws and initial values
     a_random = μ_a + σ_a * random.normal(keys[0], (T, M))
     b_random = μ_b + σ_b * random.normal(keys[1], (T, M))
@@ -792,11 +789,11 @@ def generate_draws_lax(μ_a=-0.5,
     # Define the function for each update
     def update_s(s, a_b_e_draws):
         a, b, e = a_b_e_draws
-        res = jnp.where(s < s_bar, 
-                        jnp.exp(e), 
+        res = jnp.where(s < s_bar,
+                        jnp.exp(e),
                         jnp.exp(a) * s + jnp.exp(b))
         return res, res
-    
+
     # Use lax.scan to perform the calculations on all states
     s_final, _ = lax.scan(update_s, s, (a_random, b_random, e_random))
     return s_final
@@ -877,4 +874,4 @@ plt.show()
 ```
 
 ```{solution-end}
-```
+```

lectures/wealth_dynamics.md

@@ -19,14 +19,14 @@ kernelspec:
 
 # Wealth Distribution Dynamics
 
-```{admonition} GPU in use
+```{admonition} GPU
 :class: warning
 
 This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming.
 
 Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.
 
-Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax#pip-installation-gpu-cuda) for installing JAX with GPU support.
+Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax) for installing JAX with GPU support. If you would like to install jax running on the `cpu` only you can use `pip install jax[cpu]`
 ```
 
 ```{contents} Contents
@@ -40,9 +40,7 @@ In addition to what's in Anaconda, this lecture will need the following librarie
 tags: [hide-output]
 ---
 !pip install quantecon
-# If your machine has CUDA support, please follow the guide in GPU Warning.
-# Otherwise, run the line below:
-!pip install --upgrade "jax[CPU]"
+!pip install myst-nb
 ```
 
 ## Overview