From 8260f21cab7a81972e5ecbacada31d434ecf404b Mon Sep 17 00:00:00 2001
From: Longye Tian <longye.tian@outlook.com>
Date: Wed, 20 Aug 2025 19:52:34 +1000
Subject: [PATCH 01/11] Update mccall_fitted_vfi.md

---
 lectures/mccall_fitted_vfi.md | 238 +++++++++++++++++-----------------
 1 file changed, 119 insertions(+), 119 deletions(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 2253aaf8e..087727b77 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -50,12 +50,13 @@ We will use the following imports:
 
 ```{code-cell} ipython3
 import matplotlib.pyplot as plt
-import numpy as np
-from numba import jit, float64
-from numba.experimental import jitclass
+import jax
+import jax.numpy as jnp
+from typing import NamedTuple
+import quantecon as qe
 ```
 
-## The Algorithm
+## The algorithm
 
 The model is the same as the McCall model with job separation we {doc}`studied before <mccall_model_with_separation>`, except that the wage offer distribution is continuous.
 
@@ -91,7 +92,7 @@ The function $q$ in {eq}`bell1mcmc` is the density of the wage offer distributio
 
 Its support is taken as equal to $\mathbb R_+$.
 
-### Value Function Iteration
+### Value function iteration
 
 In theory, we should now proceed as follows:
 
@@ -111,7 +112,7 @@ is to record its value $v'(w)$ for every $w \in \mathbb R_+$.
 
 Clearly, this is impossible.
 
-### Fitted Value Function Iteration
+### Fitted value function iteration
 
 What we will do instead is use **fitted value function iteration**.
 
@@ -145,21 +146,21 @@ This method
    {cite}`gordon1995stable` or {cite}`stachurski2008continuous`) and
 1. preserves useful shape properties such as monotonicity and concavity/convexity.
 
-Linear interpolation will be implemented using [numpy.interp](https://numpy.org/doc/stable/reference/generated/numpy.interp.html).
+Linear interpolation will be implemented using JAX's interpolation function `jnp.interp`.
 
 The next figure illustrates piecewise linear interpolation of an arbitrary
 function on grid points $0, 0.2, 0.4, 0.6, 0.8, 1$.
 
 ```{code-cell} python3
 def f(x):
-    y1 = 2 * np.cos(6 * x) + np.sin(14 * x)
+    y1 = 2 * jnp.cos(6 * x) + jnp.sin(14 * x)
     return y1 + 2.5
 
-c_grid = np.linspace(0, 1, 6)
-f_grid = np.linspace(0, 1, 150)
+c_grid = jnp.linspace(0, 1, 6)
+f_grid = jnp.linspace(0, 1, 150)
 
 def Af(x):
-    return np.interp(x, c_grid, f(c_grid))
+    return jnp.interp(x, c_grid, f(c_grid))
 
 fig, ax = plt.subplots()
 
@@ -175,8 +176,7 @@ plt.show()
 
 ## Implementation
 
-The first step is to build a jitted class for the McCall model with separation and
-a continuous wage offer distribution.
+The first step is to build a JAX-compatible structure for the McCall model with separation and a continuous wage offer distribution.
 
 We will take the utility function to be the log function for this application, with $u(c) = \ln c$.
 
@@ -184,114 +184,114 @@ We will adopt the lognormal distribution for wages, with $w = \exp(\mu + \sigma
 when $z$ is standard normal and $\mu, \sigma$ are parameters.
 
 ```{code-cell} python3
-@jit
 def lognormal_draws(n=1000, μ=2.5, σ=0.5, seed=1234):
-    np.random.seed(seed)
-    z = np.random.randn(n)
-    w_draws = np.exp(μ + σ * z)
+    key = jax.random.PRNGKey(seed)
+    z = jax.random.normal(key, (n,))
+    w_draws = jnp.exp(μ + σ * z)
     return w_draws
 ```
 
-Here's our class.
+Here's our model structure using a NamedTuple.
 
 ```{code-cell} python3
-mccall_data_continuous = [
-    ('c', float64),          # unemployment compensation
-    ('α', float64),          # job separation rate
-    ('β', float64),          # discount factor
-    ('w_grid', float64[:]),  # grid of points for fitted VFI
-    ('w_draws', float64[:])  # draws of wages for Monte Carlo
-]
-
-@jitclass(mccall_data_continuous)
-class McCallModelContinuous:
-
-    def __init__(self,
-                 c=1,
-                 α=0.1,
-                 β=0.96,
-                 grid_min=1e-10,
-                 grid_max=5,
-                 grid_size=100,
-                 w_draws=lognormal_draws()):
-
-        self.c, self.α, self.β = c, α, β
-
-        self.w_grid = np.linspace(grid_min, grid_max, grid_size)
-        self.w_draws = w_draws
-
-    def update(self, v, d):
-
-        # Simplify names
-        c, α, β = self.c, self.α, self.β
-        w = self.w_grid
-        u = lambda x: np.log(x)
-
-        # Interpolate array represented value function
-        vf = lambda x: np.interp(x, w, v)
-
-        # Update d using Monte Carlo to evaluate integral
-        d_new = np.mean(np.maximum(vf(self.w_draws), u(c) + β * d))
-
-        # Update v
-        v_new = u(w) + β * ((1 - α) * v + α * d)
-
-        return v_new, d_new
+class McCallModelContinuous(NamedTuple):
+    c: float              # unemployment compensation
+    α: float              # job separation rate
+    β: float              # discount factor
+    w_grid: jnp.ndarray   # grid of points for fitted VFI
+    w_draws: jnp.ndarray  # draws of wages for Monte Carlo
+
+def create_mccall_model(c=1,
+                        α=0.1,
+                        β=0.96,
+                        grid_min=1e-10,
+                        grid_max=5,
+                        grid_size=100,
+                        μ=2.5,
+                        σ=0.5,
+                        mc_size=1000,
+                        seed=1234):
+    """Factory function to create a McCall model instance."""
+    w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
+    w_grid = jnp.linspace(grid_min, grid_max, grid_size)
+    return McCallModelContinuous(c=c, α=α, β=β, w_grid=w_grid, w_draws=w_draws)
+
+@jax.jit
+def update(model, v, d):
+    """Update value function and continuation value."""
+    # Unpack model parameters
+    c, α, β, w_grid, w_draws = model
+    u = jnp.log
+    
+    # Interpolate array represented value function
+    vf = lambda x: jnp.interp(x, w_grid, v)
+    
+    # Update d using Monte Carlo to evaluate integral
+    d_new = jnp.mean(jnp.maximum(vf(w_draws), u(c) + β * d))
+    
+    # Update v
+    v_new = u(w_grid) + β * ((1 - α) * v + α * d)
+    
+    return v_new, d_new
 ```
 
 We then return the current iterate as an approximate solution.
 
 ```{code-cell} python3
-@jit
-def solve_model(mcm, tol=1e-5, max_iter=2000):
+@jax.jit
+def solve_model(model, tol=1e-5, max_iter=2000):
     """
     Iterates to convergence on the Bellman equations
 
-    * mcm is an instance of McCallModel
+    * model is an instance of McCallModelContinuous
     """
-
-    v = np.ones_like(mcm.w_grid)    # Initial guess of v
-    d = 1                           # Initial guess of d
-    i = 0
-    error = tol + 1
-
-    while error > tol and i < max_iter:
-        v_new, d_new = mcm.update(v, d)
-        error_1 = np.max(np.abs(v_new - v))
-        error_2 = np.abs(d_new - d)
-        error = max(error_1, error_2)
-        v = v_new
-        d = d_new
-        i += 1
-
-    return v, d
+    
+    # Initial guesses
+    v = jnp.ones_like(model.w_grid)
+    d = 1.0
+    
+    def body_fun(state):
+        v, d, i, error = state
+        v_new, d_new = update(model, v, d)
+        error_1 = jnp.max(jnp.abs(v_new - v))
+        error_2 = jnp.abs(d_new - d)
+        error = jnp.maximum(error_1, error_2)
+        return v_new, d_new, i + 1, error
+    
+    def cond_fun(state):
+        v, d, i, error = state
+        return (error > tol) & (i < max_iter)
+    
+    initial_state = (v, d, 0, tol + 1)
+    v_final, d_final, _, _ = jax.lax.while_loop(cond_fun, body_fun, initial_state)
+    
+    return v_final, d_final
 ```
 
 Here's a function `compute_reservation_wage` that takes an instance of `McCallModelContinuous`
 and returns the associated reservation wage.
 
-If $v(w) < h$ for all $w$, then the function returns np.inf
+If $v(w) < h$ for all $w$, then the function returns `jnp.inf`
 
 ```{code-cell} python3
-@jit
-def compute_reservation_wage(mcm):
+@jax.jit
+def compute_reservation_wage(model):
     """
     Computes the reservation wage of an instance of the McCall model
     by finding the smallest w such that v(w) >= h.
 
-    If no such w exists, then w_bar is set to np.inf.
+    If no such w exists, then w_bar is set to inf.
     """
-    u = lambda x: np.log(x)
-
-    v, d = solve_model(mcm)
-    h = u(mcm.c) + mcm.β * d
-
-    w_bar = np.inf
-    for i, wage in enumerate(mcm.w_grid):
-        if v[i] > h:
-            w_bar = wage
-            break
-
+    c, α, β, w_grid, w_draws = model
+    u = jnp.log
+    
+    v, d = solve_model(model)
+    h = u(c) + β * d
+    
+    # Find the first wage where v(w) >= h
+    indices = jnp.where(v >= h, size=1, fill_value=-1)
+    w_bar = jnp.where(indices[0] >= 0, w_grid[indices[0]], jnp.inf)
+    
     return w_bar
 ```
 
@@ -305,7 +305,7 @@ The exercises ask you to explore the solution and how it changes with parameters
 Use the code above to explore what happens to the reservation wage when the wage parameter $\mu$
 changes.
 
-Use the default parameters and $\mu$ in `mu_vals = np.linspace(0.0, 2.0, 15)`.
+Use the default parameters and $\mu$ in `mu_vals = jnp.linspace(0.0, 2.0, 15)`.
 
 Is the impact on the reservation wage as you expected?
 ```
@@ -317,21 +317,18 @@ Is the impact on the reservation wage as you expected?
 Here is one solution
 
 ```{code-cell} python3
-mcm = McCallModelContinuous()
-mu_vals = np.linspace(0.0, 2.0, 15)
-w_bar_vals = np.empty_like(mu_vals)
-
-fig, ax = plt.subplots()
+def compute_res_wage_given_mu(μ):
+    model = create_mccall_model(μ=μ)
+    w_bar = compute_reservation_wage(model)
+    return w_bar
 
-for i, m in enumerate(mu_vals):
-    mcm.w_draws = lognormal_draws(μ=m)
-    w_bar = compute_reservation_wage(mcm)
-    w_bar_vals[i] = w_bar
+mu_vals = jnp.linspace(0.0, 2.0, 15)
+w_bar_vals = jax.vmap(compute_res_wage_given_mu)(mu_vals)
 
+fig, ax = plt.subplots()
 ax.set(xlabel='mean', ylabel='reservation wage')
 ax.plot(mu_vals, w_bar_vals, label=r'$\bar w$ as a function of $\mu$')
 ax.legend()
-
 plt.show()
 ```
 
@@ -354,11 +351,11 @@ support.
 
 (This is a form of *mean-preserving spread*.)
 
-Use `s_vals = np.linspace(1.0, 2.0, 15)` and `m = 2.0`.
+Use `s_vals = jnp.linspace(1.0, 2.0, 15)` and `m = 2.0`.
 
 State how you expect the reservation wage to vary with $s$.
 
-Now compute it.  Is this as you expected?
+Now compute it - is this as you expected?
 ```
 
 ```{solution-start} mfv_ex2
@@ -368,23 +365,26 @@ Now compute it.  Is this as you expected?
 Here is one solution
 
 ```{code-cell} python3
-mcm = McCallModelContinuous()
-s_vals = np.linspace(1.0, 2.0, 15)
-m = 2.0
-w_bar_vals = np.empty_like(s_vals)
-
-fig, ax = plt.subplots()
-
-for i, s in enumerate(s_vals):
+def compute_res_wage_given_s(s, m=2.0, seed=1234):
     a, b = m - s, m + s
-    mcm.w_draws = np.random.uniform(low=a, high=b, size=10_000)
-    w_bar = compute_reservation_wage(mcm)
-    w_bar_vals[i] = w_bar
+    key = jax.random.PRNGKey(seed)
+    uniform_draws = jax.random.uniform(key, shape=(10_000,), minval=a, maxval=b)
+    # Create model with default parameters but replace wage draws
+    model = create_mccall_model()
+    model = model._replace(w_draws=uniform_draws)
+    w_bar = compute_reservation_wage(model)
+    return w_bar
 
+s_vals = jnp.linspace(1.0, 2.0, 15)
+# Use vmap with different seeds for each s value
+seeds = jnp.arange(len(s_vals))
+compute_vectorized = jax.vmap(compute_res_wage_given_s, in_axes=(0, None, 0))
+w_bar_vals = compute_vectorized(s_vals, 2.0, seeds)
+
+fig, ax = plt.subplots()
 ax.set(xlabel='volatility', ylabel='reservation wage')
 ax.plot(s_vals, w_bar_vals, label=r'$\bar w$ as a function of wage volatility')
 ax.legend()
-
 plt.show()
 ```
 

From b3993a056fee756c5c8508834295f44464cdddff Mon Sep 17 00:00:00 2001
From: Humphrey Yang <u6474961@anu.edu.au>
Date: Wed, 20 Aug 2025 22:02:28 +1000
Subject: [PATCH 02/11] minor fixes

---
 lectures/mccall_fitted_vfi.md | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 087727b77..3b5d91ecb 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -58,7 +58,7 @@ import quantecon as qe
 
 ## The algorithm
 
-The model is the same as the McCall model with job separation we {doc}`studied before <mccall_model_with_separation>`, except that the wage offer distribution is continuous.
+The model is the same as the McCall model with job separation that we {doc}`studied before <mccall_model_with_separation>`, except that the wage offer distribution is continuous.
 
 We are going to start with the two Bellman equations we obtained for the model with job separation after {ref}`a simplifying transformation <ast_mcm>`.
 
@@ -83,9 +83,9 @@ v(w) = u(w) + \beta
 
 The unknowns here are the function $v$ and the scalar $d$.
 
-The difference between these and the pair of Bellman equations we previously worked on are
+The differences between these and the pair of Bellman equations we previously worked on are
 
-1. in {eq}`bell1mcmc`, what used to be a sum over a finite number of wage values is an integral over an infinite set.
+1. In {eq}`bell1mcmc`, what used to be a sum over a finite number of wage values is an integral over an infinite set.
 1. The function $v$ in {eq}`bell2mcmc` is defined over all $w \in \mathbb R_+$.
 
 The function $q$ in {eq}`bell1mcmc` is the density of the wage offer distribution.
@@ -107,7 +107,7 @@ The iterates of the value function can neither be calculated exactly nor stored
 
 To see the issue, consider {eq}`bell2mcmc`.
 
-Even if $v$ is a known function,  the only way to store its update $v'$
+Even if $v$ is a known function, the only way to store its update $v'$
 is to record its value $v'(w)$ for every $w \in \mathbb R_+$.
 
 Clearly, this is impossible.
@@ -142,7 +142,7 @@ One good choice from both respects is continuous piecewise linear interpolation.
 
 This method
 
-1. combines well with value function iteration (see., e.g.,
+1. combines well with value function iteration (see, e.g.,
    {cite}`gordon1995stable` or {cite}`stachurski2008continuous`) and
 1. preserves useful shape properties such as monotonicity and concavity/convexity.
 

From 23ded25519625298ff35ebe802ad4e31ecd51a48 Mon Sep 17 00:00:00 2001
From: Longye Tian <133612246+longye-tian@users.noreply.github.com>
Date: Thu, 21 Aug 2025 10:39:22 +1000
Subject: [PATCH 03/11] Update lectures/mccall_fitted_vfi.md

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>
---
 lectures/mccall_fitted_vfi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 3b5d91ecb..276663694 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -317,7 +317,7 @@ Is the impact on the reservation wage as you expected?
 Here is one solution
 
 ```{code-cell} python3
-def compute_res_wage_given_mu(μ):
+def compute_res_wage_given_μ(μ):
     model = create_mccall_model(μ=μ)
     w_bar = compute_reservation_wage(model)
     return w_bar

From 2a19ed451aa033fbf9fceba95a2e0daec3e83007 Mon Sep 17 00:00:00 2001
From: Longye Tian <133612246+longye-tian@users.noreply.github.com>
Date: Thu, 21 Aug 2025 10:39:29 +1000
Subject: [PATCH 04/11] Update lectures/mccall_fitted_vfi.md

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>
---
 lectures/mccall_fitted_vfi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 276663694..7341fe1ce 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -322,7 +322,7 @@ def compute_res_wage_given_μ(μ):
     w_bar = compute_reservation_wage(model)
     return w_bar
 
-mu_vals = jnp.linspace(0.0, 2.0, 15)
+μ_vals = jnp.linspace(0.0, 2.0, 15)
 w_bar_vals = jax.vmap(compute_res_wage_given_mu)(mu_vals)
 
 fig, ax = plt.subplots()

From ee86afa41d6498c50c869c57e333df60f4b3da1f Mon Sep 17 00:00:00 2001
From: Longye Tian <133612246+longye-tian@users.noreply.github.com>
Date: Thu, 21 Aug 2025 10:39:35 +1000
Subject: [PATCH 05/11] Update lectures/mccall_fitted_vfi.md

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>
---
 lectures/mccall_fitted_vfi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 7341fe1ce..47cb15fa1 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -305,7 +305,7 @@ The exercises ask you to explore the solution and how it changes with parameters
 Use the code above to explore what happens to the reservation wage when the wage parameter $\mu$
 changes.
 
-Use the default parameters and $\mu$ in `mu_vals = jnp.linspace(0.0, 2.0, 15)`.
+Use the default parameters and $\mu$ in `μ_vals = jnp.linspace(0.0, 2.0, 15)`.
 
 Is the impact on the reservation wage as you expected?
 ```

From 398e26666d523b4ff4fcab6f1e3a3b4681a961bc Mon Sep 17 00:00:00 2001
From: Longye Tian <133612246+longye-tian@users.noreply.github.com>
Date: Thu, 21 Aug 2025 10:40:05 +1000
Subject: [PATCH 06/11] Update lectures/mccall_fitted_vfi.md

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>
---
 lectures/mccall_fitted_vfi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 47cb15fa1..2c5efb989 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -259,7 +259,7 @@ def solve_model(model, tol=1e-5, max_iter=2000):
         return v_new, d_new, i + 1, error
     
     def cond_fun(state):
-        v, d, i, error = state
+        _, _, i, error = state
         return (error > tol) & (i < max_iter)
     
     initial_state = (v, d, 0, tol + 1)

From 4f0b4e64d54454abf868b2b223642afed0049891 Mon Sep 17 00:00:00 2001
From: Longye Tian <longye.tian@outlook.com>
Date: Thu, 21 Aug 2025 11:03:19 +1000
Subject: [PATCH 07/11] Update mccall_fitted_vfi.md

---
 lectures/mccall_fitted_vfi.md | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 2c5efb989..fbd445483 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -210,9 +210,13 @@ def create_mccall_model(c=1,
                         μ=2.5,
                         σ=0.5,
                         mc_size=1000,
-                        seed=1234):
+                        seed=1234,
+                        w_draws=None):
     """Factory function to create a McCall model instance."""
-    w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
+    if w_draws is None:
+        # Generate wage draws if not provided
+         w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
+
     w_grid = jnp.linspace(grid_min, grid_max, grid_size)
     return McCallModelContinuous(c=c, α=α, β=β, w_grid=w_grid, w_draws=w_draws)
 
@@ -323,11 +327,11 @@ def compute_res_wage_given_μ(μ):
     return w_bar
 
 μ_vals = jnp.linspace(0.0, 2.0, 15)
-w_bar_vals = jax.vmap(compute_res_wage_given_mu)(mu_vals)
+w_bar_vals = jax.vmap(compute_res_wage_given_μ)(μ_vals)
 
 fig, ax = plt.subplots()
 ax.set(xlabel='mean', ylabel='reservation wage')
-ax.plot(mu_vals, w_bar_vals, label=r'$\bar w$ as a function of $\mu$')
+ax.plot(μ_vals, w_bar_vals, label=r'$\bar w$ as a function of $\mu$')
 ax.legend()
 plt.show()
 ```
@@ -370,8 +374,7 @@ def compute_res_wage_given_s(s, m=2.0, seed=1234):
     key = jax.random.PRNGKey(seed)
     uniform_draws = jax.random.uniform(key, shape=(10_000,), minval=a, maxval=b)
     # Create model with default parameters but replace wage draws
-    model = create_mccall_model()
-    model = model._replace(w_draws=uniform_draws)
+    model = create_mccall_model(w_draws=uniform_draws)
     w_bar = compute_reservation_wage(model)
     return w_bar
 

From ac020a079b0f61e85af7e48740aeb720182e72ed Mon Sep 17 00:00:00 2001
From: Longye Tian <133612246+longye-tian@users.noreply.github.com>
Date: Thu, 21 Aug 2025 15:59:25 +1000
Subject: [PATCH 08/11] Update lectures/mccall_fitted_vfi.md

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 lectures/mccall_fitted_vfi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index fbd445483..039d823a3 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -215,7 +215,7 @@ def create_mccall_model(c=1,
     """Factory function to create a McCall model instance."""
     if w_draws is None:
         # Generate wage draws if not provided
-         w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
+        w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
 
     w_grid = jnp.linspace(grid_min, grid_max, grid_size)
     return McCallModelContinuous(c=c, α=α, β=β, w_grid=w_grid, w_draws=w_draws)

From 95f5ef672fdb826c44ec688ddc8963b14c2ef6d1 Mon Sep 17 00:00:00 2001
From: Humphrey Yang <u6474961@anu.edu.au>
Date: Thu, 21 Aug 2025 16:49:02 +1000
Subject: [PATCH 09/11] testing cpu run time

---
 lectures/mccall_fitted_vfi.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 039d823a3..83c400f7a 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -54,6 +54,9 @@ import jax
 import jax.numpy as jnp
 from typing import NamedTuple
 import quantecon as qe
+
+# Set JAX to use CPU
+jax.config.update('jax_platform_name', 'cpu')
 ```
 
 ## The algorithm

From 96508697a44120cc3bc07d99d2db9f2367d81882 Mon Sep 17 00:00:00 2001
From: Matt McKay <mmcky@users.noreply.github.com>
Date: Fri, 22 Aug 2025 14:24:24 +1000
Subject: [PATCH 10/11] Update lectures/mccall_fitted_vfi.md

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

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 83c400f7a..312e9555b 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -217,6 +217,7 @@ def create_mccall_model(c=1,
                         w_draws=None):
     """Factory function to create a McCall model instance."""
     if w_draws is None:
+
         # Generate wage draws if not provided
         w_draws = lognormal_draws(n=mc_size, μ=μ, σ=σ, seed=seed)
 

From 438eda2f877a27b2e0f7e4eae1486e6a6705f681 Mon Sep 17 00:00:00 2001
From: Matt McKay <mmcky@users.noreply.github.com>
Date: Fri, 22 Aug 2025 14:25:14 +1000
Subject: [PATCH 11/11] Apply suggestions from code review

pep8 compliance
---
 lectures/mccall_fitted_vfi.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/lectures/mccall_fitted_vfi.md b/lectures/mccall_fitted_vfi.md
index 312e9555b..b24867836 100644
--- a/lectures/mccall_fitted_vfi.md
+++ b/lectures/mccall_fitted_vfi.md
@@ -227,6 +227,7 @@ def create_mccall_model(c=1,
 @jax.jit
 def update(model, v, d):
     """Update value function and continuation value."""
+
     # Unpack model parameters
     c, α, β, w_grid, w_draws = model
     u = jnp.log
@@ -377,12 +378,14 @@ def compute_res_wage_given_s(s, m=2.0, seed=1234):
     a, b = m - s, m + s
     key = jax.random.PRNGKey(seed)
     uniform_draws = jax.random.uniform(key, shape=(10_000,), minval=a, maxval=b)
+
     # Create model with default parameters but replace wage draws
     model = create_mccall_model(w_draws=uniform_draws)
     w_bar = compute_reservation_wage(model)
     return w_bar
 
 s_vals = jnp.linspace(1.0, 2.0, 15)
+
 # Use vmap with different seeds for each s value
 seeds = jnp.arange(len(s_vals))
 compute_vectorized = jax.vmap(compute_res_wage_given_s, in_axes=(0, None, 0))