From 0dc04c0e21776ce2282e08942ec3c67787a414ab Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Wed, 28 Dec 2022 22:33:39 +1100
Subject: [PATCH 01/15] porting to jax

---
 lectures/kesten_processes.md | 76 +++++++++++++++++++++++++++++++-----
 1 file changed, 67 insertions(+), 9 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 7cf89d387..834137dff 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -673,14 +673,15 @@ s_init = 1.0      # initial condition for each firm
 :class: dropdown
 ```
 
-Here's one solution.  First we generate the observations:
+Here's one solution in Jax. 
 
-```{code-cell} ipython3
-from numba import njit, prange
-from numpy.random import randn
+First we generate the observations:
 
+```{code-cell} ipython3
+import jax.numpy as jnp
+from jax import jit, random
 
-@njit(parallel=True)
+@jit
 def generate_draws(μ_a=-0.5,
                    σ_a=0.1,
                    μ_b=0.0,
@@ -690,7 +691,66 @@ def generate_draws(μ_a=-0.5,
                    s_bar=1.0,
                    T=500,
                    M=1_000_000,
-                   s_init=1.0):
+                   s_init=1.0,
+                   seed=123):
+    key = random.PRNGKey(seed)
+    # Generate arrays of random numbers
+    a_random = μ_a + σ_a * random.normal(key, (M, T))
+    b_random = μ_b + σ_b * random.normal(key, (M, T))
+    e_random = μ_e + σ_e * random.normal(key, (M, T))
+
+    # Initialize the array of s values with the initial value
+    s = jnp.full((M, T+1), s_init)
+
+    # Perform the calculations in a vectorized manner
+    for t in range(T):
+        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar, 
+                             jnp.exp(e_random[:, t]), 
+                             jnp.exp(a_random[:, t]) * s[:, t] + jnp.exp(b_random[:, t])))
+    
+    return s[:, -1]
+
+%time data = generate_draws().block_until_ready()
+```
+
+Now we produce the rank-size plot:
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+
+rank_data, size_data = qe.rank_size(data, c=0.01)
+ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)
+ax.set_xlabel("log rank")
+ax.set_ylabel("log size")
+
+plt.show()
+```
+The plot produces a straight line, consistent with a Pareto tail.
+
+Since we called `jax.jit` on the function, it runs even faster when we call the function again
+
+```{code-cell} ipython3
+%time data = generate_draws().block_until_ready()
+```
+
+We can also use Numba with `for` loops:
+
+```{code-cell} ipython3
+from numba import njit, prange
+from numpy.random import randn
+
+
+@njit(parallel=True)
+def generate_draws_numba(μ_a=-0.5,
+                         σ_a=0.1,
+                         μ_b=0.0,
+                         σ_b=0.5,
+                         μ_e=0.0,
+                         σ_e=0.5,
+                         s_bar=1.0,
+                         T=500,
+                         M=1_000_000,
+                         s_init=1.0):
 
     draws = np.empty(M)
     for m in prange(M):
@@ -707,7 +767,7 @@ def generate_draws(μ_a=-0.5,
 
     return draws
 
-data = generate_draws()
+%time data = generate_draws_numba()
 ```
 
 Now we produce the rank-size plot:
@@ -723,7 +783,5 @@ ax.set_ylabel("log size")
 plt.show()
 ```
 
-The plot produces a straight line, consistent with a Pareto tail.
-
 ```{solution-end}
 ```
\ No newline at end of file

From a5fce1178ed134660263dfa7f80bb73073653724 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Wed, 28 Dec 2022 23:21:23 +1100
Subject: [PATCH 02/15] a faster version using lax.scan

---
 lectures/kesten_processes.md | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 834137dff..e23d58c6b 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -680,6 +680,7 @@ First we generate the observations:
 ```{code-cell} ipython3
 import jax.numpy as jnp
 from jax import jit, random
+from jax.lax import scan
 
 @jit
 def generate_draws(μ_a=-0.5,
@@ -693,22 +694,25 @@ def generate_draws(μ_a=-0.5,
                    M=1_000_000,
                    s_init=1.0,
                    seed=123):
+  
     key = random.PRNGKey(seed)
-    # Generate arrays of random numbers
+    
+    # Generate random draws and initialize states
     a_random = μ_a + σ_a * random.normal(key, (M, T))
     b_random = μ_b + σ_b * random.normal(key, (M, T))
     e_random = μ_e + σ_e * random.normal(key, (M, T))
-
-    # Initialize the array of s values with the initial value
     s = jnp.full((M, T+1), s_init)
 
-    # Perform the calculations in a vectorized manner
-    for t in range(T):
-        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar, 
-                             jnp.exp(e_random[:, t]), 
-                             jnp.exp(a_random[:, t]) * s[:, t] + jnp.exp(b_random[:, t])))
-    
-    return s[:, -1]
+    # Define the function for each update
+    def update_s(s, a_b_e_draws):
+        a, b, e = a_b_e_draws
+        return s, jnp.where(s < s_bar, 
+                            jnp.exp(e), 
+                            jnp.exp(a) * s + jnp.exp(b))
+     
+    # Use scan to perform the calculations on all states
+    _, s = scan(update_s, s_init, (a_random, b_random, e_random))
+    return s[:,-1]
 
 %time data = generate_draws().block_until_ready()
 ```
@@ -739,7 +743,6 @@ We can also use Numba with `for` loops:
 from numba import njit, prange
 from numpy.random import randn
 
-
 @njit(parallel=True)
 def generate_draws_numba(μ_a=-0.5,
                          σ_a=0.1,

From 3b42f849dc94c053d4e1672e55883e0d068e1f86 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Wed, 28 Dec 2022 23:35:04 +1100
Subject: [PATCH 03/15] reduce redundancy

---
 lectures/kesten_processes.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index e23d58c6b..13bbba2ae 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -701,7 +701,7 @@ def generate_draws(μ_a=-0.5,
     a_random = μ_a + σ_a * random.normal(key, (M, T))
     b_random = μ_b + σ_b * random.normal(key, (M, T))
     e_random = μ_e + σ_e * random.normal(key, (M, T))
-    s = jnp.full((M, T+1), s_init)
+    s = jnp.full(M, s_init)
 
     # Define the function for each update
     def update_s(s, a_b_e_draws):

From e92ff6dd37ed175ff89d16421df648189b0f6038 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Thu, 29 Dec 2022 16:09:18 +1100
Subject: [PATCH 04/15] debugging

---
 lectures/kesten_processes.md | 31 ++++++++++++++++---------------
 1 file changed, 16 insertions(+), 15 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 13bbba2ae..6995b05e2 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -679,8 +679,7 @@ First we generate the observations:
 
 ```{code-cell} ipython3
 import jax.numpy as jnp
-from jax import jit, random
-from jax.lax import scan
+from jax import jit, random, lax
 
 @jit
 def generate_draws(μ_a=-0.5,
@@ -696,23 +695,25 @@ def generate_draws(μ_a=-0.5,
                    seed=123):
   
     key = random.PRNGKey(seed)
+    key, *subkeys = random.split(key, 4)
     
-    # Generate random draws and initialize states
-    a_random = μ_a + σ_a * random.normal(key, (M, T))
-    b_random = μ_b + σ_b * random.normal(key, (M, T))
-    e_random = μ_e + σ_e * random.normal(key, (M, T))
-    s = jnp.full(M, s_init)
+    # Generate random draws and initial values
+    a_random = μ_a + σ_a * random.normal(subkeys[0], (T, M))
+    b_random = μ_b + σ_b * random.normal(subkeys[1], (T, M))
+    e_random = μ_e + σ_e * random.normal(subkeys[2], (T, M))
+    s = jnp.full((M, ), s_init)
 
     # Define the function for each update
     def update_s(s, a_b_e_draws):
         a, b, e = a_b_e_draws
-        return s, jnp.where(s < s_bar, 
-                            jnp.exp(e), 
-                            jnp.exp(a) * s + jnp.exp(b))
-     
-    # Use scan to perform the calculations on all states
-    _, s = scan(update_s, s_init, (a_random, b_random, e_random))
-    return s[:,-1]
+        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
 
 %time data = generate_draws().block_until_ready()
 ```
@@ -731,7 +732,7 @@ plt.show()
 ```
 The plot produces a straight line, consistent with a Pareto tail.
 
-Since we called `jax.jit` on the function, it runs even faster when we call the function again
+Since we applied `jax.jit` on the function, it runs even faster when we call the function again
 
 ```{code-cell} ipython3
 %time data = generate_draws().block_until_ready()

From 9d0d6d491db6508497930f34339776f98eda98d4 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Mon, 2 Jan 2023 10:07:46 +1100
Subject: [PATCH 05/15] changes according to comments

---
 lectures/kesten_processes.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 6995b05e2..4028fd5bd 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -673,7 +673,7 @@ s_init = 1.0      # initial condition for each firm
 :class: dropdown
 ```
 
-Here's one solution in Jax. 
+Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html). 
 
 First we generate the observations:
 
@@ -695,12 +695,12 @@ def generate_draws(μ_a=-0.5,
                    seed=123):
   
     key = random.PRNGKey(seed)
-    key, *subkeys = random.split(key, 4)
+    keys = random.split(key, 3)
     
     # Generate random draws and initial values
-    a_random = μ_a + σ_a * random.normal(subkeys[0], (T, M))
-    b_random = μ_b + σ_b * random.normal(subkeys[1], (T, M))
-    e_random = μ_e + σ_e * random.normal(subkeys[2], (T, M))
+    a_random = μ_a + σ_a * random.normal(keys[0], (T, M))
+    b_random = μ_b + σ_b * random.normal(keys[1], (T, M))
+    e_random = μ_e + σ_e * random.normal(keys[2], (T, M))
     s = jnp.full((M, ), s_init)
 
     # Define the function for each update
@@ -738,7 +738,7 @@ Since we applied `jax.jit` on the function, it runs even faster when we call the
 %time data = generate_draws().block_until_ready()
 ```
 
-We can also use Numba with `for` loops:
+We can also use Numba with `for` loops to generate the observations (replicating the results we obtained with JAX):
 
 ```{code-cell} ipython3
 from numba import njit, prange

From 57077ed3da588023752c800240a36ad408068f70 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Wed, 4 Jan 2023 13:20:53 +1100
Subject: [PATCH 06/15] adding for loop

---
 lectures/kesten_processes.md | 98 +++++++++++++++++++++++++++++++-----
 1 file changed, 86 insertions(+), 12 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 4028fd5bd..2212e3594 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -675,13 +675,24 @@ s_init = 1.0      # initial condition for each firm
 
 Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html). 
 
-First we generate the observations:
+First let's do a quick setup and check the backend for JAX
 
 ```{code-cell} ipython3
+import jax
 import jax.numpy as jnp
-from jax import jit, random, lax
+from jax import random
 
-@jit
+# Use 64 bit floats with JAX in order to match NumPy/Numba code
+jax.config.update("jax_enable_x64", True)
+
+# Check if JAX is using GPU
+print(f"jax backend: {jax.devices()[0].platform}")
+```
+
+Now we can generate the observations:
+
+```{code-cell} ipython3
+@jax.jit
 def generate_draws(μ_a=-0.5,
                    σ_a=0.1,
                    μ_b=0.0,
@@ -693,6 +704,68 @@ def generate_draws(μ_a=-0.5,
                    M=1_000_000,
                    s_init=1.0,
                    seed=123):
+
+    key = random.PRNGKey(seed)
+    keys = random.split(key, 3)
+
+    # Generate arrays of random numbers
+    a_random = μ_a + σ_a * random.normal(keys[0], (T, M))
+    b_random = μ_b + σ_b * random.normal(keys[1], (T, M))
+    e_random = μ_e + σ_e * random.normal(keys[2], (T, M))
+
+    # Initialize the array of s values with the initial value
+    s = jnp.full((M, T+1), s_init)
+
+    # Perform the calculations in a vectorized manner for T periods
+    for t in range(T):
+        s = s.at[:, t+1].set(jnp.where(s[:, t] < s_bar, 
+                             jnp.exp(e_random[t, :]), 
+                             jnp.exp(a_random[t, :]) * s[:, t] + jnp.exp(b_random[t, :])))
+    
+    return s[:, -1]
+
+%time data = generate_draws().block_until_ready()
+```
+
+Since we applied `jax.jit` on the function, it runs even faster when we call the function again
+
+```{code-cell} ipython3
+%time data = generate_draws().block_until_ready()
+```
+
+Now we produce the rank-size plot to check the distribution:
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+
+rank_data, size_data = qe.rank_size(data, c=0.01)
+ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)
+ax.set_xlabel("log rank")
+ax.set_ylabel("log size")
+
+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) 
+to reduce the loop overhead in the compilation
+
+```{code-cell} ipython3
+from jax import lax
+
+@jax.jit
+def generate_draws_lax(μ_a=-0.5,
+                       σ_a=0.1,
+                       μ_b=0.0,
+                       σ_b=0.5,
+                       μ_e=0.0,
+                       σ_e=0.5,
+                       s_bar=1.0,
+                       T=500,
+                       M=1_000_000,
+                       s_init=1.0,
+                       seed=123):
   
     key = random.PRNGKey(seed)
     keys = random.split(key, 3)
@@ -715,10 +788,16 @@ def generate_draws(μ_a=-0.5,
     s_final, _ = lax.scan(update_s, s, (a_random, b_random, e_random))
     return s_final
 
-%time data = generate_draws().block_until_ready()
+%time data = generate_draws_lax().block_until_ready()
 ```
 
-Now we produce the rank-size plot:
+The compiled function is even faster
+
+```{code-cell} ipython3
+%time data = generate_draws_lax().block_until_ready()
+```
+
+Here we produce the same rank-size plot:
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
@@ -730,15 +809,10 @@ ax.set_ylabel("log size")
 
 plt.show()
 ```
-The plot produces a straight line, consistent with a Pareto tail.
-
-Since we applied `jax.jit` on the function, it runs even faster when we call the function again
 
-```{code-cell} ipython3
-%time data = generate_draws().block_until_ready()
-```
+We can also use Numba with `for` loops to generate the observations (replicating the results we obtained with JAX).
 
-We can also use Numba with `for` loops to generate the observations (replicating the results we obtained with JAX):
+The results will be slightly different since the pseudo random number generation is implemented [differently in JAX](https://www.kaggle.com/code/aakashnain/tf-jax-tutorials-part-6-prng-in-jax/notebook)
 
 ```{code-cell} ipython3
 from numba import njit, prange

From ef9b8c9e0db487b7e6a30a83ecf74e13ba45c869 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Wed, 4 Jan 2023 21:26:47 +1100
Subject: [PATCH 07/15] remove 64-bit config

---
 lectures/kesten_processes.md | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 2212e3594..1d913bbf2 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -675,16 +675,13 @@ s_init = 1.0      # initial condition for each firm
 
 Here's one solution in [JAX](https://python-programming.quantecon.org/jax_intro.html). 
 
-First let's do a quick setup and check the backend for JAX
+First let's import the necessary modules and check the backend for JAX
 
 ```{code-cell} ipython3
 import jax
 import jax.numpy as jnp
 from jax import random
 
-# Use 64 bit floats with JAX in order to match NumPy/Numba code
-jax.config.update("jax_enable_x64", True)
-
 # Check if JAX is using GPU
 print(f"jax backend: {jax.devices()[0].platform}")
 ```

From b669a867cb311f0d4a435fc4133e0a79fda1f395 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Mon, 9 Jan 2023 12:34:47 +1100
Subject: [PATCH 08/15] update wording

---
 lectures/kesten_processes.md | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 1d913bbf2..8be772bee 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -730,7 +730,7 @@ Since we applied `jax.jit` on the function, it runs even faster when we call the
 %time data = generate_draws().block_until_ready()
 ```
 
-Now we produce the rank-size plot to check the distribution:
+Let's produce the rank-size plot and check the distribution:
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
@@ -746,7 +746,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) 
-to reduce the loop overhead in the compilation
+to reduce the loop overhead in the compilation of the jitted function
 
 ```{code-cell} ipython3
 from jax import lax
@@ -845,7 +845,9 @@ def generate_draws_numba(μ_a=-0.5,
 %time data = generate_draws_numba()
 ```
 
-Now we produce the rank-size plot:
+We can see that JAX and vectorization of the code have sped up the computation significantly compared to the Numba version.
+
+We produce the rank-size plot again using the data, and it shows the same pattern we saw before:
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()

From a2e621efba7cfcabb07c3d30b19a493dce598730 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Mon, 9 Jan 2023 22:17:33 +1100
Subject: [PATCH 09/15] add gpu tag

---
 lectures/kesten_processes.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 8be772bee..e8e7faf48 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -19,6 +19,14 @@ kernelspec:
 
 # Kesten Processes and Firm Dynamics
 
+```{note}
+This lecture is built using [hardware](status:machine-details) that
+has access to a GPU. This means that 
+
+1. the lecture might be significantly slower when running on your machine, and
+2. the code is well-suited to execution with [Google colab](https://colab.research.google.com/github/QuantEcon/lecture-python-programming.notebooks/blob/master/jax_intro.ipynb)
+```
+
 ```{index} single: Linear State Space Models
 ```
 
@@ -34,6 +42,12 @@ tags: [hide-output]
 ---
 !pip install quantecon
 !pip install --upgrade yfinance
+
+# If you are running on a machine without CUDA support, then run the line below:
+# !pip install --upgrade "jax[CPU]"
+
+# If you have CUDA support on your machine, then run the following:
+# !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 ## Overview

From 662c01014035c90b3ba707fac7a16b69f7d8b0ec Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Mon, 9 Jan 2023 22:44:29 +1100
Subject: [PATCH 10/15] update status

---
 lectures/status.md | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/lectures/status.md b/lectures/status.md
index b8a415a8f..7e7d4e821 100644
--- a/lectures/status.md
+++ b/lectures/status.md
@@ -16,5 +16,8 @@ This table contains the latest execution statistics.
 ```{nb-exec-table}
 ```
 
-These lectures are built on `linux` instances through `github actions` so are
-executed using the following [hardware specifications](https://docs.github.com/en/actions/using-github-hosted-runners/about-github-hosted-runners#supported-runners-and-hardware-resources)
\ No newline at end of file
+(status:machine-details)=
+
+These lectures are built on `linux` instances through `github actions`  and `amazon web services (aws)` to 
+enable access to a `gpu`. These lectures are built on a [p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/)
+that has access to `8 vcpu's`, a `V100 NVIDIA Tesla GPU`, and `61 Gb` of memory. 
\ No newline at end of file

From dc8a3d9984ee39de515fb0cdb34f8f74ce33f7e2 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Tue, 10 Jan 2023 23:37:50 +1100
Subject: [PATCH 11/15] update GPU warning

---
 lectures/kesten_processes.md | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index e8e7faf48..bfd2e3f0a 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -19,13 +19,16 @@ kernelspec:
 
 # Kesten Processes and Firm Dynamics
 
-```{note}
+````{admonition} GPU Warning
+:class: dropdown, warning
+
 This lecture is built using [hardware](status:machine-details) that
-has access to a GPU. This means that 
+has access to a GPU and uses JAX for GPU programming. This means that 
 
-1. the lecture might be significantly slower when running on your machine, and
-2. the code is well-suited to execution with [Google colab](https://colab.research.google.com/github/QuantEcon/lecture-python-programming.notebooks/blob/master/jax_intro.ipynb)
-```
+1. this lecture will be significantly slower when running on a machine **without CUDA support**.
+In this case, the code is well-suited to execution with Google Colab with runtime type set to GPU. [Click here](https://colab.research.google.com/github/QuantEcon/lecture-python.notebooks/blob/master/kesten_processes.ipynb) to load this lecture on Colab;
+2. if you have CUDA support for your machine, you should follow the [instructions](https://github.com/google/jax#installation) for JAX with GPU support.
+````
 
 ```{index} single: Linear State Space Models
 ```
@@ -43,11 +46,9 @@ tags: [hide-output]
 !pip install quantecon
 !pip install --upgrade yfinance
 
-# If you are running on a machine without CUDA support, then run the line below:
+!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+# If you do not have CUDA support, run the line below instead
 # !pip install --upgrade "jax[CPU]"
-
-# If you have CUDA support on your machine, then run the following:
-# !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 ## Overview

From dc98a37857913c0ab21150b3713d8a3b2dbba172 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Tue, 10 Jan 2023 23:39:13 +1100
Subject: [PATCH 12/15] add a comment

---
 lectures/kesten_processes.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index bfd2e3f0a..7ca15d895 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -47,7 +47,7 @@ tags: [hide-output]
 !pip install --upgrade yfinance
 
 !pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# If you do not have CUDA support, run the line below instead
+# If you do not have CUDA support on your machine, run the line below instead:
 # !pip install --upgrade "jax[CPU]"
 ```
 

From af6b8e7e08cca30fccaff346dcbb167ee47e3281 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Tue, 10 Jan 2023 23:43:22 +1100
Subject: [PATCH 13/15] change the default pip

---
 lectures/kesten_processes.md | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 7ca15d895..fdc7f7a27 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -45,10 +45,9 @@ tags: [hide-output]
 ---
 !pip install quantecon
 !pip install --upgrade yfinance
-
-!pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-# If you do not have CUDA support on your machine, run the line below instead:
-# !pip install --upgrade "jax[CPU]"
+# If your machine has CUDA support, please follow the guide in GPU Warning. 
+# Otherwise, run the line below:
+!pip install --upgrade "jax[CPU]"
 ```
 
 ## Overview

From 01412e31921dacbe1157361b6d2d1243eae01a00 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Sat, 14 Jan 2023 13:59:24 +1100
Subject: [PATCH 14/15] update GPU warning

---
 lectures/kesten_processes.md | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index fdc7f7a27..73cc88662 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -19,16 +19,15 @@ kernelspec:
 
 # Kesten Processes and Firm Dynamics
 
-````{admonition} GPU Warning
+```{admonition} GPU Warning
 :class: dropdown, warning
 
-This lecture is built using [hardware](status:machine-details) that
-has access to a GPU and uses JAX for GPU programming. This means that 
+This lecture is built using [hardware](status:machine-details) that has access to a GPU and uses JAX for GPU programming. As a result, the lecture will be slower when running on a machine without a GPU.
 
-1. this lecture will be significantly slower when running on a machine **without CUDA support**.
-In this case, the code is well-suited to execution with Google Colab with runtime type set to GPU. [Click here](https://colab.research.google.com/github/QuantEcon/lecture-python.notebooks/blob/master/kesten_processes.ipynb) to load this lecture on Colab;
-2. if you have CUDA support for your machine, you should follow the [instructions](https://github.com/google/jax#installation) for JAX with GPU support.
-````
+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.
+```
 
 ```{index} single: Linear State Space Models
 ```

From f28b8eb55b86d35bdb54053507e4db7040885d17 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <humphrey.yang@anu.edu.au>
Date: Tue, 17 Jan 2023 11:58:36 +1100
Subject: [PATCH 15/15] update gpu warning and reduce the length

---
 lectures/kesten_processes.md | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 73cc88662..8e4da868e 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -19,10 +19,10 @@ kernelspec:
 
 # Kesten Processes and Firm Dynamics
 
-```{admonition} GPU Warning
-:class: dropdown, warning
+```{admonition} GPU in use
+:class: warning
 
-This lecture is built using [hardware](status:machine-details) that has access to a GPU and uses JAX for GPU programming. As a result, the lecture will be slower when running on a machine without a GPU.
+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.
 
@@ -728,9 +728,12 @@ def generate_draws(μ_a=-0.5,
 
     # Perform the calculations in a vectorized manner for T periods
     for t in range(T):
+        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, 
-                             jnp.exp(e_random[t, :]), 
-                             jnp.exp(a_random[t, :]) * s[:, t] + jnp.exp(b_random[t, :])))
+                             exp_e,
+                             exp_a * s[:, t] + exp_b))
     
     return s[:, -1]