diff --git a/lectures/_config.yml b/lectures/_config.yml
index fb896570..4f5b4445 100644
--- a/lectures/_config.yml
+++ b/lectures/_config.yml
@@ -35,6 +35,7 @@ latex:
 sphinx:
   extra_extensions: [sphinx_multitoc_numbering, sphinxext.rediraffe, sphinxcontrib.youtube, sphinx.ext.todo, sphinx_exercise, sphinx_togglebutton, sphinx_tojupyter]
   config:
+    bibtex_reference_style: author_year
     # myst-nb config
     nb_render_image_options:
       width: 80%
@@ -81,8 +82,8 @@ sphinx:
       twitter: quantecon
       twitter_logo_url: https://assets.quantecon.org/img/qe-twitter-logo.png
       og_logo_url: https://assets.quantecon.org/img/qe-og-logo.png
-      description: This website presents a set of lectures on quantitative economic modeling, designed and written by Thomas J. Sargent and John Stachurski.
-      keywords: Python, QuantEcon, Quantitative Economics, Economics, Sloan, Alfred P. Sloan Foundation, Tom J. Sargent, John Stachurski
+      description: This website presents a set of lectures on quantitative economic modeling using JAX, designed and written by Thomas J. Sargent and John Stachurski.
+      keywords: Python, JAX, QuantEcon, Quantitative Economics, Economics, Sloan, Alfred P. Sloan Foundation, Tom J. Sargent, John Stachurski
       analytics:
         google_analytics_id: G-K1NYBSC1CZ
       launch_buttons:
@@ -95,7 +96,7 @@ sphinx:
     mathjax_path: https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
     rediraffe_redirects:
       index_toc.md: intro.md
-    tojupyter_static_file_path: ["source/_static", "_static"]
+    tojupyter_static_file_path: ["_static"]
     tojupyter_target_html: true
     tojupyter_urlpath: "https://jax.quantecon.org/"
     tojupyter_image_urlpath: "https://jax.quantecon.org/_static/"
diff --git a/lectures/_static/quant-econ.bib b/lectures/_static/quant-econ.bib
index dd686944..f5dc4db8 100644
--- a/lectures/_static/quant-econ.bib
+++ b/lectures/_static/quant-econ.bib
@@ -2503,4 +2503,17 @@ @article{Are08
   pages   = {690--712},
   title   = {{Default risk and income fluctuations in emerging economies}},
   year    = {2008}
-}
\ No newline at end of file
+}
+
+@article{Bianchi2011,
+  Author = {Bianchi, Javier},
+  Title = {Overborrowing and Systemic Externalities in the Business Cycle},
+  Journal = {American Economic Review},
+  Volume = {101},
+  Number = {7},
+  Year = {2011},
+  Month = {December},
+  Pages = {3400-3426},
+  DOI = {10.1257/aer.101.7.3400},
+  URL = {https://www.aeaweb.org/articles?id=10.1257/aer.101.7.3400}
+}
diff --git a/lectures/_toc.yml b/lectures/_toc.yml
index 19f2f4a6..042ad8a6 100644
--- a/lectures/_toc.yml
+++ b/lectures/_toc.yml
@@ -21,15 +21,21 @@ parts:
 - caption: Dynamic Programming
   numbered: true
   chapters:
+  - file: job_search
   - file: opt_savings_1
   - file: opt_savings_2
   - file: short_path
   - file: opt_invest
   - file: inventory_ssd
   - file: ifp_egm
+  - file: cake_eating_numerical
+- caption: Macroeconomic Models
+  numbered: true
+  chapters:
   - file: arellano
   - file: aiyagari_jax
-  - file: cake_eating_numerical
+  - file: hopenhayn
+  - file: overborrowing
 - caption: Data and Empirics
   numbered: true
   chapters:
diff --git a/lectures/hopenhayn.md b/lectures/hopenhayn.md
new file mode 100644
index 00000000..5995d092
--- /dev/null
+++ b/lectures/hopenhayn.md
@@ -0,0 +1,878 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.16.1
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# The Hopenhayn Entry-Exit Model
+
+```{include} _admonition/gpu.md
+```
+
+## Outline
+
+The Hopenhayn (1992, ECMA) entry-exit model is a highly influential (partial equilibrium) heterogeneous agent model where 
+
+* the agents receiving idiosyncratic shocks are firms, and
+* the model produces a non-trivial firm size distribution and a range of predictions for firm dynamics.
+
+We are going to study an extension of the basic model that has unbounded productivity.
+
+Among other things, this allows us to match the heavy tail observed in the firm size distribution data.
+
+In addition to what’s in Anaconda, this lecture will need the following libraries:
+
+```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install quantecon
+```
+
+We will use the following imports.
+
+```{code-cell} ipython3
+import jax.numpy as jnp
+import jax
+import quantecon as qe
+import matplotlib.pyplot as plt
+from collections import namedtuple
+```
+
+## The Model
+
+
+### Environment
+
+There is a single good produced by a continuum of competitive firms.
+
+The productivity $\varphi_t = \varphi_t^i$ of each firm evolves randomly on $\mathbb R_+$.
+
+* the $i$ superscript indicates the $i$-th firm and we drop it in what follows
+
+A firm with productivity $\varphi$ facing output price $p$ and wage rate $w$ earns current profits
+
+$$
+    \pi(\varphi, p) = p \varphi n^\theta - c  - wn
+$$
+
+Here 
+
+- $\theta \in (0,1)$ is a productivity parameter, 
+- $c$ is a fixed cost and 
+- $n$ is labor input.
+
+Maximizing over $n$ leads to
+
+$$
+    \pi(\varphi, p) = (1 - \theta) (p \varphi)^\eta
+    \left(\frac{\theta}{w} \right)^{\theta / \eta} - c
+$$
+
+where $\eta := 1/(1-\theta)$.
+
+Output is
+
+$$
+    q(\varphi, p) = \varphi^\eta
+    \left( \frac{p \theta}{w}\right)^{\theta / \eta}
+$$  
+
+Productivity of incumbents is independent across firms and grows according to 
+
+
+$$ \varphi_{t+1} \sim \Gamma(\varphi_t, d \varphi') $$
+
+Here $\Gamma$ is a Markov transition kernel, so that $\Gamma(\varphi, d
+                                                             \varphi')$ is the
+distribution over $\mathbb R_+$ given $\varphi \in \mathbb R_+$.
+
+For example,
+
+$$ 
+   \int \pi(\varphi', p) \Gamma(\varphi, d \varphi')
+    = \mathbb E[ \pi(\varphi_{t+1}, p) \,|\, \varphi_t = \varphi ]
+$$
+
+New entrants obtain their productivity draw from a fixed distribution $\gamma$.
+
+
+Demand is 
+
+$$D(p) = 1/p $$
+
+
+
+### Firm decisions
+
+The intertemporal firm decision problem is choosing when to exit.
+
+The timing is:
+
+1. produce and receive current profits $\pi(\varphi, p)$
+2. decide whether to continue or exit.
+
+Scrap value is set to zero and the discount rate is $\beta \in (0, 1)$.
+
+The firm makes their stop-continue choice by first solving for the value function, which is the unique solution in a class of continuous functions $\mathscr C$ (details omitted) to
+
+\begin{equation}
+    v(\varphi, p) = \pi(\varphi, p) + \beta 
+        \max \left\{ 0, \int v(\varphi', p) \Gamma(\varphi, d \varphi') \right\}.
+\end{equation}
+
+Let $\bar v$ be the unique solution to this functional equation.
+
+
+
+Given $\bar v$, we let $\bar \varphi$ be the exit threshold function defined by
+%
+\begin{equation}
+    \bar \varphi(p) := 
+    \min
+    \left\{ 
+        \varphi \geq 0  
+        \;\; \big| \;
+        \int \bar v(\varphi', p) \, \Gamma(\varphi, d \varphi') \geq 0
+    \right\}.
+\end{equation}
+%
+With the convention that incumbents who are indifferent 
+remain rather than exit, an incumbent with productivity $\varphi$ exits
+if and only if $\varphi < \bar \varphi (p)$.
+
+
+
+### Equilibrium 
+
+Now we describe a stationary recursive equilibrium in the model, where
+
+* the goods market clears (supply equals demand)
+* the firm size distribution is stationary over time (an invariance condition)
+* the mass of entrants equals the mass of exiting firms, and
+* entering firms satisfy a break-even condition.
+
+The break-even condition means that expected lifetime value given an initial draw from $\gamma$ is just equal to the fixed cost of entry, which we denote by $c_e$.
+
+Let's now state this more formally
+
+To begin, let $\mathscr B$ be the Borel subsets of $\mathbb R_+$ and $\mathscr M$ be all measures
+on $\mathscr B$.  
+
+Taking $\bar v$ and $\bar \varphi$ as defined in the previous
+section, a **stationary recursive equilibrium**  is a triple 
+%
+\begin{equation*}
+    (p, M, \mu) \quad \text{ in } \quad
+    \mathscr E := (0, \infty) \times (0, \infty) \times \mathscr M,
+\end{equation*}
+%
+with $p$ understood as price, $M$ as mass of entrants, and $\mu$ as a
+distribution of firms over productivity levels, such that the goods market clears, or
+%
+\begin{equation}
+    \int q( \varphi, p) \mu(d \varphi) = D(p),
+\end{equation}
+%
+the *invariance condition* over the firm distribution
+%
+\begin{equation}
+    \mu (B) 
+    = \int \Gamma(\varphi, B) \, \mathbf{1}  \{\varphi \geq \bar \varphi(p)\} \,
+    \mu(d \varphi)
+    + M \, \gamma(B) 
+    \text{ for all } B \in \mathscr B,
+\end{equation}
+%
+holds (see below for intuition), the *equilibrium entry condition*
+%
+\begin{equation}
+    \int \bar v( \varphi, p) \, \gamma(d \varphi) = c_e
+\end{equation}
+%
+holds, and the *balanced entry and exit condition*
+%
+\begin{equation}
+    M = \mu \{ \varphi < \bar \varphi(p) \}
+\end{equation}
+%
+is verified.
+
+The invariance condition says that, for any subset $B$ of the state space, the mass of firms with productivity in set $B$ today (the left-hand side) is equal to the mass of firms with productivity in set $B$ tomorrow (the right-hand side).
+
+
+
+### Computing equilibrium
+
+We compute the equilibrium as follows:
+
+For the purposes of this section, we insert balanced entry and
+exit into the time invariance condition, yielding
+%
+\begin{equation}
+    \mu (B) 
+    = \int \Pi(\varphi, B) \mu (d \varphi)
+    \quad \text{for all } B \in \mathscr B,
+\end{equation}
+%
+where $\Pi$ is the transition kernel on $\mathbb R_+$ defined by
+%
+\begin{equation}
+    \Pi(\varphi, B) 
+    = \Gamma(\varphi, B) \mathbf{1} \{\varphi \geq \bar \varphi(p) \} + 
+    \mathbf{1} \{\varphi < \bar \varphi(p)\} \gamma(B).
+\end{equation}
+%
+
+Under our assumptions, for each $p > 0$, there exists a unique $\mu$ satisfying
+this invariance law.
+
+Let $\mathscr P$ be the probability measures on $\mathbb R_+$.
+
+The unique stationary equilibrium can be computed as follows:
+
+1. Obtain $\bar v$ as the unique solution to the Bellman equation, and
+        then calculate the exit threshold function $\bar \varphi$.
+2. Solve for the equilibrium entry price $p^*$ by solving the entry condition.
+3. Define $\Pi$ as above, using $p^*$ as the price, and compute $\mu$
+        as the unique solution to the invariance condition.
+4. Rescale $\mu$ by setting $s := D(p^*)/ \int q( \varphi, p^*)
+        \mu(d \varphi)$ and then $\mu^* := s \, \mu$.
+5. Obtain the mass of entrants via $M^* = \mu^* \{ \varphi < \bar \varphi(p^*) \}$.
+
+
+
+When we compute the distribution $\mu$ in step 3, we will use the fact that it
+is the stationary distribution of the Markov transition kernel $\Pi(\varphi, d
+                                                                    \varphi)$.
+
+This transition kernel turns out to be *ergodic*, which means that if we simulate a cross-section of firms according to $\Pi$ for a large number of periods, the resulting sample (cross-section of productivities) will approach a set of IID draws from $\mu$.  
+
+This allows us to 
+
+1. compute integrals with respect to the distribution using Monte Carlo, and
+2. investigate the shape and properties of the stationary distribution.
+
+
+
+### Specification of dynamics
+
+Before solving the model we need to specify $\Gamma$ and $\gamma$.
+
+We assume $\Gamma(\varphi, d \varphi')$ is given by
+
+$$
+    \varphi_{t+1} = A_{t+1} \varphi_t 
+$$
+
+We assume that $(A_t)$ is IID over time, independent across firms, and lognormal $LN(m_a, \sigma_a)$.
+
+(This means that incumbents follow Gibrat's law, which is a reasonable assumption for medium to large firms -- and hence for incumbents.)
+
+New entrants are drawn from a lognormal distribution $LN(m_e, \sigma_e)$.
+
+
+
+## Code
+
+
+
+We use 64 bit floats for extra precision.
+
+```{code-cell} ipython3
+jax.config.update("jax_enable_x64", True)
+```
+
+We store the parameters, grids and Monte Carlo draws in a `namedtuple`.
+
+```{code-cell} ipython3
+Parameters = namedtuple("Parameters", 
+    ("β",             # discount factor        
+     "θ",             # labor productivity
+     "c",             # fixed cost in production
+     "c_e",           # entry cost
+     "w",             # wages
+     "m_a",           # productivity shock location parameter
+     "σ_a",           # productivity shock scale parameter
+     "m_e",           # new entrant location parameter
+     "σ_e"))          # new entrant scale parameter
+```
+
+```{code-cell} ipython3
+Grids = namedtuple("Grids",
+    ("φ_grid",        # productivity grid
+     "E_draws",       # entry size draws for Monte Carlo
+     "A_draws"))      # productivity shock draws for Monte Carlo
+```
+
+```{code-cell} ipython3
+Model = namedtuple("Model",
+    ("parameters",    # instance of Parameters
+     "grids"))        # instance of Grids
+```
+
+```{code-cell} ipython3
+def create_model(β=0.95,             # discount factor
+                 θ=0.3,              # labor productivity
+                 c=4.0,              # fixed cost in production
+                 c_e=1.0,            # entry cost
+                 w=1.0,              # wages
+                 m_a=-0.012,         # productivity shock location parameter
+                 σ_a=0.1,            # productivity shock scale parameter
+                 m_e=1.0,            # new entrant location parameter
+                 σ_e=0.2,            # new entrant scale parameter
+                 φ_grid_max=5,       # productivity grid max
+                 φ_grid_size=100,    # productivity grid size
+                 E_draw_size=200,    # entry MC integration size
+                 A_draw_size=200,    # prod shock MC integration size
+                 seed=1234):         # Seed for MC draws
+    """
+    Create an instance of the `namedtuple` Model using default parameter values.
+    """
+    
+    # Test stability
+    assert m_a + σ_a**2 / (2 * (1 - θ)) < 0, "Stability condition fails"
+    # Build grids and initialize random number generator
+    φ_grid = jnp.linspace(0, φ_grid_max, φ_grid_size)
+    key, subkey = jax.random.split(jax.random.PRNGKey(seed))
+    # Generate a sample of draws of A for Monte Carlo integration
+    A_draws = jnp.exp(m_a + σ_a * jax.random.normal(key, (A_draw_size,)))
+    # Generate a sample of draws from γ for Monte Carlo
+    E_draws = jnp.exp(m_e + σ_e * jax.random.normal(subkey, (E_draw_size,)))
+    # Build namedtuple and return
+    parameters = Parameters(β, θ, c, c_e, w, m_a, σ_a, m_e, σ_e)
+    grids = Grids(φ_grid, E_draws, A_draws)
+    model = Model(parameters, grids)
+    return model
+```
+
+Let us write down functions for profits and output.
+
+```{code-cell} ipython3
+@jax.jit
+def π(φ, p, parameters):
+    """ Profits. """
+    # Unpack
+    β, θ, c, c_e, w, m_a, σ_a, m_e, σ_e = parameters
+    # Compute profits
+    return (1 - θ) * (p * φ)**(1/(1 - θ)) * (θ/w)**(θ/(1 - θ)) - c 
+```
+
+```{code-cell} ipython3
+@jax.jit
+def q(φ, p, parameters):
+    """ Output. """
+    # Unpack
+    β, θ, c, c_e, w, m_a, σ_a, m_e, σ_e = parameters
+    # Compute output
+    return φ**(1/(1 - θ)) * (p * θ/w)**(θ/(1 - θ)) 
+```
+
+Let's write code to simulate a cross-section of firms given a particular value for the exit threshold (rather than an exit threshold function).
+
+Firms that exit are immediately replaced by a new entrant, drawn from $\gamma$.
+
+Our first function updates by one step
+
+```{code-cell} ipython3
+def update_cross_section(φ_bar, φ_vec, key, parameters, num_firms):
+    # Unpack
+    β, θ, c, c_e, w, m_a, σ_a, m_e, σ_e = parameters
+    # Update
+    Z = jax.random.normal(key, (2, num_firms))  # long rows for row-major arrays
+    incumbent_draws = φ_vec * jnp.exp(m_a + σ_a * Z[0, :])
+    new_firm_draws = jnp.exp(m_e + σ_e * Z[1, :])
+    return jnp.where(φ_vec >= φ_bar, incumbent_draws, new_firm_draws)
+```
+
+```{code-cell} ipython3
+update_cross_section = jax.jit(update_cross_section, static_argnums=(4,))
+```
+
+Our next function runs the cross-section forward in time `sim_length` periods.
+
+```{code-cell} ipython3
+def simulate_firms(φ_bar, parameters, grids, 
+                   sim_length=200, num_firms=1_000_000, seed=12):
+    """
+    Simulate a cross-section of firms when the exit threshold is φ_bar.
+
+    """
+    # Set initial conditions to the threshold value
+    φ_vec = jnp.ones((num_firms,)) * φ_bar 
+    key = jax.random.PRNGKey(seed)
+    # Iterate forward in time
+    for t in range(sim_length):
+        key, subkey = jax.random.split(key)
+        φ_vec = update_cross_section(φ_bar, φ_vec, subkey, parameters, num_firms)
+    return φ_vec
+```
+
+Here's a utility function to compute the expected value
+
+$$
+    \int v(\varphi') \Gamma(\varphi, d \varphi') = \mathbb E v(A_{t+1} \varphi)
+$$
+
+given $\varphi$
+
+```{code-cell} ipython3
+@jax.jit
+def _compute_exp_value_at_phi(v, φ, grids):
+    """
+    Compute 
+    
+        E[v(φ')| φ] = Ev(A φ) 
+        
+    using linear interpolation and Monte Carlo. 
+    """
+    # Unpack
+    φ_grid, E_draws, A_draws = grids
+    # Set up V
+    Aφ = A_draws * φ 
+    vAφ  = jnp.interp(Aφ, φ_grid, v)  # v(A_j φ) for all j
+    # Return mean 
+    return jnp.mean(vAφ)     # (1/n) Σ_j v(A_j φ)
+```
+
+Now let's vectorize this function in $\varphi$ and then write another function that computes the expected value across all $\varphi$ in `φ_grid`
+
+```{code-cell} ipython3
+compute_exp_value_at_phi = jax.vmap(_compute_exp_value_at_phi, (None, 0, None))
+```
+
+```{code-cell} ipython3
+@jax.jit
+def compute_exp_value(v, grids):
+    """
+    Compute 
+    
+        E[v(φ_prime)| φ] = Ev(A φ) for all φ, as a vector
+
+    """
+    # Unpack
+    φ_grid, E_draws, A_draws = grids
+    return compute_exp_value_at_phi(v, φ_grid, grids)
+```
+
+Here is the Bellman operator $T$.
+
+```{code-cell} ipython3
+@jax.jit
+def T(v, p, parameters, grids):
+    """ Bellman operator. """
+    # Unpack
+    β, θ, c, c_e, w, m_a, σ_a, m_e, σ_e = parameters
+    φ_grid, E_draws, A_draws = grids
+    # Compute Tv
+    EvAφ = compute_exp_value(v, grids)
+    return π(φ_grid, p, parameters) + β * jnp.maximum(0.0, EvAφ)
+```
+
+The next function takes $v, p$ as inputs and, conditional on the value function
+$v$, computes the value $\bar \varphi(p)$ that corresponds to the exit value.
+
+```{code-cell} ipython3
+@jax.jit
+def get_threshold(v, grids):
+    """ Compute the exit threshold. """
+    # Unpack
+    φ_grid, E_draws, A_draws = grids
+    # Compute exit threshold: φ such that E v(A φ) = 0
+    EvAφ = compute_exp_value(v, grids)
+    i = jnp.searchsorted(EvAφ, 0.0)
+    return φ_grid[i]
+```
+
+We use value function iteration (VFI) to compute the value function $\bar v(\cdot, p)$, taking $p$ as given.
+
+VFI is relatively cheap and simple in this setting.
+
+```{code-cell} ipython3
+@jax.jit
+def vfi(p, v_init, parameters, grids, tol=1e-6, max_iter=10_000):
+    """
+    Implement value function iteration to solve for the value function.
+    """
+    # Unpack
+    φ_grid, E_draws, A_draws = grids
+    # Set up
+    def cond_function(state):
+        i, v, error = state
+        return jnp.logical_and(i < max_iter, error > tol)
+
+    def body_function(state):
+        i, v, error = state
+        new_v = T(v, p, parameters, grids)
+        error = jnp.max(jnp.abs(v - new_v))
+        i += 1
+        return i, new_v, error
+
+    # Loop till convergence
+    init_state = 0, v_init, tol + 1
+    state = jax.lax.while_loop(cond_function, body_function, init_state) 
+    i, v, error = state
+    return v
+```
+
+```{code-cell} ipython3
+@jax.jit
+def compute_net_entry_value(p, v_init, parameters, grids):
+    """
+    Returns the net value of entry, which is 
+        
+         \int v_bar(φ, p) γ(d φ) - c_e
+
+    This is the break-even condition for new entrants.  The argument
+    v_init is used as an initial condition when computing v_bar for VFI.
+    """
+    c_e = parameters.c_e
+    φ_grid = grids.φ_grid
+    E_draws = grids.E_draws
+    v_bar = vfi(p, v_init, parameters, grids)
+    v_φ = jnp.interp(E_draws, φ_grid, v_bar)
+    Ev_φ = jnp.mean(v_φ)
+    return Ev_φ - c_e, v_bar
+```
+
+We need to solve for the equilibrium price, which is the $p$ satisfying
+
+$$
+\int \bar v(\varphi', p) \gamma(d \varphi') = c_e
+$$
+
+At each price $p$, we need to recompute $\bar v(\cdot, p)$ and then take the expectation.
+
+The technique we will use is bisection.
+
+We will write our own bisection routine because, when we shift to a new price, we want to update the initial condition for value function iteration to the value function from the previous price.
+
+```{code-cell} ipython3
+def compute_p_star(parameters, grids, p_min=1.0, p_max=2.0, tol=10e-5):
+    """
+    Compute the equilibrium entry price p^* via bisection.
+
+    Return both p* and the corresponding value function v_bar, which is
+    computed as a byproduct.
+    
+    Implements the bisection root finding algorithm to find p_star
+
+    """
+    φ_grid, E_draws, A_draws = grids
+    lower, upper = p_min, p_max
+    v_bar = jnp.zeros_like(φ_grid)  # Initial condition at first price guess
+
+    while upper - lower > tol:
+        mid = 0.5 * (upper + lower)
+        entry_val, v_bar = compute_net_entry_value(mid, v_bar, parameters, grids)
+        if entry_val > 0:   # Root is between lower and mid
+            lower, upper = lower, mid
+        else:               # Root is between mid and upper
+            lower, upper = mid, upper
+
+    p_star = 0.5 * (upper + lower)
+    return p_star, v_bar
+```
+
+We are now ready to compute all elements of the stationary recursive equilibrium.
+
+```{code-cell} ipython3
+def compute_equilibrium_prices_and_quantities(model):
+    """
+    Compute 
+
+        1. The equilibrium outcomes for p*, v* and φ*, where φ* is the
+           equilibrium exit threshold φ_bar(p*).
+        1. The scaling factor necessary to convert the stationary probability
+           distribution μ into the equilibrium firm distribution μ* = s μ.
+        2. The equilibrium mass of entrants M* = μ*{ φ < φ*}
+
+    """
+    # Unpack
+    parameters, grids = model
+    # Compute prices and values
+    p_star, v_bar = compute_p_star(parameters, grids)
+    # Get φ_star = φ_bar(p_star), the equilibrium exit threshold 
+    φ_star = get_threshold(v_bar, grids)
+    # Generate an array of draws from μ, the normalized stationary distribution.
+    φ_sample = simulate_firms(φ_star, parameters, grids)
+    # Compute s to scale μ
+    demand = 1 / p_star
+    pre_normalized_supply = jnp.mean(q(φ_sample, p_star, parameters))
+    s = demand / pre_normalized_supply
+    # Compute M* = μ*{ φ < φ_star}
+    m_star = s * jnp.mean(φ_sample < φ_star)
+    # return computed objects
+    return p_star, v_bar, φ_star, φ_sample, s, m_star
+```
+
+## Solving the model
+
+### Preliminary calculations
+
+
+
+Let's create an instance of the model.
+
+```{code-cell} ipython3
+model = create_model()
+parameters, grids = model
+```
+
+Let's see how long it takes to compute the value function at a given price
+from a cold start.
+
+```{code-cell} ipython3
+p = 2.0
+v_init = jnp.zeros_like(grids.φ_grid)            # Initial condition 
+%time v_bar = vfi(p, v_init, parameters, grids) 
+```
+
+```{code-cell} ipython3
+%time v_bar = vfi(p, v_init, parameters, grids)  
+```
+
+Let's have a look at the net entry value as a function of price
+
+The root is the equilibrium price at the given parameters
+
+```{code-cell} ipython3
+p_min, p_max, p_size = 1.0, 2.0, 20
+p_vec = jnp.linspace(p_min, p_max, p_size)
+entry_vals = []
+v_bar = jnp.zeros_like(grids.φ_grid)  # Initial condition at first price guess
+for i, p in enumerate(p_vec):
+    entry_val, v_bar = compute_net_entry_value(p, v_bar, parameters, grids)
+    entry_vals.append(entry_val)
+fig, ax = plt.subplots()
+ax.plot(p_vec, entry_vals, label="net value of entry")
+ax.plot(p_vec, jnp.zeros_like(p_vec), 'k', ls='--', label="break even")
+ax.legend()
+ax.set_xlabel("price")
+ax.set_ylabel("value")
+plt.show()
+```
+
+Below we solve for the zero of this function to calculate $p*$.
+
+From the figure it looks like $p^*$ will be close to 1.5.
+
+
+
+### Computing the equilibrium
+
+
+
+Now let's try computing the equilibrium
+
+```{code-cell} ipython3
+%%time
+p_star, v_bar, φ_star, φ_sample, s, m_star = \
+        compute_equilibrium_prices_and_quantities(model)
+```
+
+Let's check that $p^*$ is close to 1.5
+
+```{code-cell} ipython3
+p_star
+```
+
+Here is a plot of the value function $\bar v(\cdot, p^*)$.
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+ax.plot(grids.φ_grid, v_bar, label=r'$\varphi \mapsto \bar v(\varphi, p^*)$')
+ax.set_xlabel("productivity")
+ax.set_ylabel("firm value")
+ax.legend()
+plt.show()
+```
+
+Let's have a look at the firm size distribution, with firm size measured by output.
+
+```{code-cell} ipython3
+output_dist = q(φ_sample, p_star, parameters)
+```
+
+```{code-cell} ipython3
+fig, ax = plt.subplots(figsize=(7, 4))
+ax.hist(jnp.log(output_dist), bins=100, density=True,
+        label="firm size distribution")
+ax.set_xlabel("log output")
+ax.set_ylabel("frequency")
+ax.legend()
+plt.show()
+```
+
+## Pareto tails
+
+
+
+The firm size distribution shown above appears to have a long right tail.
+
+This matches the observed firm size distribution.
+
+In fact the firm size distribution obeys a **power law**.
+
+More mathematically, the distribution of firm size has a Pareto right hand tail, so that there exist constants $k, \alpha > 0$ with
+
+$$
+   \mu((x, \infty)) \approx kx^{-\alpha} \text{ when $x$ is large}
+$$
+
+Here $\alpha$ is called the tail index.
+
+Does the model replicate this feature?
+
+
+
+One option is to look at the empirical counter CDF (cumulative distribution).
+
+The idea is as follows: The counter CDF of a random variable $X$ is 
+
+$$
+   G(x) := \mathbb P\{X > x\}
+$$
+
+In the case of a Pareto tailed distribution we have $\mathbb P\{X > x\} \approx k x^{-\alpha}$ for large $x$.
+
+Hence, for large $x$,
+
+$$
+    \ln G(x) \approx \ln k - \alpha \ln x
+$$
+
+The empirical counterpart of $G$ given sample $X_1, \ldots, X_n$ is 
+
+$$
+    G_n (x) := \frac{1}{n} \sum_{i=1}^n \mathbf 1 \{X_i > x \}
+$$
+
+For large $k$ (implying $G_n \approx G$) and large $x$, we expect that, for a Pareto-tailed sample, $\ln G_n$ is approximately linear.
+
+
+
+Here's a function to compute the empirical counter CDF:
+
+```{code-cell} ipython3
+def ECCDF(data):
+    """
+    Return a function that implements the ECCDF given the data.
+    """
+    def eccdf(x):
+        return jnp.mean(data > x)
+    return eccdf
+```
+
+Let's plot the empirical counter CDF of the output distribution.
+
+```{code-cell} ipython3
+eccdf = ECCDF(output_dist)
+
+ϵ = 10e-10
+x_grid = jnp.linspace(output_dist.min() + ϵ, output_dist.max() - ϵ, 200)
+y = [eccdf(x) for x in x_grid]
+
+fix, ax = plt.subplots()
+ax.loglog(x_grid, y, 'o', label="ECCDF")
+ax.set_xlabel("productivity")
+ax.set_ylabel("counter CDF")
+plt.show()
+```
+
+## Exercise
+
+
+```{exercise-start}
+:label: hopen_ex1
+```
+
+Plot the same output distribution, but this time using a rank-size plot.
+
+If the rank-size plot is approximately linear, the data suggests a Pareto tail.
+
+You can use QuantEcon's `rank_size` function --- here's an example of usage.
+
+```{code-cell} ipython3
+x = jnp.abs(jax.random.normal(jax.random.PRNGKey(42), (1_000_000,)))
+rank_data, size_data = qe.rank_size(x, c=0.1)
+fig, ax = plt.subplots(figsize=(7,4))
+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()
+```
+
+This plot is not linear --- as expected, since we are using a folded normal distribution.
+
+
+```{exercise-end}
+```
+
+```{solution-start} hopen_ex1
+:class: dropdown
+```
+
+```{code-cell} ipython3
+rank_data, size_data = qe.rank_size(output_dist, c=0.1)
+fig, ax = plt.subplots(figsize=(7,4))
+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()
+```
+
+This looks very linear --- the model generates Pareto tails.
+
+(In fact it's possible to prove this.)
+
+```{solution-end}
+```
+
+
+```{exercise-start}
+:label: hopen_ex2
+```
+
+As an exercise, let's look at the fixed cost paid by incumbents each period and how it relates to the equilibrium price.
+
+We expect that a higher fixed cost will reduce supply and hence increase the market price.
+
+For the fixed costs, use
+
+```{code-cell} ipython3
+c_values = jnp.linspace(2.5, 5.0, 10)
+```
+
+```{exercise-end}
+```
+
+```{solution-start} hopen_ex2
+:class: dropdown
+```
+
+```{code-cell} ipython3
+eq_prices = []
+for i, c in enumerate(c_values):
+    model = create_model(c=c)
+    p_star, v_bar, φ_star, φ_sample, s, m_star = \
+        compute_equilibrium_prices_and_quantities(model)
+    eq_prices.append(p_star)
+    print(f"Equilibrium price when c = {c:.2} is {p_star:.2}")
+
+fig, ax = plt.subplots()
+ax.plot(c_values, eq_prices, label="$p^*$")
+ax.set_xlabel("fixed cost for incumbents")
+ax.set_ylabel("price")
+ax.legend()
+plt.show()
+```
+
+```{solution-end}
+```
diff --git a/lectures/job_search.md b/lectures/job_search.md
new file mode 100644
index 00000000..911a6d8a
--- /dev/null
+++ b/lectures/job_search.md
@@ -0,0 +1,336 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.16.1
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# Job Search
+
+```{include} _admonition/gpu.md
+```
+
+
+In this lecture we study a basic infinite-horizon job search with Markov wage
+draws 
+
+The exercise at the end asks you to add recursive preferences and compare
+the result.
+
+In addition to what’s in Anaconda, this lecture will need the following libraries:
+
+```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install quantecon
+```
+
+We use the following imports.
+
+```{code-cell} ipython3
+import matplotlib.pyplot as plt
+import quantecon as qe
+import jax
+import jax.numpy as jnp
+from collections import namedtuple
+
+jax.config.update("jax_enable_x64", True)
+```
+
+## Model
+
+We study an elementary model where 
+
+* jobs are permanent 
+* unemployed workers receive current compensation $c$
+* the wage offer distribution $\{W_t\}$ is Markovian
+* the horizon is infinite
+* an unemployment agent discounts the future via discount factor $\beta \in (0,1)$
+
+The wage process obeys
+
+$$
+    W_{t+1} = \rho W_t + \nu Z_{t+1},
+    \qquad \{Z_t\} \text{ is IID and } N(0, 1)
+$$
+
+We discretize this using Tauchen's method to produce a stochastic matrix $P$
+
+Since jobs are permanent, the return to accepting wage offer $w$ today is
+
+$$
+    w + \beta w + \beta^2 w + \cdots = \frac{w}{1-\beta}
+$$
+
+The Bellman equation is
+
+$$
+    v(w) = \max
+    \left\{
+            \frac{w}{1-\beta}, c + \beta \sum_{w'} v(w') P(w, w')
+    \right\}
+$$
+
+We solve this model using value function iteration.
+
+
+Let's set up a `namedtuple` to store information needed to solve the model.
+
+```{code-cell} ipython3
+Model = namedtuple('Model', ('n', 'w_vals', 'P', 'β', 'c'))
+```
+
+The function below holds default values and populates the namedtuple.
+
+```{code-cell} ipython3
+def create_js_model(
+        n=500,       # wage grid size
+        ρ=0.9,       # wage persistence
+        ν=0.2,       # wage volatility
+        β=0.99,      # discount factor
+        c=1.0        # unemployment compensation
+    ):
+    "Creates an instance of the job search model with Markov wages."
+    mc = qe.tauchen(n, ρ, ν)
+    w_vals, P = jnp.exp(mc.state_values), mc.P
+    P = jnp.array(P)
+    return Model(n, w_vals, P, β, c)
+```
+
+Here's the Bellman operator.
+
+```{code-cell} ipython3
+@jax.jit
+def T(v, model):
+    """
+    The Bellman operator Tv = max{e, c + β E v} with 
+
+        e(w) = w / (1-β) and (Ev)(w) = E_w[ v(W')]
+
+    """
+    n, w_vals, P, β, c = model
+    h = c + β * P @ v
+    e = w_vals / (1 - β)
+
+    return jnp.maximum(e, h)
+```
+
+The next function computes the optimal policy under the assumption that $v$ is
+                 the value function.
+
+The policy takes the form
+
+$$
+    \sigma(w) = \mathbf 1 
+        \left\{
+            \frac{w}{1-\beta} \geq c + \beta \sum_{w'} v(w') P(w, w')
+        \right\}
+$$
+
+Here $\mathbf 1$ is an indicator function.
+
+The statement above means that the worker accepts ($\sigma(w) = 1$) when the value of stopping
+is higher than the value of continuing.
+
+```{code-cell} ipython3
+@jax.jit
+def get_greedy(v, model):
+    """Get a v-greedy policy."""
+    n, w_vals, P, β, c = model
+    e = w_vals / (1 - β)
+    h = c + β * P @ v
+    σ = jnp.where(e >= h, 1, 0)
+    return σ
+```
+
+Here's a routine for value function iteration.
+
+```{code-cell} ipython3
+def vfi(model, max_iter=10_000, tol=1e-4):
+    """Solve the infinite-horizon Markov job search model by VFI."""
+
+    print("Starting VFI iteration.")
+    v = jnp.zeros_like(model.w_vals)    # Initial guess
+    i = 0
+    error = tol + 1
+
+    while error > tol and i < max_iter:
+        new_v = T(v, model)
+        error = jnp.max(jnp.abs(new_v - v))
+        i += 1
+        v = new_v
+
+    v_star = v
+    σ_star = get_greedy(v_star, model)
+    return v_star, σ_star
+```
+
+### Computing the solution
+
+Let's set up and solve the model.
+
+```{code-cell} ipython3
+model = create_js_model()
+n, w_vals, P, β, c = model
+
+qe.tic()
+v_star, σ_star = vfi(model)
+vfi_time = qe.toc()
+```
+
+We compute the reservation wage as the first $w$ such that $\sigma(w)=1$.
+
+```{code-cell} ipython3
+res_wage = w_vals[jnp.searchsorted(σ_star, 1.0)]
+```
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+ax.plot(w_vals, v_star, alpha=0.8, label="value function")
+ax.vlines((res_wage,), 150, 400, 'k', ls='--', label="reservation wage")
+ax.legend(frameon=False, fontsize=12, loc="lower right")
+ax.set_xlabel("$w$", fontsize=12)
+plt.show()
+```
+
+## Exercise
+
+```{exercise-start}
+:label: job_search_1
+```
+
+In the setting above, the agent is risk-neutral vis-a-vis future utility risk.
+
+Now solve the same problem but this time assuming that the agent has risk-sensitive
+preferences, which are a type of nonlinear recursive preferences.
+
+The Bellman equation becomes
+
+$$
+    v(w) = \max
+    \left\{
+            \frac{w}{1-\beta}, 
+            c + \frac{\beta}{\theta}
+            \ln \left[ 
+                      \sum_{w'} \exp(\theta v(w')) P(w, w')
+                \right]
+    \right\}
+$$
+
+
+When $\theta < 0$ the agent is risk sensitive.
+
+Solve the model when $\theta = -0.1$ and compare your result to the risk neutral
+case.
+
+Try to interpret your result.
+
+```{exercise-end}
+```
+
+```{solution-start} job_search_1
+:class: dropdown
+```
+
+```{code-cell} ipython3
+RiskModel = namedtuple('Model', ('n', 'w_vals', 'P', 'β', 'c', 'θ'))
+
+def create_risk_sensitive_js_model(
+        n=500,       # wage grid size
+        ρ=0.9,       # wage persistence
+        ν=0.2,       # wage volatility
+        β=0.99,      # discount factor
+        c=1.0,       # unemployment compensation
+        θ=-0.1       # risk parameter
+    ):
+    "Creates an instance of the job search model with Markov wages."
+    mc = qe.tauchen(n, ρ, ν)
+    w_vals, P = jnp.exp(mc.state_values), mc.P
+    P = jnp.array(P)
+    return RiskModel(n, w_vals, P, β, c, θ)
+
+
+@jax.jit
+def T_rs(v, model):
+    """
+    The Bellman operator Tv = max{e, c + β R v} with 
+
+        e(w) = w / (1-β) and
+
+        (Rv)(w) = (1/θ) ln{E_w[ exp(θ v(W'))]}
+
+    """
+    n, w_vals, P, β, c, θ = model
+    h = c + (β / θ) * jnp.log(P @ (jnp.exp(θ * v)))
+    e = w_vals / (1 - β)
+
+    return jnp.maximum(e, h)
+
+
+@jax.jit
+def get_greedy_rs(v, model):
+    " Get a v-greedy policy."
+    n, w_vals, P, β, c, θ = model
+    e = w_vals / (1 - β)
+    h = c + (β / θ) * jnp.log(P @ (jnp.exp(θ * v)))
+    σ = jnp.where(e >= h, 1, 0)
+    return σ
+
+
+
+def vfi(model, max_iter=10_000, tol=1e-4):
+    "Solve the infinite-horizon Markov job search model by VFI."
+    print("Starting VFI iteration.")
+    v = jnp.zeros_like(model.w_vals)    # Initial guess
+    i = 0
+    error = tol + 1
+
+    while error > tol and i < max_iter:
+        new_v = T_rs(v, model)
+        error = jnp.max(jnp.abs(new_v - v))
+        i += 1
+        v = new_v
+
+    v_star = v
+    σ_star = get_greedy_rs(v_star, model)
+    return v_star, σ_star
+
+
+
+model_rs = create_risk_sensitive_js_model()
+
+n, w_vals, P, β, c, θ = model_rs
+
+qe.tic()
+v_star_rs, σ_star_rs = vfi(model_rs)
+vfi_time = qe.toc()
+```
+
+```{code-cell} ipython3
+res_wage_rs = w_vals[jnp.searchsorted(σ_star_rs, 1.0)]
+```
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+ax.plot(w_vals, v_star,  alpha=0.8, label="RN $v$")
+ax.plot(w_vals, v_star_rs, alpha=0.8, label="RS $v$")
+ax.vlines((res_wage,), 150, 400,  ls='--', color='darkblue', alpha=0.5, label=r"RV $\bar w$")
+ax.vlines((res_wage_rs,), 150, 400, ls='--', color='orange', alpha=0.5, label=r"RS $\bar w$")
+ax.legend(frameon=False, fontsize=12, loc="lower right")
+ax.set_xlabel("$w$", fontsize=12)
+plt.show()
+```
+
+The figure shows that the reservation wage under risk sensitive preferences (RS $\bar w$) shifts down.
+
+This makes sense -- the agent does not like risk and hence is more inclined to
+accept the current offer, even when it's lower.
+
+```{solution-end}
+```
diff --git a/lectures/overborrowing.md b/lectures/overborrowing.md
new file mode 100644
index 00000000..576b3f25
--- /dev/null
+++ b/lectures/overborrowing.md
@@ -0,0 +1,915 @@
+---
+jupytext:
+  text_representation:
+    extension: .md
+    format_name: myst
+    format_version: 0.13
+    jupytext_version: 1.16.1
+kernelspec:
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+---
+
+# Bianchi Overborrowing Model
+
+```{include} _admonition/gpu.md
+```
+
+This lecture provides a JAX implementation of "Overborrowing and Systemic Externalities" {cite:p}`Bianchi2011` by [Javier Bianchi](http://www.javierbianchi.com/).
+
+In addition to what’s in Anaconda, this lecture will need the following libraries:
+
+```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install quantecon
+```
+
+We use the following imports.
+
+```{code-cell} ipython3
+import time
+import jax
+import numba
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+import quantecon as qe
+import scipy as sp
+import matplotlib.pyplot as plt
+import seaborn
+```
+
+## Description of the model
+
+The model seeks to explain sudden stops in emerging market economies, where
+painful financial disruption follows a period of sustained heavy borrowing.
+
+A representative household chooses how much to borrow on international markets and how much to consume.
+
+The household is credit constrained, with the constraint depending on both current income and the real exchange rate.
+
+The model shows that household overborrow because they do not internalize the
+effect of borrowing on the credit constraint.
+
+This overborrowing leaves them vulnerable to bad shocks in current income.
+
+In essence, the model works as follows
+
+1. During good times, households borrow more and consume more
+2. Increased consumption pushes up the price of nontradables and hence the real exchange rate
+3. A rising real exchange rate loosens the credit constraint and encourages more borrowing
+4. This leads to excessive borrowing relative to a planner.
+
+This overborrowing leads to vulnerability vis-a-vis bad shocks.
+
+1. When a bad shock hits, borrowing is restricted.
+2. Consumption now falls, pushing down the real exchange rate.
+3. This fall in the exchange rate further tightens the borrowing constraint, amplifying the shock
+
+
+
+### Decentralized equilibrium
+
+The model contains a representative household that seeks to maximize
+an expected sum of discounted utility where 
+
+* the flow utility function $u$ is CRRA, with $u(c) = c^{1-\sigma}/(1-\sigma)$
+  and
+* $c = (\omega c_t^{-\eta} + (1-\omega) c_n^{-\eta})^{-1/\eta}$
+
+Here $c_t$ ($c_n$) is consumption of tradables (nontradables).
+
+
+
+The household maximizes subject to the budget constraint
+
+$$
+    b' + c_t + p_n c_n = b (1+r) + y_t + p_n y_n
+$$
+
+where
+
+* $b$ is bond holdings (positive values denote assets)
+* primes denote next period values
+* the interest rate $r$ is exogenous
+* $p_n$ is the price of nontradables, while the price of tradables is normalized
+  to 1
+* $y_t$ and $y_n$ are current tradable and nontradable income
+
+The process for $y := (y_t, y_n)$ is first-order Markov.
+
+
+
+The household also faces the credit constraint
+
+$$
+    b' \geq - \kappa (y_t + p_n y_n)
+$$
+
+Market clearing implies
+
+$$
+    c_n = y_n
+    \quad \text{and} \quad
+    c_t = y_t + (1+r)b - b'
+$$
+
+The household takes the aggregate timepath for bonds as given by $B' = H(B, y)$
+and solves
+
+$$
+    V(b, B, y)
+    = \max_{c, b'} 
+    \left\{
+        u(c) + \beta \mathbb{E}_y v(b', B', y')
+    \right\}
+$$
+
+subject to the budget and credit constraints.
+
+
+
+A decentralized equilibrium is a law of motion $H$ such that the implied savings 
+policy $b' = b'(b, B, y)$ verifies 
+
+$$
+    b'(B, B, y) = H(B, y)
+    \quad \text{for all } B, y
+$$
+
+
+
+### Notation
+
+Let
+
+$$
+        w(c_t, y_n) = \frac{c^{1 - σ}}{1 - σ}
+        \quad \text{where} \quad
+        c = [ω c_t^{- η} + (1 - ω) y_n^{- η}]^{-1/η}
+$$
+
+Using the market clearing conditions, we can write the
+household problem as
+
+$$
+    V(b, B, y)
+    = \max_{b'} 
+    \left\{
+        w((1 + r)  b + y_t - b', y_n) + \beta \mathbb{E}_y v(b', H(B, y), y')
+    \right\}
+$$
+
+subject to 
+
+$$
+    - κ (p_n y_n + y_t) 
+    \leq
+    b' \leq (1 + r) b + y_t
+$$
+
+where $p_n$ is given by
+
+$$
+    p_n = ((1 - ω) / ω)  (C / y_n)^{η + 1}
+    \quad \text{with} \quad
+    C := (1 + r) B + y_t - H(B, y)
+$$
+
+
+
+### Constrained planner
+
+The constrained planner solves
+
+$$
+    V(b, B, y)
+    = \max_{c, b'} 
+    \left\{
+        u(c) + \beta \mathbb{E}_y v(b', B', y')
+    \right\}
+$$
+
+subject to the market clearing conditions and 
+the same constraint
+
+$$
+     - \kappa (y_t + p_n y_n) \leq b' \leq (1+r) b + y_t
+$$
+
+although the price of nontradable is now given by
+
+$$
+    p_n = ((1 - ω) / ω) (c_t / y_n)^{η + 1}
+    \quad \text{with} \quad
+    c_t := (1 + r) b + y_t - b'
+$$
+
+We see that the planner internalizes the impact of the savings choice $b'$ on
+the price of nontradables and hence the credit constraint.
+
+
+
+## Markov dynamics
+
+We develop some functions for working with the VAR process
+
+$$
+   \ln y' = A \ln y + u'   
+   \quad \text{(prime indicates next period value)}
+$$
+
+where
+
+* $y = (y_t, y_n) = $ (tradables, nontradables)
+* $A$ is $2 \times 2$
+* $u' \sim N(0, \Omega)$
+* the log function is applied pointwise
+
+We use the following estimated values, reported on p. 12 of [Yamada (2023)](https://jxiv.jst.go.jp/index.php/jxiv/preprint/view/514).
+
+```{code-cell} ipython3
+A = [[0.2425,   0.3297],
+     [-0.1984,  0.7576]]
+Ω = [[0.0052, 0.002],
+     [0.002,  0.0059]]
+```
+
+We'll store the data in $\Omega$ using its square root:
+
+```{code-cell} ipython3
+C = sp.linalg.sqrtm(Ω)
+A = np.array(A)
+```
+
+Here's a function to convert the VAR process to a Markov chain evolving on a
+rectilinear grid of points in $\mathbb R^2$.
+
+Under the hood, this function uses the QuantEcon function `discrete_var`.
+
+```{code-cell} ipython3
+def discretize_income_var(A=A, C=C, grid_size=4, seed=1234):
+    """
+    Discretize the VAR model, returning
+
+        y_t_nodes, a grid of y_t values
+        y_n_nodes, a grid of y_n values
+        Q, a Markov operator
+
+    Let n = grid_size. The format is that Q is n x n x n x n, with
+
+        Q[i, j, i', j'] = one step transition prob from 
+        (y_t_nodes[i], y_n_nodes[j]) to (y_t_nodes[i'], y_n_nodes[j'])
+
+    """
+    
+    n = grid_size
+    rng = np.random.default_rng(seed)
+    mc = qe.markov.discrete_var(A, C, (n, n),
+                                sim_length=1_000_000,
+                                std_devs=np.sqrt(3),
+                                random_state=rng)
+    y_nodes, Q = np.exp(mc.state_values), mc.P
+    # The array y_nodes is currently an array listing all 2 x 1 state pairs
+    # (y_t, y_n), so that y_nodes[i] is the i-th such pair, while Q[l, m] 
+    # is the probability of transitioning from state l to state m in one step. 
+    # We switch the representation to the one described in the docstring.
+    y_t_nodes = [y_nodes[n*i, 0] for i in range(n)]  
+    y_n_nodes = y_nodes[0:4, 1]                      
+    Q = np.reshape(Q, (n, n, n, n))
+    
+    return y_t_nodes, y_n_nodes, Q
+```
+
+Here's code for sampling from the Markov chain.
+
+```{code-cell} ipython3
+def generate_discrete_var(A=A, C=C, grid_size=4, seed=1234, 
+                          ts_length=1_000_000,
+                          indices=False):
+    """
+    Generate a time series from the discretized model, returning y_t_series and
+    y_n_series.  If `indices=True`, then these series are returned as grid
+    indices.
+    """
+    
+    
+    n = grid_size
+    rng = np.random.default_rng(seed)
+    mc = qe.markov.discrete_var(A, C, (n, n),
+                                sim_length=1_000_000,
+                                std_devs=np.sqrt(3),
+                                random_state=rng)
+    if indices:
+        y_series = mc.simulate_indices(ts_length=ts_length)
+        y_t_series, y_n_series = y_series % grid_size, y_series // grid_size
+    else:
+        y_series = np.exp(mc.simulate(ts_length=ts_length))
+        y_t_series, y_n_series = y_series[:, 0], y_series[:, 1]
+    return y_t_series, y_n_series
+```
+
+Here's code for generating the original VAR process, which can be used for
+testing.
+
+```{code-cell} ipython3
+@numba.jit
+def generate_var_process(A=A, C=C, ts_length=1_000_000):
+    """
+    Generate the original VAR process.
+
+    """
+    y_series = np.empty((ts_length, 2))
+    y_series[0, :] = np.zeros(2)
+    for t in range(ts_length-1):
+        y_series[t+1, :] = A @ y_series[t, :] + C @ np.random.randn(2)
+    y_t_series = np.exp(y_series[:, 0])
+    y_n_series = np.exp(y_series[:, 1])
+    return y_t_series, y_n_series
+```
+
+Let's check some statistics for both the original and the discretized processes.
+
+```{code-cell} ipython3
+def corr(x, y):
+    m_x, m_y = x.mean(), y.mean()
+    s_xy = np.sqrt(np.sum((x - m_x)**2) * np.sum((y - m_y)**2))
+    return np.sum((x - m_x) * (y - m_y)) / (s_xy)
+```
+
+```{code-cell} ipython3
+def print_stats(y_t_series, y_n_series):
+    print(f"Std dev of y_t is {y_t_series.std()}")
+    print(f"Std dev of y_n is {y_n_series.std()}")
+    print(f"corr(y_t, y_n) is {corr(y_t_series, y_n_series)}")
+    print(f"auto_corr(y_t) is {corr(y_t_series[:-1], y_t_series[1:])}")
+    print(f"auto_corr(y_n) is {corr(y_n_series[:-1], y_n_series[1:])}")
+    print("\n")
+```
+
+```{code-cell} ipython3
+print("Statistics for original process.\n")
+print_stats(*generate_var_process())
+```
+
+```{code-cell} ipython3
+print("Statistics for discretized process.\n")
+print_stats(*generate_discrete_var())
+```
+
+## Overborrowing Model
+
+In what follows
+
+* `y` = `(y_t, y_n)` is the exogenous state process
+
+Individual states and actions are
+
+* `c` = consumption of tradables (`c` rather than `c_t`)
+* `b` = household savings (bond holdings)
+* `bp` = household savings decision 
+
+Aggregate quantities and prices are
+
+* `p` = price of nontradables (`p` rather than `p_n`)
+* `B` = aggregate savings (bond holdings)
+* `C` = aggregate consumption 
+* `H` = current guess of update rule as an array of the form $H(B, y)$
+
+Here's code to create three tuples that store model data relevant for computation.
+
+```{code-cell} ipython3
+def create_overborrowing_model(
+        σ=2,                 # CRRA utility parameter
+        η=(1/0.83)-1,        # Elasticity = 0.83, η = 0.2048
+        β=0.91,              # Discount factor
+        ω=0.31,              # Aggregation constant
+        κ=0.3235,            # Constraint parameter
+        r=0.04,              # Interest rate
+        b_size=400,          # Bond grid size
+        b_grid_min=-1.02,    # Bond grid min
+        b_grid_max=-0.2      # Bond grid max (originally -0.6 to match fig)
+    ):    
+    """
+    Creates an instance of the overborrowing model using 
+
+        * default parameter values from Bianchi (2011)
+        * Markov dynamics from Yamada (2023)
+
+    The Markov kernel Q has the interpretation
+
+        Q[i, j, ip, jp] = one step prob of moving 
+
+            (y_t[i], y_n[j]) -> (y_t[ip], y_n[jp])
+
+    """
+    # Read in Markov data and shift to JAX arrays
+    data = discretize_income_var()
+    y_t_nodes, y_n_nodes, Q = tuple(map(jnp.array, data))
+    # Set up grid for bond holdings
+    b_grid = jnp.linspace(b_grid_min, b_grid_max, b_size)
+    # Pack and return
+    parameters = σ, η, β, ω, κ, r
+    sizes = b_size, len(y_t_nodes)
+    arrays = b_grid, y_t_nodes, y_n_nodes, Q
+    return parameters, sizes, arrays
+```
+
+Here's flow utility.
+
+```{code-cell} ipython3
+@jax.jit
+def w(parameters, c, y_n):
+    """ 
+    Current utility when c_t = c and c_n = y_n.
+
+        a = [ω c^(- η) + (1 - ω) y_n^(- η)]^(-1/η)
+
+        w(c, y_n) := a^(1 - σ) / (1 - σ)
+
+    """
+    σ, η, β, ω, κ, r = parameters
+    a = (ω * c**(-η) + (1 - ω) * y_n**(-η))**(-1/η)
+    return a**(1 - σ) / (1 - σ)
+```
+
+We need code to generate an initial guess of $H$.
+
+```{code-cell} ipython3
+def generate_initial_H(parameters, sizes, arrays, at_constraint=False):
+    """
+    Compute an initial guess for H. Repeat the indices for b_grid over y_t and
+    y_n axes.
+
+    """
+    b_size, y_size = sizes
+    b_indices = jnp.arange(b_size)
+    O = jnp.ones((b_size, y_size, y_size), dtype=int)
+    return  O * jnp.reshape(b_indices, (b_size, 1, 1)) 
+```
+
+```{code-cell} ipython3
+generate_initial_H = jax.jit(generate_initial_H, static_argnums=(1,))
+```
+
+We need to construct the Bellman operator for the household.
+
+Our first function returns the (unmaximized) RHS of the Bellman equation.
+
+```{code-cell} ipython3
+@jax.jit
+def T_generator(v, H, parameters, arrays, i_b, i_B, i_y_t, i_y_n, i_bp):
+    """
+    Given current state (b, B, y_t, y_n) with indices (i_b, i_B, i_y_t, i_y_n),
+    compute the unmaximized right hand side (RHS) of the Bellman equation as a
+    function of the next period choice bp = b', with index i_bp.  
+    """
+    # Unpack
+    σ, η, β, ω, κ, r = parameters
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    # Compute next period aggregate bonds given H
+    i_Bp = H[i_B, i_y_t, i_y_n]
+    # Evaluate states and actions at indices
+    B, Bp, b, bp = b_grid[i_B], b_grid[i_Bp], b_grid[i_b], b_grid[i_bp]
+    y_t = y_t_nodes[i_y_t]
+    y_n = y_n_nodes[i_y_n]
+    # Compute price of nontradables using aggregates
+    C = (1 + r) * B + y_t - Bp
+    p = ((1 - ω) / ω) * (C / y_n)**(η + 1)
+    # Compute household flow utility
+    c = (1 + r) * b + y_t - bp
+    utility = w(parameters, c, y_n)
+    # Compute expected value Σ_{y'} v(b', B', y') Q(y, y')
+    EV = jnp.sum(v[i_bp, i_Bp, :, :] * Q[i_y_t, i_y_n, :, :])
+    # Set up constraints 
+    credit_constraint_holds = - κ * (p * y_n + y_t) <= bp
+    budget_constraint_holds = bp <= (1 + r) * b + y_t
+    constraints_hold = jnp.logical_and(credit_constraint_holds, 
+                                     budget_constraint_holds)
+    # Compute and return
+    return jnp.where(constraints_hold, utility + β * EV, -jnp.inf)
+```
+
+Let's now vectorize and jit-compile this map.
+
+```{code-cell} ipython3
+# Vectorize over the control bp and all the current states
+T_vec_1 = jax.vmap(T_generator,
+    in_axes=(None, None, None, None, None, None, None, None, 0))
+T_vec_2 = jax.vmap(T_vec_1, 
+    in_axes=(None, None, None, None, None, None, None, 0, None))
+T_vec_3 = jax.vmap(T_vec_2, 
+    in_axes=(None, None, None, None, None, None, 0, None, None))
+T_vec_4 = jax.vmap(T_vec_3, 
+    in_axes=(None, None, None, None, None, 0, None, None, None))
+T_vectorized = jax.vmap(T_vec_4, 
+    in_axes=(None, None, None, None, 0, None, None, None, None))
+```
+
+Now we can set up the Bellman operator by maximizing over the choice variable
+$b'$.
+
+```{code-cell} ipython3
+def T(parameters, sizes, arrays, v, H):
+    """
+    Evaluate the RHS of the Bellman equation at all states and actions and then
+    maximize with respect to actions.
+
+    Return 
+
+        * Tv as an array of shape (b_size, b_size, y_size, y_size).
+
+    """
+    b_size, y_size = sizes
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    b_indices, y_indices = jnp.arange(b_size), jnp.arange(y_size)
+    val = T_vectorized(v, H, parameters, arrays,
+                     b_indices, b_indices, y_indices, y_indices, b_indices)
+    # Maximize over bp
+    return jnp.max(val, axis=-1)
+```
+
+```{code-cell} ipython3
+T = jax.jit(T, static_argnums=(1,))
+```
+
+Here's a function that computes a greedy policy (best response to $v$).
+
+```{code-cell} ipython3
+def get_greedy(parameters, sizes, arrays, v, H):
+    """
+    Compute the greedy policy for the household, which maximizes the right hand
+    side of the Bellman equation given v and H.  The greedy policy is recorded
+    as an array giving the index i in b_grid such that b_grid[i] is the optimal
+    choice, for every state.
+
+    Return 
+
+        * bp_policy as an array of shape (b_size, b_size, y_size, y_size).
+
+    """
+    b_size, y_size = sizes
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    b_indices, y_indices = jnp.arange(b_size), jnp.arange(y_size)
+    val = T_vectorized(v, H, parameters, arrays,
+                       b_indices, b_indices, y_indices, y_indices, b_indices)
+    return jnp.argmax(val, axis=-1)
+```
+
+```{code-cell} ipython3
+get_greedy = jax.jit(get_greedy, static_argnums=(1,))
+```
+
+Here's some code for value function iteration (VFI).
+
+```{code-cell} ipython3
+def vfi(T, v_init, max_iter=10_000, tol=1e-5):
+    """
+    Use successive approximation to compute the fixed point of T, starting from
+    v_init.
+
+    """
+    v = v_init
+
+    def cond_fun(state):
+        error, i, v = state
+        return (error > tol) & (i < max_iter)
+    
+    def body_fun(state):
+        error, i, v = state
+        v_new = T(v)
+        error = jnp.max(jnp.abs(v_new - v))
+        return error, i+1, v_new
+
+    error, i, v_new = jax.lax.while_loop(cond_fun, body_fun,
+                                                    (tol+1, 0, v))
+    return v_new, i
+
+vfi = jax.jit(vfi, static_argnums=(0,))
+```
+
+This is how we update our guess of $H$, using the current policy $b'$
+and a damped fixed point iteration scheme.
+
+```{code-cell} ipython3
+def update_H(parameters, sizes, arrays, H, α):
+    """
+    Update guess of the aggregate update rule.
+
+    """
+    # Set up
+    b_size, y_size = sizes
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    b_indices = jnp.arange(b_size)
+    # Compute household response to current guess H
+    v_init = jnp.ones((b_size, b_size, y_size, y_size))
+    _T = lambda v: T(parameters, sizes, arrays, v, H)
+    v, vfi_num_iter = vfi(_T, v_init)
+    bp_policy = get_greedy(parameters, sizes, arrays, v, H)
+    # Switch policy arrays to values rather than indices
+    H_vals = b_grid[H]
+    bp_vals = b_grid[bp_policy]
+    # Update guess
+    new_H_vals = α * bp_vals[b_indices, b_indices, :, :] + (1 - α) * H_vals
+    # Switch back to indices
+    new_H = jnp.searchsorted(b_grid, new_H_vals)
+    return new_H, vfi_num_iter
+```
+
+```{code-cell} ipython3
+update_H = jax.jit(update_H, static_argnums=(1,))
+```
+
+Now we can write code to compute an equilibrium law of motion $H$.
+
+```{code-cell} ipython3
+def compute_equilibrium(parameters, sizes, arrays,
+                          α=0.5, tol=0.005, max_iter=500):
+    """
+    Compute the equilibrium law of motion.
+
+    """
+    H = generate_initial_H(parameters, sizes, arrays)
+    error = tol + 1
+    i = 0
+    msgs = []
+    while error > tol and i < max_iter:
+        H_new, vfi_num_iter = update_H(parameters, sizes, arrays, H, α)
+        msgs.append(f"VFI terminated after {vfi_num_iter} iterations.")
+        error = jnp.max(jnp.abs(b_grid[H] - b_grid[H_new]))
+        msgs.append(f"Updated H at iteration {i} with error {error}.")
+        H = H_new
+        i += 1
+    if i == max_iter:
+        msgs.append("Warning: Equilibrium search iteration hit upper bound.")
+    print("\n".join(msgs))
+    return H
+```
+
+## Planner problem
+
+Now we switch to the planner problem.
+
+Our first function returns the (unmaximized) RHS of the Bellman equation.
+
+```{code-cell} ipython3
+@jax.jit
+def planner_T_generator(v, parameters, arrays, i_b, i_y_t, i_y_n, i_bp):
+    """
+    Given current state (b, y_t, y_n) with indices (i_b, i_y_t, i_y_n),
+    compute the unmaximized right hand side (RHS) of the Bellman equation as a
+    function of the next period choice bp = b'.  
+    """
+    σ, η, β, ω, κ, r = parameters
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    y_t = y_t_nodes[i_y_t]
+    y_n = y_n_nodes[i_y_n]
+    b, bp = b_grid[i_b], b_grid[i_bp]
+    # Compute price of nontradables using aggregates
+    c = (1 + r) * b + y_t - bp
+    p = ((1 - ω) / ω) * (c / y_n)**(η + 1)
+    # Compute household flow utility
+    utility = w(parameters, c, y_n)
+    # Compute expected value (continuation)
+    EV = jnp.sum(v[i_bp, :, :] * Q[i_y_t, i_y_n, :, :])
+    # Set up constraints and evaluate 
+    credit_constraint_holds = - κ * (p * y_n + y_t) <= bp
+    budget_constraint_holds = bp <= (1 + r) * b + y_t
+    return jnp.where(jnp.logical_and(credit_constraint_holds, 
+                                     budget_constraint_holds), 
+                     utility + β * EV,
+                     -jnp.inf)
+```
+
+```{code-cell} ipython3
+# Vectorize over the control bp and all the current states
+planner_T_vec_1 = jax.vmap(planner_T_generator,
+    in_axes=(None, None, None, None, None, None, 0))
+planner_T_vec_2 = jax.vmap(planner_T_vec_1, 
+    in_axes=(None, None, None, None, None, 0, None))
+planner_T_vec_3 = jax.vmap(planner_T_vec_2, 
+    in_axes=(None, None, None, None, 0, None, None))
+planner_T_vectorized = jax.vmap(planner_T_vec_3, 
+    in_axes=(None, None, None, 0, None, None, None))
+```
+
+Now we construct the Bellman operator.
+
+```{code-cell} ipython3
+def planner_T(parameters, sizes, arrays, v):
+    b_size, y_size = sizes
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    b_indices, y_indices = jnp.arange(b_size), jnp.arange(y_size)
+    # Evaluate RHS of Bellman equation at all states and actions
+    val = planner_T_vectorized(v, parameters, arrays,
+                     b_indices, y_indices, y_indices, b_indices)
+    # Maximize over bp
+    return jnp.max(val, axis=-1)
+```
+
+```{code-cell} ipython3
+planner_T = jax.jit(planner_T, static_argnums=(1,))
+```
+
+Here's a function that computes a greedy policy (best response to $v$).
+
+```{code-cell} ipython3
+def planner_get_greedy(parameters, sizes, arrays, v):
+    b_size, y_size = sizes
+    b_grid, y_t_nodes, y_n_nodes, Q = arrays
+    b_indices, y_indices = jnp.arange(b_size), jnp.arange(y_size)
+    # Evaluate RHS of Bellman equation at all states and actions
+    val = planner_T_vectorized(v, parameters, arrays,
+                     b_indices, y_indices, y_indices, b_indices)
+    # Maximize over bp
+    return jnp.argmax(val, axis=-1)
+```
+
+```{code-cell} ipython3
+planner_get_greedy = jax.jit(planner_get_greedy, static_argnums=(1,))
+```
+
+Computing the planner solution is straightforward value function iteration:
+
+```{code-cell} ipython3
+def compute_planner_solution(model):
+    """
+    Compute the constrained planner solution.
+
+    """
+    parameters, sizes, arrays = model
+    b_size, y_size = sizes
+    b_indices = jnp.arange(b_size)
+    v_init = jnp.ones((b_size, y_size, y_size))
+    _T = lambda v: planner_T(parameters, sizes, arrays, v)
+    # Compute household response to current guess H
+    v, vfi_num_iter = vfi(_T, v_init)
+    bp_policy = planner_get_greedy(parameters, sizes, arrays, v)
+    return v, bp_policy, vfi_num_iter
+```
+
+## Numerical solution
+
+Let's now solve the model and compare the decentralized and planner solutions
+
+### Generating solutions
+
+Here we compute the two solutions.
+
+```{code-cell} ipython3
+model = create_overborrowing_model()
+parameters, sizes, arrays = model
+b_size, y_size = sizes
+b_grid, y_t_nodes, y_n_nodes, Q = arrays
+```
+
+```{code-cell} ipython3
+print("Computing decentralized solution.")
+in_time = time.time()
+H_eq = compute_equilibrium(parameters, sizes, arrays)
+out_time = time.time()
+diff = out_time - in_time
+print(f"Computed decentralized equilibrium in {diff} seconds")
+```
+
+```{code-cell} ipython3
+print("Computing planner's solution.")
+in_time = time.time()
+planner_v, H_plan, vfi_num_iter = compute_planner_solution(model)
+out_time = time.time()
+diff = out_time - in_time
+print(f"Computed decentralized equilibrium in {diff} seconds")
+```
+
+### Policy plots
+
+We produce a policy plot that is similar to Figure 1 in {cite:p}`Bianchi2011`.
+
+```{code-cell} ipython3
+i, j = 1, 3
+y_t, y_n = y_t_nodes[i], y_n_nodes[j]
+fig, ax = plt.subplots()
+ax.plot(b_grid, b_grid[H_eq[:, i, j]], label='decentralized equilibrium')
+ax.plot(b_grid, b_grid[H_plan[:, i, j]], ls='--', label='social planner')
+ax.plot(b_grid, b_grid, color='black', lw=0.5)
+ax.set_ylim((-1.0, -0.6))
+ax.set_xlim((-1.0, -0.6))
+ax.set_xlabel("current bond holdings")
+ax.set_ylabel("next period bond holdings")
+ax.set_title(f"policy when $y_t = {y_t:.2}$ and $y_n = {y_n:.2}$")
+ax.legend()
+plt.show()
+```
+
+The match is not exact because we use a different estimate of the Markov
+dynamics for income.
+
+Nonetheless, it is qualitatively similar.
+
+## Exercise
+
+
+```{exercise-start}
+:label: overborrow_1
+```
+
+Your task is to examine the ergodic distribution of borrowing in the decentralized and
+planner equilibria.
+
+We recommend that you use simulation and a kernel density estimate.
+
+Here's a function for generating the borrowing sequence.
+
+We use Numba because we want to compile a long for loop.
+
+```{code-cell} ipython3
+@numba.jit
+def generate_borrowing_sequence(H, y_t_series, y_n_series):
+    """
+    Generate the borrowing sequence B' = H(B, y_t, y_n).
+
+    * H is a policy array
+    * y_t_series and y_n_series are simulated time paths 
+
+    Both y_t_series and y_n_series are stored as indices rather than values.
+
+    """
+    B = np.empty_like(y_t_series)
+    B[0] = 0
+    for t in range(len(y_t_series)-1):
+        B[t+1] = H[B[t], y_t_series[t], y_n_series[t]]
+    return B
+```
+
+Note that you will have to convert JAX arrays into NumPy arrays if you want to use this function.
+
+From here you will need to 
+
+* generate a time path for income $y = (y_t, y_n)$ using one of the functions provided above.
+* use the function `generate_borrowing_sequence` plus `H_eq` and `H_plan` to calculate bond holdings for the planner and the decentralized equilibrium
+* produce a kernel density plot for each of these data sets
+
+If you are successful, your plot should look something like Fig 2 of {cite:p}`Bianchi2011` --- although not exactly the same, due to the alternative specification of the Markov process.
+
+To generate a kernel density plot, we recommend that you use `kdeplot` from the package `seaborn`, which is included in Anaconda.
+
+```{exercise-end}
+```
+
+```{solution-start} overborrow_1
+:class: dropdown
+```
+
+```{code-cell} ipython3
+sim_length = 100_000
+y_t_series, y_n_series = generate_discrete_var(ts_length=sim_length,
+                                               indices=True)
+```
+
+We convert JAX arrays to NumPy arrays in order to use Numba.
+
+```{code-cell} ipython3
+y_t_series, y_n_series, H_eq, H_plan = \
+        [np.array(v) for v in (y_t_series, y_n_series, H_eq, H_plan)]
+```
+
+Now let's compute borrowing for the decentralized equilibrium and the planner.
+
+```{code-cell} ipython3
+B_eq = generate_borrowing_sequence(H_eq, y_t_series, y_n_series)
+eq_b_sequence = b_grid[B_eq]
+B_plan = generate_borrowing_sequence(H_plan, y_t_series, y_n_series)
+plan_b_sequence = b_grid[B_plan]
+```
+
+We suppress some annoying warnings produced by the current version of `seaborn`
+
+```{code-cell} ipython3
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+```
+
+Now let's plot the distributions using a kernel density estimator.
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+seaborn.kdeplot(eq_b_sequence, ax=ax, label='decentralized')
+seaborn.kdeplot(plan_b_sequence, ax=ax, label='planner')
+ax.legend()
+ax.set_xlim((-1, -0.5))
+ax.set_xlabel("probability")
+ax.set_ylabel("bond holdings")
+plt.show()
+```
+
+This corresponds to Figure 2 in {cite:p}`Bianchi2011`.
+
+Again, the match is not exact but it is qualitatively similar.
+
+Asset holding has a longer left hand tail under the decentralized equilibrium,
+leaving the economy more vulnerable to bad shocks.
+
+```{solution-end}
+```
diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md
index 19256fc2..914284cf 100644
--- a/lectures/wealth_dynamics.md
+++ b/lectures/wealth_dynamics.md
@@ -16,17 +16,41 @@ kernelspec:
 ```{include} _admonition/gpu.md
 ```
 
-This lecture is the extended JAX implementation of [this lecture](https://python.quantecon.org/wealth_dynamics.html). 
+In this lecture we examine wealth dynamics in large cross-section of agents who
+are subject to both 
 
-Please refer that lecture for all background and notation.
+* idiosyncratic shocks, which affect labor income and returns, and 
+* an aggregate shock, which also impacts on labor income and returns
 
-We will use the following imports.
+In most macroeconomic models savings and consumption are determined by optimization.
+
+Here savings and consumption behavior is taken as given -- you can plug in your
+favorite model to obtain savings behavior and then analyze distribution dynamics
+using the techniques described below.
+
+One of our interests will be how different aspects of wealth dynamics -- such
+as labor income and the rate of return on investments -- feed into measures of
+inequality, such as the Gini coefficient.
+
+In addition to JAX and Anaconda, this lecture will need the following libraries:
 
 ```{code-cell} ipython3
+:tags: [hide-output]
+
+!pip install quantecon
+```
+
+We will use the following imports:
+
+```{code-cell} ipython3
+import numba
+import pandas as pd
+import numpy as np
 import matplotlib.pyplot as plt
+import quantecon as qe
 import jax
 import jax.numpy as jnp
-from collections import namedtuple
+from time import time
 ```
 
 Let's check the GPU we are running
@@ -35,603 +59,751 @@ Let's check the GPU we are running
 !nvidia-smi
 ```
 
-## Lorenz Curves and the Gini Coefficient
+## Wealth dynamics
 
-Before we investigate wealth dynamics, we briefly review some measures of
-inequality.
 
-### Lorenz Curves
+Wealth evolves as follows:
 
-One popular graphical measure of inequality is the [Lorenz curve](https://en.wikipedia.org/wiki/Lorenz_curve).
+```{math}
+    w_{t+1} = (1 + r_{t+1}) s(w_t) + y_{t+1}
+```
 
-To illustrate, let us define a function `lorenz_curve_jax` that returns the cumulative share of people and the cumulative share of income earned.
+Here
 
-```{code-cell} ipython3
-@jax.jit
-def lorenz_curve_jax(y):
-    n = y.shape[0]
-    y = jnp.sort(y)
-    s = jnp.concatenate((jnp.zeros(1), jnp.cumsum(y)))
-    _cum_p = jnp.arange(1, n + 1) / n
-    cum_income = s / s[n]
-    cum_people = jnp.concatenate((jnp.zeros(1), _cum_p))
-    return cum_people, cum_income
+- $w_t$ is wealth at time $t$ for a given household,
+- $r_t$ is the rate of return of financial assets,
+- $y_t$ is labor income and
+- $s(w_t)$ is savings (current wealth minus current consumption)
+
+
+There is an aggregate state process 
+
+$$
+    z_{t+1} = a z_t + b + \sigma_z \epsilon_{t+1}
+$$
+
+that affects the interest rate and labor income.
+
+In particular, the gross interest rates obey
+
+$$
+    R_t := 1 + r_t = c_r \exp(z_t) + \exp(\mu_r + \sigma_r \xi_t)
+$$
+
+while
+
+$$
+    y_t = c_y \exp(z_t) + \exp(\mu_y + \sigma_y \zeta_t)
+$$
+
+The tuple $\{ (\epsilon_t, \xi_t, \zeta_t) \}$ is IID and standard normal in $\mathbb R^3$.
+
+(Each household receives their own idiosyncratic shocks.)
+
+Regarding the savings function $s$, our default model will be
+
+```{math}
+:label: sav_ah
+
+s(w) = s_0 w \cdot \mathbb 1\{w \geq \hat w\}
 ```
 
-Let's suppose that
+where $s_0$ is a positive constant.
+
+Thus, 
+
+* for $w < \hat w$, the household saves nothing, while
+* for $w \geq \bar w$, the household saves a fraction $s_0$ of their wealth.
+
+## Implementation 
+
+### Numba implementation
+
+Here's a function that collects parameters and useful constants
 
 ```{code-cell} ipython3
-n = 10_000                                              # Size of sample
-rand_key = jax.random.PRNGKey(101)                      # Set random key
-w = jnp.exp(jax.random.normal(rand_key, shape=(n,)))    # Lognormal draws
+def create_wealth_model(w_hat=1.0,   # Savings parameter
+                        s_0=0.75,    # Savings parameter
+                        c_y=1.0,     # Labor income parameter
+                        μ_y=1.0,     # Labor income parameter
+                        σ_y=0.2,     # Labor income parameter
+                        c_r=0.05,    # Rate of return parameter
+                        μ_r=0.1,     # Rate of return parameter
+                        σ_r=0.5,     # Rate of return parameter
+                        a=0.5,       # Aggregate shock parameter
+                        b=0.0,       # Aggregate shock parameter
+                        σ_z=0.1):    # Aggregate shock parameter
+    """
+    Create a wealth model with given parameters. 
+
+    Return a tuple model = (household_params, aggregate_params), where
+    household_params collects household information and aggregate_params
+    collects information relevant to the aggregate shock process.
+    
+    """
+    # Mean and variance of z process
+    z_mean = b / (1 - a)
+    z_var = σ_z**2 / (1 - a**2)
+    exp_z_mean = np.exp(z_mean + z_var / 2)
+    # Mean of R and y processes
+    R_mean = c_r * exp_z_mean + np.exp(μ_r + σ_r**2 / 2)
+    y_mean = c_y * exp_z_mean + np.exp(μ_y + σ_y**2 / 2)
+    # Test stability condition ensuring wealth does not diverge
+    # to infinity.
+    α = R_mean * s_0
+    if α >= 1:
+        raise ValueError("Stability condition failed.")
+    # Pack values into tuples and return them
+    household_params = (w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean)
+    aggregate_params = (a, b, σ_z, z_mean, z_var)
+    model = household_params, aggregate_params
+    return model
 ```
 
-is data representing the wealth of 10,000 households.
+Here's a function that generates the aggregate state process
 
-We can compute and plot the Lorenz curve as follows:
+```{code-cell} ipython3
+@numba.jit
+def generate_aggregate_state_sequence(aggregate_params, length=100):
+    a, b, σ_z, z_mean, z_var = aggregate_params 
+    z = np.empty(length+1)
+    z[0] = z_mean   # Initialize at z_mean
+    for t in range(length):
+        z[t+1] = a * z[t] + b + σ_z * np.random.randn()
+    return z
+```
+
+Here's a function that updates household wealth by one period, taking the
+current value of the aggregate shock
 
 ```{code-cell} ipython3
-%%time
-f_vals, l_vals = lorenz_curve_jax(w)
+@numba.jit
+def update_wealth(household_params, w, z):
+    """
+    Generate w_{t+1} given w_t and z_{t+1}.
+    """
+    # Unpack
+    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params
+    # Update wealth
+    y = c_y * np.exp(z) + np.exp(μ_y + σ_y * np.random.randn())
+    wp = y
+    if w >= w_hat:
+        R = c_r * np.exp(z) + np.exp(μ_r + σ_r * np.random.randn())
+        wp += R * s_0 * w
+    return wp
 ```
 
+Here's a function to simulate the time series of wealth for an individual household
+
 ```{code-cell} ipython3
-%%time
-# This will be much faster as it will use the jitted function
-f_vals, l_vals = lorenz_curve_jax(w)
+@numba.jit
+def wealth_time_series(model, w_0, sim_length):
+    """
+    Generate a single time series of length sim_length for wealth given initial
+    value w_0.  The function generates its own aggregate shock sequence.
+
+    """
+    # Unpack
+    household_params, aggregate_params = model
+    a, b, σ_z, z_mean, z_var = aggregate_params 
+    # Initialize and update
+    z = generate_aggregate_state_sequence(aggregate_params, 
+                                          length=sim_length)
+    w = np.empty(sim_length)
+    w[0] = w_0
+    for t in range(sim_length-1):
+        w[t+1] = update_wealth(household_params, w[t], z[t+1])
+    return w
+```
+
+Let's look at the wealth dynamics of an individual household
+
+```{code-cell} ipython3
+model = create_wealth_model()
+household_params, aggregate_params = model
+w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params
+a, b, σ_z, z_mean, z_var = aggregate_params 
+ts_length = 200
+w = wealth_time_series(model, y_mean, ts_length)
 ```
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
-ax.plot(f_vals, l_vals, label='Lorenz curve, lognormal sample')
-ax.plot(f_vals, f_vals, label='Lorenz curve, equality')
-ax.legend()
+ax.plot(w)
 plt.show()
 ```
 
-Here is another example, which shows how the Lorenz curve shifts as the
-underlying distribution changes.
+Notice the large spikes in wealth over time.
 
-We generate 10,000 observations using the Pareto distribution with a range of
-parameters, and then compute the Lorenz curve corresponding to each set of
-observations.
+Such spikes are related to heavy tails in the wealth distribution, which we
+discuss below.
+
+Here's a function to simulate a cross section of households forward in time.
+
+Note the use of parallelization to speed up computation.
 
 ```{code-cell} ipython3
-a_vals = (1, 2, 5)              # Pareto tail index
-n = 10_000                      # size of each sample
+@numba.jit(parallel=True)
+def update_cross_section(model, w_distribution, z_sequence):
+    """
+    Shifts a cross-section of households forward in time
+
+    Takes 
+
+        * a current distribution of wealth values as w_distribution and
+        * an aggregate shock sequence z_sequence
+
+    and updates each w_t in w_distribution to w_{t+j}, where
+    j = len(z_sequence).
+
+    Returns the new distribution.
+
+    """
+    # Unpack
+    household_params, aggregate_params = model
+
+    num_households = len(w_distribution)
+    new_distribution = np.empty_like(w_distribution)
+    z = z_sequence
+
+    # Update each household
+    for i in numba.prange(num_households):
+        w = w_distribution[i]
+        for t in range(sim_length):
+            w = update_wealth(household_params, w, z[t])
+        new_distribution[i] = w
+    return new_distribution
 ```
 
+Parallelization works in the function above because the time path of each
+household can be calculated independently once the path for the aggregate state
+is known.
+
+Let's see how long it takes to shift a large cross-section of households forward
+200 periods
+
 ```{code-cell} ipython3
-fig, ax = plt.subplots()
-for a in a_vals:
-    rand_key = jax.random.PRNGKey(a*100)
-    u = jax.random.uniform(rand_key, shape=(n,))
-    y = u**(-1/a)                    # distributed as Pareto with tail index a
-    f_vals, l_vals = lorenz_curve_jax(y)
-    ax.plot(f_vals, l_vals, label=f'$a = {a}$')
-
-ax.plot(f_vals, f_vals, label='equality')
-ax.legend()
-plt.show()
+sim_length = 200
+num_households = 10_000_000
+ψ_0 = np.full(num_households, y_mean)  # Initial distribution
+z_sequence = generate_aggregate_state_sequence(aggregate_params,
+                                               length=sim_length)
+print("Generating cross-section using Numba")
+start_time = time()
+ψ_star = update_cross_section(model, ψ_0, z_sequence)
+numba_time = time() - start_time
+print(f"Generated cross-section in {numba_time} seconds.\n")
 ```
 
-You can see that, as the tail parameter of the Pareto distribution increases, inequality decreases.
+### JAX implementation
 
-This is to be expected, because a higher tail index implies less weight in the tail of the Pareto distribution.
+Let's redo some of the preceding calculations using JAX and see how execution
+speed compares
 
-### The Gini Coefficient
+```{code-cell} ipython3
+def update_cross_section_jax(model, w_distribution, z_sequence, key):
+    """
+    Shifts a cross-section of households forward in time
 
-The definition and interpretation of the Gini coefficient can be found on the corresponding [Wikipedia page](https://en.wikipedia.org/wiki/Gini_coefficient).
+    Takes 
 
-We can test it on the Weibull distribution with parameter $a$, where the Gini coefficient is known to be
+        * a current distribution of wealth values as w_distribution and
+        * an aggregate shock sequence z_sequence
 
-$$
-G = 1 - 2^{-1/a}
-$$
+    and updates each w_t in w_distribution to w_{t+j}, where
+    j = len(z_sequence).
 
-Let's define a function to compute the Gini coefficient.
+    Returns the new distribution.
 
-```{code-cell} ipython3
-@jax.jit
-def gini_jax(y):
-    n = y.shape[0]
-    g_sum = 0
+    """
+    # Unpack, simplify names
+    household_params, aggregate_params = model
+    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params
+    w = w_distribution
+    n = len(w)
 
-    def sum_y_gini(i, g_sum):
-        g_sum += jnp.sum(jnp.abs(y[i] - y))
-        return g_sum
-    
-    g_sum = jax.lax.fori_loop(0, n, sum_y_gini, 0)
-    return g_sum / (2 * n * jnp.sum(y))
+    # Update wealth
+    for t, z in enumerate(z_sequence):
+        U = jax.random.normal(key, (2, n))
+        y = c_y * jnp.exp(z) + jnp.exp(μ_y + σ_y * U[0, :])
+        R = c_r * jnp.exp(z) + jnp.exp(μ_r + σ_r * U[1, :])
+        w = y + jnp.where(w < w_hat, 0.0, R * s_0 * w) 
+        key, subkey = jax.random.split(key)
+
+    return w
 ```
 
-Let's see if the Gini coefficient computed from a simulated sample matches
-this at each fixed value of $a$.
+Let's see how long it takes to shift the cross-section of households forward
+using JAX
 
 ```{code-cell} ipython3
-a_vals = range(1, 20)
-ginis = []
-ginis_theoretical = []
-n = 100
+sim_length = 200
+num_households = 10_000_000
+ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution
+z_sequence = generate_aggregate_state_sequence(aggregate_params,
+                                               length=sim_length)
+z_sequence = jnp.array(z_sequence)
+```
 
-for a in a_vals:
-    rand_key = jax.random.PRNGKey(a)
-    y = jax.random.weibull_min(rand_key, 1, a, shape=(n,))
-    ginis.append(gini_jax(y))
-    ginis_theoretical.append(1 - 2**(-1/a))
+```{code-cell} ipython3
+print("Generating cross-section using JAX")
+key = jax.random.PRNGKey(1234)
+start_time = time()
+ψ_star = update_cross_section_jax(model, ψ_0, z_sequence, key)
+jax_time = time() - start_time
+print(f"Generated cross-section in {jax_time} seconds.\n")
 ```
 
 ```{code-cell} ipython3
-fig, ax = plt.subplots()
-ax.plot(a_vals, ginis, label='estimated gini coefficient')
-ax.plot(a_vals, ginis_theoretical, label='theoretical gini coefficient')
-ax.legend()
-ax.set_xlabel("Weibull parameter $a$")
-ax.set_ylabel("Gini coefficient")
-plt.show()
+print("Repeating without compile time.")
+key = jax.random.PRNGKey(1234)
+start_time = time()
+ψ_star = update_cross_section_jax(model, ψ_0, z_sequence, key)
+jax_time = time() - start_time
+print(f"Generated cross-section in {jax_time} seconds")
 ```
 
-The simulation shows that the fit is good.
+And let's see how long it takes if we compile the loop.
 
-## A Model of Wealth Dynamics
+```{code-cell} ipython3
+def update_cross_section_jax_compiled(model, 
+                                      w_distribution, 
+                                      w_size,
+                                      z_sequence, 
+                                      key):
+    """
+    Shifts a cross-section of households forward in time
 
-Having discussed inequality measures, let us now turn to wealth dynamics.
+    Takes 
 
-The model we will study is
+        * a current distribution of wealth values as w_distribution and
+        * an aggregate shock sequence z_sequence
 
-```{math}
-:label: wealth_dynam_ah
-
-w_{t+1} = (1 + r_{t+1}) s(w_t) + y_{t+1}
-```
+    and updates each w_t in w_distribution to w_{t+j}, where
+    j = len(z_sequence).
 
-where
+    Returns the new distribution.
 
-- $w_t$ is wealth at time $t$ for a given household,
-- $r_t$ is the rate of return of financial assets,
-- $y_t$ is current non-financial (e.g., labor) income and
-- $s(w_t)$ is current wealth net of consumption
+    """
+    # Unpack, simplify names
+    household_params, aggregate_params = model
+    w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params
+    w = w_distribution
+    n = len(w)
+    z = z_sequence
+    sim_length = len(z)
 
-## Implementation using JAX
+    def body_function(t, state):
+        key, w = state
+        key, subkey = jax.random.split(key)
+        U = jax.random.normal(subkey, (2, n))
+        y = c_y * jnp.exp(z[t]) + jnp.exp(μ_y + σ_y * U[0, :])
+        R = c_r * jnp.exp(z[t]) + jnp.exp(μ_r + σ_r * U[1, :])
+        w = y + jnp.where(w < w_hat, 0.0, R * s_0 * w) 
+        return key, w
 
-Let's define a model to represent the wealth dynamics.
+    key, w = jax.lax.fori_loop(0, sim_length, body_function, (key, w))
+    return w
+```
 
 ```{code-cell} ipython3
-# NamedTuple Model
+update_cross_section_jax_compiled = jax.jit(
+        update_cross_section_jax_compiled, static_argnums=(2,)
+)
+```
 
-Model = namedtuple("Model", ("w_hat", "s_0", "c_y", "μ_y",
-                             "σ_y", "c_r", "μ_r", "σ_r", "a",
-                             "b", "σ_z", "z_mean", "z_var", "y_mean"))
+```{code-cell} ipython3
+print("Generating cross-section using JAX with compiled loop")
+key = jax.random.PRNGKey(1234)
+start_time = time()
+ψ_star = update_cross_section_jax_compiled(
+        model, ψ_0, num_households, z_sequence, key
+)
+jax_fori_time = time() - start_time
+print(f"Generated cross-section in {jax_fori_time} seconds.\n")
 ```
 
-Here's a function to create the Model with the given parameters
+```{code-cell} ipython3
+print("Repeating without compile time")
+key = jax.random.PRNGKey(1234)
+start_time = time()
+ψ_star = update_cross_section_jax_compiled(
+        model, ψ_0, num_households, z_sequence, key
+)
+jax_fori_time = time() - start_time
+print(f"Generated cross-section in {jax_fori_time} seconds")
+```
 
 ```{code-cell} ipython3
-def create_wealth_model(w_hat=1.0,
-                        s_0=0.75,
-                        c_y=1.0,
-                        μ_y=1.0,
-                        σ_y=0.2,
-                        c_r=0.05,
-                        μ_r=0.1,
-                        σ_r=0.5,
-                        a=0.5,
-                        b=0.0,
-                        σ_z=0.1):
-    """
-    Create a wealth model with given parameters and return
-    and instance of NamedTuple Model.
-    """
-    z_mean = b / (1 - a)
-    z_var = σ_z**2 / (1 - a**2)
-    exp_z_mean = jnp.exp(z_mean + z_var / 2)
-    R_mean = c_r * exp_z_mean + jnp.exp(μ_r + σ_r**2 / 2)
-    y_mean = c_y * exp_z_mean + jnp.exp(μ_y + σ_y**2 / 2)
-    # Test a stability condition that ensures wealth does not diverge
-    # to infinity.
-    α = R_mean * s_0
-    if α >= 1:
-        raise ValueError("Stability condition failed.")
-    return Model(w_hat=w_hat, s_0=s_0, c_y=c_y, μ_y=μ_y,
-                 σ_y=σ_y, c_r=c_r, μ_r=μ_r, σ_r=σ_r, a=a,
-                 b=b, σ_z=σ_z, z_mean=z_mean, z_var=z_var, y_mean=y_mean)
+print(f"JAX is {numba_time / jax_fori_time:.4f} times faster.\n")
 ```
 
-The following function updates one period with the given current wealth and persistent state.
+### Pareto tails
 
-```{code-cell} ipython3
-def update_states_jax(arrays, wdy, size, rand_key):
-    """
-    Update one period, given current wealth w and persistent
-    state z. They are stored in the form of tuples under the arrays argument
-    """
-    # Unpack w and z
-    w, z = arrays
+In most countries, the cross-sectional distribution of wealth exhibits a Pareto
+tail (power law).
 
-    rand_key, *subkey = jax.random.split(rand_key, 3)
-    zp = wdy.a * z + wdy.b + wdy.σ_z * jax.random.normal(rand_key, shape=size)
+Let's see if our model can replicate this stylized fact by running a simulation
+that generates a cross-section of wealth and generating a suitable rank-size plot.
 
-    # Update wealth
-    y = wdy.c_y * jnp.exp(zp) + jnp.exp(
-                        wdy.μ_y + wdy.σ_y * jax.random.normal(subkey[0], shape=size))
-    wp = y
+We will use the function `rank_size` from `quantecon` library.
+
+In the limit, data that obeys a power law generates a straight line.
+
+```{code-cell} ipython3
+model = create_wealth_model()
+key = jax.random.PRNGKey(1234)
+ψ_star = update_cross_section_jax_compiled(
+        model, ψ_0, num_households, z_sequence, key
+)
+fig, ax = plt.subplots()
 
-    R = wdy.c_r * jnp.exp(zp) + jnp.exp(
-                        wdy.μ_r + wdy.σ_r * jax.random.normal(subkey[1], shape=size))
-    wp += (w >= wdy.w_hat) * R * wdy.s_0 * w
-    return wp, zp
+rank_data, size_data = qe.rank_size(ψ_star, c=0.001)
+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()
 ```
 
-Here’s function to simulate the time series of wealth for individual households using a `for` loop and JAX.
+### Lorenz curves and Gini coefficients
 
-```{code-cell} ipython3
-# Using JAX and for loop
+To study the impact of parameters on inequality, we examine Lorenz curves
+and the Gini coefficients at different parameters.
 
-def wealth_time_series_for_loop_jax(w_0, n, wdy, size, rand_seed=1):
-    """
-    Generate a single time series of length n for wealth given
-    initial value w_0.
+QuantEcon provides functions to compute Lorenz curves and Gini coefficients that are accelerated using Numba.
 
-    * This implementation uses a `for` loop.
+Here we provide JAX-based functions that do the same job and are faster for large data sets on parallel hardware.
 
-    The initial persistent state z_0 for each household is drawn from
-    the stationary distribution of the AR(1) process.
 
-        * wdy: NamedTuple Model
-        * w_0: scalar/vector
-        * n: int
-        * size: size/shape of the w_0
-        * rand_seed: int (Used to generate PRNG key)
-    """
-    rand_key = jax.random.PRNGKey(rand_seed)
-    rand_key, *subkey = jax.random.split(rand_key, n)
+#### Lorenz curve
 
-    w_0 = jax.device_put(w_0).reshape(size)
+Recall that, for sorted data $w_1, \ldots, w_n$, the Lorenz curve 
+generates data points $(x_i, y_i)_{i=0}^n$  according to
 
-    z = wdy.z_mean + jnp.sqrt(wdy.z_var) * jax.random.normal(rand_key, shape=size)
-    w = [w_0]
-    for t in range(n-1):
-        w_, z = update_states_jax((w[t], z), wdy, size, subkey[t])
-        w.append(w_)
-    return jnp.array(w)
+$$
+    x_0 = y_0 = 0
+    \qquad \text{and, for $i \geq 1$,} \quad
+    x_i = \frac{i}{n},
+    \qquad
+    y_i =
+       \frac{\sum_{j \leq i} w_j}{\sum_{j \leq n} w_j}  
+$$
+
+```{code-cell} ipython3
+def _lorenz_curve_jax(w, w_size):
+    n = w.shape[0]
+    w = jnp.sort(w)
+    x = jnp.arange(n + 1) / n
+    s = jnp.concatenate((jnp.zeros(1), jnp.cumsum(w)))
+    y = s / s[n]
+    return x, y
+
+lorenz_curve_jax = jax.jit(_lorenz_curve_jax, static_argnums=(1,))
 ```
 
-Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution using the above function.
+Let's test
 
 ```{code-cell} ipython3
-wdy = create_wealth_model() # default model
-ts_length = 200
-size = (1,)
+sim_length = 200
+num_households = 1_000_000
+ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution
+z_sequence = generate_aggregate_state_sequence(aggregate_params,
+                                               length=sim_length)
+z_sequence = jnp.array(z_sequence)
+```
+
+```{code-cell} ipython3
+key = jax.random.PRNGKey(1234)
+ψ_star = update_cross_section_jax_compiled(
+        model, ψ_0, num_households, z_sequence, key
+)
 ```
 
 ```{code-cell} ipython3
-%%time
-w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean,
-                    ts_length, wdy, size).block_until_ready()
+%time _ = lorenz_curve_jax(ψ_star, num_households)
 ```
 
 ```{code-cell} ipython3
-%%time
-w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean,
-                    ts_length, wdy, size).block_until_ready()
+%time x, y = lorenz_curve_jax(ψ_star, num_households)
 ```
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
-ax.plot(w_jax_result)
+ax.plot(x, y, label="Lorenz curve at defaults")
+ax.plot(x, x, 'k-', lw=1)
+ax.legend()
 plt.show()
 ```
 
-We can further try to optimize and speed up the compile time of the above function by replacing `for` loop with [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html).
+#### Gini Coefficient
+
+Recall that, for sorted data $w_1, \ldots, w_n$, the Gini coefficient takes the form
+
+
+$$
+G :=
+\frac
+    {\sum_{i=1}^n \sum_{j = 1}^n |w_j - w_i|}
+    {2n\sum_{i=1}^n w_i}.
+$$ (eq:gini)
+
+Here's a function that computes the Gini coefficient using vectorization.
 
 ```{code-cell} ipython3
-def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1):
-    """
-    Generate a single time series of length n for wealth given
-    initial value w_0.
+def _gini_jax(w, w_size):
+    w_1 = jnp.reshape(w, (w_size, 1))
+    w_2 = jnp.reshape(w, (1, w_size))
+    g_sum = jnp.sum(jnp.abs(w_1 - w_2))
+    return g_sum / (2 * w_size * jnp.sum(w))
 
-    * This implementation uses `jax.lax.scan`.
+gini_jax = jax.jit(_gini_jax, static_argnums=(1,))
+```
 
-    The initial persistent state z_0 for each household is drawn from
-    the stationary distribution of the AR(1) process.
+```{code-cell} ipython3
+%time gini = gini_jax(ψ_star, num_households).block_until_ready()
+```
 
-        * wdy: NamedTuple Model
-        * w_0: scalar/vector
-        * n: int
-        * size: size/shape of the w_0
-        * rand_seed: int (Used to generate PRNG key)
-    """
-    rand_key = jax.random.PRNGKey(rand_seed)
-    rand_key, *subkey = jax.random.split(rand_key, n)
+```{code-cell} ipython3
+%time gini = gini_jax(ψ_star, num_households).block_until_ready()
+gini
+```
 
-    w_0 = jax.device_put(w_0).reshape(size)
-    z_init = wdy.z_mean + jnp.sqrt(wdy.z_var) * jax.random.normal(rand_key, shape=size)
-    arrays = w_0, z_init
-    rand_sub_keys = jnp.array(subkey)
+## Exercises
 
-    w_final = jnp.array([w_0])
+```{exercise}
+:label: wd_ex1
 
-    # Define the function for each update
-    def update_w_z(arrays, rand_sub_key):
-        wp, zp = update_states_jax(arrays, wdy, size, rand_sub_key)
-        return (wp, zp), wp
+In this exercise, write an alternative version of `gini_jax` that uses `vmap` instead of reshaping and broadcasting.
 
-    arrays_last, w_values = jax.lax.scan(update_w_z, arrays, rand_sub_keys)
-    return jnp.concatenate((w_final, w_values))
+Test with the same array to see if you can obtain the same output
+```
 
-# Create the jit function
-wealth_time_series_jax = jax.jit(wealth_time_series_jax, static_argnums=(1,3,))
+```{solution-start} wd_ex1
+:class: dropdown
 ```
 
-Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution and also observe the difference in time between `wealth_time_series_jax` and `wealth_time_series_for_loop_jax`.
+Here's one solution:
 
 ```{code-cell} ipython3
-wdy = create_wealth_model() # default model
-ts_length = 200
-size = (1,)
+@jax.jit
+def gini_jax_vmap(w):
+
+    def _inner_sum(x):
+        return jnp.sum(jnp.abs(x - w))
+    
+    inner_sum = jax.vmap(_inner_sum)
+    
+    full_sum = jnp.sum(inner_sum(w))
+    return full_sum / (2 * len(w) * jnp.sum(w))
 ```
 
 ```{code-cell} ipython3
-%%time
-w_jax_result = wealth_time_series_jax(wdy.y_mean, 
-                    ts_length, wdy, size).block_until_ready()
+%time gini = gini_jax_vmap(ψ_star).block_until_ready()
+gini
 ```
 
-Running the above function again will be even faster because of JAX's JIT.
-
 ```{code-cell} ipython3
-%%time
-# 2nd time is expected to be very fast because of JIT
-w_jax_result = wealth_time_series_jax(wdy.y_mean, 
-                    ts_length, wdy, size).block_until_ready()
+%time gini = gini_jax_vmap(ψ_star).block_until_ready()
+gini
 ```
 
-```{code-cell} ipython3
-fig, ax = plt.subplots()
-ax.plot(w_jax_result)
-plt.show()
+```{solution-end}
 ```
 
-Now here’s function to simulate a cross section of households forward in time.
 
-```{code-cell} ipython3
-def update_cross_section_jax(w_distribution, shift_length, wdy, size, rand_seed=2):
-    """
-    Shifts a cross-section of household forward in time
+```{exercise-start}
+:label: wd_ex2
+```
 
-    * wdy: NamedTuple Model
-    * w_distribution: array_like, represents current cross-section
+In this exercise we investigate how the parameters determining the rate of return on assets and labor income shape inequality.
 
-    Takes a current distribution of wealth values as w_distribution
-    and updates each w_t in w_distribution to w_{t+j}, where
-    j = shift_length.
+In doing so we recall that
 
-    Returns the new distribution.
-    """
-    new_dist = wealth_time_series_jax(w_distribution, shift_length, wdy, size, rand_seed)
-    new_distribution = new_dist[-1, :]
-    return new_distribution
+$$
+    R_t := 1 + r_t = c_r \exp(z_t) + \exp(\mu_r + \sigma_r \xi_t)
+$$
 
+while
 
-# Create the jit function
-update_cross_section_jax = jax.jit(update_cross_section_jax, static_argnums=(1,3,))
-```
+$$
+    y_t = c_y \exp(z_t) + \exp(\mu_y + \sigma_y \zeta_t)
+$$
 
-## Applications
+Investigate how the Lorenz curves and the Gini coefficient associated with the wealth distribution change as return to savings varies.
 
-Let's try simulating the model at different parameter values and investigate
-the implications for the wealth distribution.
+In particular, plot Lorenz curves for the following three different values of $\mu_r$
 
+```{code-cell} ipython3
+μ_r_vals = (0.0, 0.025, 0.05)
+```
 
-### Inequality Measures
+Use the following as your initial cross-sectional distribution
 
-Let's look at how inequality varies with returns on financial assets.
+```{code-cell} ipython3
+num_households = 1_000_000
+ψ_0 = jnp.full(num_households, y_mean)  # Initial distribution
+```
 
-The next function generates a cross section and then computes the Lorenz
-curve and Gini coefficient.
+Once you have done that, plot the Gini coefficients as well.
 
-```{code-cell} ipython3
-def generate_lorenz_and_gini_jax(wdy, num_households=100_000, T=500):
-    """
-    Generate the Lorenz curve data and gini coefficient corresponding to a
-    WealthDynamics mode by simulating num_households forward to time T.
-    """
-    size = (num_households, )
-    ψ_0 = jnp.full(size, wdy.y_mean)
-    ψ_star = update_cross_section_jax(ψ_0, T, wdy, size)
-    return gini_jax(ψ_star), lorenz_curve_jax(ψ_star)
+Do the outcomes match your intuition?
 
-# Create the jit function
-generate_lorenz_and_gini_jax = jax.jit(generate_lorenz_and_gini_jax,
-                                       static_argnums=(1,2,))
+```{exercise-end}
 ```
 
-Now we investigate how the Lorenz curves associated with the wealth distribution change as return to savings varies.
+```{solution-start} wd_ex2
+:class: dropdown
+```
 
-The code below plots Lorenz curves for three different values of $\mu_r$.
+Here is one solution
 
 ```{code-cell} ipython3
-%%time
-
+key = jax.random.PRNGKey(1234)
 fig, ax = plt.subplots()
-μ_r_vals = (0.0, 0.025, 0.05)
 gini_vals = []
-
 for μ_r in μ_r_vals:
-    wdy = create_wealth_model(μ_r=μ_r)
-    gv, (f_vals, l_vals) = generate_lorenz_and_gini_jax(wdy)
-    ax.plot(f_vals, l_vals, label=f'$\psi^*$ at $\mu_r = {μ_r:0.2}$')
-    gini_vals.append(gv)
-
-ax.plot(f_vals, f_vals, label='equality')
+    model = create_wealth_model(μ_r=μ_r)
+    ψ_star = update_cross_section_jax_compiled(
+            model, ψ_0, num_households, z_sequence, key
+    )
+    x, y = lorenz_curve_jax(ψ_star, num_households)
+    g = gini_jax(ψ_star, num_households)
+    ax.plot(x, y, label=f'$\psi^*$ at $\mu_r = {μ_r:0.2}$')
+    gini_vals.append(g)
+ax.plot(x, y, label='equality')
 ax.legend(loc="upper left")
 plt.show()
 ```
 
 The Lorenz curve shifts downwards as returns on financial income rise, indicating a rise in inequality.
 
-Now let's check the Gini coefficient.
+Now let's check the Gini coefficient
 
 ```{code-cell} ipython3
 fig, ax = plt.subplots()
-ax.plot(μ_r_vals, gini_vals, label='gini coefficient')
+ax.plot(μ_r_vals, gini_vals, label='Gini coefficient')
 ax.set_xlabel("$\mu_r$")
 ax.legend()
 plt.show()
 ```
 
-Once again, we see that inequality increases as returns on financial income
-rise.
-
-Let's finish this section by investigating what happens when we change the
-volatility term $\sigma_r$ in financial returns.
-
-```{code-cell} ipython3
-%%time
-
-fig, ax = plt.subplots()
-σ_r_vals = (0.35, 0.45, 0.52)
-gini_vals = []
-
-for σ_r in σ_r_vals:
-    wdy = create_wealth_model(σ_r=σ_r)
-    gv, (f_vals, l_vals) = generate_lorenz_and_gini_jax(wdy)
-    ax.plot(f_vals, l_vals, label=f'$\psi^*$ at $\sigma_r = {σ_r:0.2}$')
-    gini_vals.append(gv)
+As expected, inequality increases as returns on financial income rise.
 
-ax.plot(f_vals, f_vals, label='equality')
-ax.legend(loc="upper left")
-plt.show()
+```{solution-end}
 ```
 
-We see that greater volatility has the effect of increasing inequality in this model.
-
-## Exercises
-
-```{exercise}
-:label: wd_ex1
-
-For a wealth or income distribution with Pareto tail, a higher tail index suggests lower inequality.
+```{exercise-start}
+:label: wd_ex3
+```
 
-Indeed, it is possible to prove that the Gini coefficient of the Pareto
-distribution with tail index $a$ is $1/(2a - 1)$.
+Now investigate what happens when we change the volatility term $\sigma_r$ in financial returns.
 
-To the extent that you can, confirm this by simulation.
+Use the same initial condition as before and the sequence
 
-In particular, generate a plot of the Gini coefficient against the tail index
-using both the theoretical value just given and the value computed from a sample via `gini_jax`.
+```{code-cell} ipython3
+σ_r_vals = (0.35, 0.45, 0.52)
+```
 
-For the values of the tail index, use `a_vals = jnp.linspace(1, 10, 25)`.
+To isolate the role of volatility, set $\mu_r = - \sigma_r^2 / 2$ at each $\sigma_r$.
 
-Use sample of size 1,000 for each $a$ and the sampling method for generating Pareto draws employed in the discussion of Lorenz curves for the Pareto distribution.
+(This holds the variance of the idiosyncratic term $\exp(\mu_r + \sigma_r \zeta)$ constant.)
 
-To the extent that you can, interpret the monotone relationship between the
-Gini index and $a$.
+```{exercise-end}
 ```
 
-```{solution-start} wd_ex1
+```{solution-start} wd_ex3
 :class: dropdown
 ```
 
-Here is one solution, which produces a good match between theory and
-simulation.
+Here's one solution
 
 ```{code-cell} ipython3
-a_vals = jnp.linspace(1, 10, 25)  # Pareto tail index
-ginis = []
+key = jax.random.PRNGKey(1234)
+fig, ax = plt.subplots()
+
+gini_vals = []
+for σ_r in σ_r_vals:
+    model = create_wealth_model(σ_r=σ_r, μ_r=(-σ_r**2/2))
+    ψ_star = update_cross_section_jax_compiled(
+            model, ψ_0, num_households, z_sequence, key
+    )
+    x, y = lorenz_curve_jax(ψ_star, num_households)
+    g = gini_jax(ψ_star, num_households)
+    ax.plot(x, y, label=f'$\psi^*$ at $\sigma_r = {σ_r:0.2}$')
+    gini_vals.append(g)
+ax.plot(x, y, label='equality')
+ax.legend(loc="upper left")
+plt.show()
+```
 
-n = 1000                         # size of each sample
+```{code-cell} ipython3
 fig, ax = plt.subplots()
-for i, a in enumerate(a_vals):
-    rand_key = jax.random.PRNGKey(i*10)
-    u = jax.random.uniform(rand_key, shape=(n,))
-    y = u**(-1/a)
-    ginis.append(gini_jax(y))
-ax.plot(a_vals, ginis, label='sampled')
-ax.plot(a_vals, 1/(2*a_vals - 1), label='theoretical')
+ax.plot(σ_r_vals, gini_vals, label='Gini coefficient')
+ax.set_xlabel("$\sigma_r$")
 ax.legend()
 plt.show()
 ```
 
-In general, for a Pareto distribution, a higher tail index implies less weight
-in the right hand tail.
-
-This means less extreme values for wealth and hence more equality.
-
-More equality translates to a lower Gini index.
-
 ```{solution-end}
 ```
 
 ```{exercise-start}
-:label: wd_ex2
+:label: wd_ex4
 ```
+In this exercise, examine which has more impact on inequality:
 
-When savings is constant, the wealth process has the same quasi-linear
-structure as a Kesten process, with multiplicative and additive shocks.
-
-The Kesten--Goldie theorem tells us that Kesten processes have Pareto tails under a range of parameterizations.
-
-The theorem does not directly apply here, since savings is not always constant and since the multiplicative and additive terms in {eq}`wealth_dynam_ah` are not IID.
-
-At the same time, given the similarities, perhaps Pareto tails will arise.
-
-To test this, run a simulation that generates a cross-section of wealth and
-generate a rank-size plot.
+- a 5% rise in volatility of the rate of return, 
+- or a 5% rise in volatility of labor income.
 
-In viewing the plot, remember that Pareto tails generate a straight line.  Is
-this what you see?
+Test this by 
 
-For sample size and initial conditions, use
+1. Shifting $\sigma_r$ up 5% from the baseline and plotting the Lorenz curve
+1. Shifting $\sigma_y$ up 5% from the baseline and plotting the Lorenz curve
 
-```{code-cell} ipython3
-num_households = 250_000
-T = 500 # Shift forward T periods
-ψ_0 = jnp.full((num_households, ), wdy.y_mean) # Initial distribution
-```
+Plot both on the same figure and examine the result.
 
 ```{exercise-end}
 ```
 
-```{solution-start} wd_ex2
+```{solution-start} wd_ex4
 :class: dropdown
 ```
 
-First let's generate the distribution:
-
-```{code-cell} ipython3
-num_households = 250_000
-T = 500  # how far to shift forward in time
-size = (num_households, )
-
-wdy = create_wealth_model()
-ψ_0 = jnp.full(size, wdy.y_mean)
-ψ_star = update_cross_section_jax(ψ_0, T, wdy, size)
-```
+Here's one solution.
 
-Let's define a function to get the rank data
+It shows that increasing volatility in financial income has a greater effect
 
 ```{code-cell} ipython3
-def rank_size(data, c=1):
-    w = -jnp.sort(-data)                  # Reverse sort
-    w = w[:int(len(w) * c)]               # extract top (c * 100)%
-    rank_data = jnp.arange(len(w)) + 1
-    size_data = w
-    return rank_data, size_data
-```
+model = create_wealth_model()
+household_params, aggregate_params = model
+w_hat, s_0, c_y, μ_y, σ_y, c_r, μ_r, σ_r, y_mean = household_params
+σ_r_default = σ_r
+σ_y_default = σ_y
 
-Now let's see the rank-size plot:
+ψ_star = update_cross_section_jax_compiled(
+            model, ψ_0, num_households, z_sequence, key
+)
+x_default, y_default = lorenz_curve_jax(ψ_star, num_households)
 
-```{code-cell} ipython3
-fig, ax = plt.subplots()
+model = create_wealth_model(σ_r=(1.05 * σ_r_default))
+ψ_star = update_cross_section_jax_compiled(
+            model, ψ_0, num_households, z_sequence, key
+)
+x_financial, y_financial = lorenz_curve_jax(ψ_star, num_households)
 
-rank_data, size_data = rank_size(ψ_star, c=0.001)
-ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)
-ax.set_xlabel("log rank")
-ax.set_ylabel("log size")
+model = create_wealth_model(σ_y=(1.05 * σ_y_default))
+ψ_star = update_cross_section_jax_compiled(
+            model, ψ_0, num_households, z_sequence, key
+)
+x_labor, y_labor = lorenz_curve_jax(ψ_star, num_households)
 
+fig, ax = plt.subplots()    
+ax.plot(x_default, x_default, 'k-', lw=1, label='equality')
+ax.plot(x_financial, y_financial, label=r'higher $\sigma_r$')
+ax.plot(x_labor, y_labor, label=r'higher $\sigma_y$')
+ax.legend()
 plt.show()
 ```
-
 ```{solution-end}
-```
+```
\ No newline at end of file