From a2cba7f07b385c43cc44a2955d8de445eb0220d2 Mon Sep 17 00:00:00 2001
From: John Stachurski <john.stachurski@gmail.com>
Date: Tue, 24 Jun 2025 15:36:55 +0900
Subject: [PATCH 1/5] misc

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

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 6948c63..0cbbb27 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -190,12 +190,12 @@ class Firm(NamedTuple):
     μ_e:   float = 0.0
     σ_e:   float = 0.5
     s_bar: float = 1.0
-
-#
-# Here's code to update a cross-section of firms according to the dynamics in
-# [](firm_dynam_ee).
 ```
 
+
+Here's code to update a cross-section of firms according to the dynamics in
+ [](firm_dynam_ee).
+
 ```{code-cell} ipython3
 @jax.jit
 def update_cross_section(s, a, b, e, firm):
@@ -250,7 +250,6 @@ data = generate_cross_section(firm).block_until_ready()
 toc()
 ```
 
-
 Let's produce the rank-size plot and check the distribution:
 
 ```{code-cell} ipython3
@@ -311,10 +310,10 @@ def generate_cross_section_lax(
         0, T, update_cross_section, initial_state
     )
     return final_s
-
-# Let's see if we got any speed gain
 ```
 
+Let's see if we get any speed gain
+
 ```{code-cell} ipython3
 tic()
 data = generate_cross_section_lax(firm).block_until_ready()
@@ -339,16 +338,20 @@ ax.set_ylabel("log size")
 
 plt.show()
 
-#
-# If the time horizon is not too large, we can also try generating all shocks at
-# once.
-#
-# Note, however, that this approach consumes more memory, as we need to have to
-# store large matrices of random draws
-#
-# Hence the code below will fail due to out-of-memory errors when `T` and `M` are large.
 ```
 
+## Exercises
+
+Try writing an alternative version of `generate_cross_section_lax()` where the entire sequence of random draws is generated at once, so that all of `a`, `b`, and `e` are of shape `(T, M)`.
+
+(The `update_cross_section()` function should not generate any random numbers.)
+
+Does it improve the runtime?
+
+What are the pros and cons of this approach.
+
+*Solution*
+
 ```{code-cell} ipython3
 @jax.jit
 def generate_cross_section_lax(
@@ -393,6 +396,8 @@ data = generate_cross_section_lax(firm).block_until_ready()
 toc()
 ```
 
-This second method might be slightly faster in some cases but in general the
+This method might be faster in some cases but in general the
 relative speed will depend on the size of the cross-section and the length of
 the simulation paths.
+
+Also, this method is far more memory intensive.

From f0675a3f692ccd8d855d5fd5e4f428a8b2f874b2 Mon Sep 17 00:00:00 2001
From: mmcky <mamckay@gmail.com>
Date: Tue, 24 Jun 2025 20:28:47 +1000
Subject: [PATCH 2/5] @mmcky edits

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

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 0cbbb27..939a2be 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -40,9 +40,8 @@ In addition to JAX and Anaconda, this lecture will need the following libraries:
 This lecture describes Kesten processes, which are an important class of
 stochastic processes, and an application of firm dynamics.
 
-The lecture draws on [an earlier QuantEcon
-lecture](https://python.quantecon.org/kesten_processes.html), which uses Numba
-to accelerate the computations.
+The lecture draws on [an earlier QuantEcon lecture](https://python.quantecon.org/kesten_processes.html), 
+which uses Numba to accelerate the computations.
 
 In that earlier lecture you can find a more detailed discussion of the concepts involved.
 
@@ -137,10 +136,8 @@ We now study the implications of this specification.
 
 #### Heavy tails
 
-If the conditions of the [Kesten--Goldie
-Theorem](https://python.quantecon.org/kesten_processes.html#the-kestengoldie-theorem)
-are satisfied, then {eq}`firm_dynam` implies that the firm size distribution
-will have Pareto tails.
+If the conditions of the [Kesten--Goldie Theorem](https://python.quantecon.org/kesten_processes.html#the-kestengoldie-theorem)
+are satisfied, then {eq}`firm_dynam` implies that the firm size distribution will have Pareto tails.
 
 This matches empirical findings across many data sets.
 
@@ -192,9 +189,8 @@ class Firm(NamedTuple):
     s_bar: float = 1.0
 ```
 
-
 Here's code to update a cross-section of firms according to the dynamics in
- [](firm_dynam_ee).
+[](firm_dynam_ee).
 
 ```{code-cell} ipython3
 @jax.jit
@@ -270,7 +266,7 @@ The plot produces a straight line, consistent with a Pareto tail.
 
 We did not JIT-compile the `for` loop above because
 acceleration of outer loops makes relatively little difference terms of
-   compute time.
+compute time.
 
 However, to maximize performance, let's try squeezing out a bit more speed
 by replacing the `for` loop with
@@ -342,6 +338,10 @@ plt.show()
 
 ## Exercises
 
+```{exercise-start}
+:label: kp_ex1
+```
+
 Try writing an alternative version of `generate_cross_section_lax()` where the entire sequence of random draws is generated at once, so that all of `a`, `b`, and `e` are of shape `(T, M)`.
 
 (The `update_cross_section()` function should not generate any random numbers.)
@@ -350,7 +350,11 @@ Does it improve the runtime?
 
 What are the pros and cons of this approach.
 
-*Solution*
+```{exercise-end}
+```
+
+```{solution-start} kp_ex1
+```
 
 ```{code-cell} ipython3
 @jax.jit
@@ -401,3 +405,6 @@ relative speed will depend on the size of the cross-section and the length of
 the simulation paths.
 
 Also, this method is far more memory intensive.
+
+```{solution-end}
+```
\ No newline at end of file

From cdacdf88bb24e91c2438785724e5457ea5ee6c36 Mon Sep 17 00:00:00 2001
From: mmcky <mamckay@gmail.com>
Date: Tue, 24 Jun 2025 20:38:45 +1000
Subject: [PATCH 3/5] enable dropdown

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

diff --git a/lectures/kesten_processes.md b/lectures/kesten_processes.md
index 939a2be..03750e4 100644
--- a/lectures/kesten_processes.md
+++ b/lectures/kesten_processes.md
@@ -354,6 +354,7 @@ What are the pros and cons of this approach.
 ```
 
 ```{solution-start} kp_ex1
+:class: dropdown
 ```
 
 ```{code-cell} ipython3

From 208ded23d2bc6e7d8900f8ed251a8facc036c97a Mon Sep 17 00:00:00 2001
From: mmcky <mamckay@gmail.com>
Date: Tue, 24 Jun 2025 20:55:03 +1000
Subject: [PATCH 4/5] upgrade CUDANN to 9.10.2

---
 .github/workflows/ci.yml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 5bb8bd0..47d1801 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -7,6 +7,13 @@ jobs:
       - uses: actions/checkout@v4
         with:
           ref: ${{ github.event.pull_request.head.sha }}
+      - name: Upgrade CUDANN
+        shell: bash -l {0}
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get -y install cudnn-cuda-12
       - name: Setup Anaconda
         uses: conda-incubator/setup-miniconda@v3
         with:

From 59f3c50a0a84f09e5828ea79aa0b06db559e33ae Mon Sep 17 00:00:00 2001
From: mmcky <mamckay@gmail.com>
Date: Tue, 24 Jun 2025 21:07:15 +1000
Subject: [PATCH 5/5] MAINT: upgrade cache and publish workflows with CUDANN
 upgrade

---
 .github/workflows/cache.yml   | 7 +++++++
 .github/workflows/publish.yml | 7 +++++++
 2 files changed, 14 insertions(+)

diff --git a/.github/workflows/cache.yml b/.github/workflows/cache.yml
index b1111b7..89003fc 100644
--- a/.github/workflows/cache.yml
+++ b/.github/workflows/cache.yml
@@ -11,6 +11,13 @@ jobs:
       - uses: actions/checkout@v4
         with:
           ref: ${{ github.event.pull_request.head.sha }}
+      - name: Upgrade CUDANN
+        shell: bash -l {0}
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get -y install cudnn-cuda-12
       - name: Setup Anaconda
         uses: conda-incubator/setup-miniconda@v3
         with:
diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index c72fff2..8e195c3 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -10,6 +10,13 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
+      - name: Upgrade CUDANN
+        shell: bash -l {0}
+        run: |
+          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+          sudo dpkg -i cuda-keyring_1.1-1_all.deb
+          sudo apt-get update
+          sudo apt-get -y install cudnn-cuda-12
       - name: Setup Anaconda
         uses: conda-incubator/setup-miniconda@v3
         with: