diff --git a/lectures/status.md b/lectures/status.md
index b8a415a8f..d37a0b808 100644
--- a/lectures/status.md
+++ b/lectures/status.md
@@ -16,5 +16,8 @@ This table contains the latest execution statistics.
 ```{nb-exec-table}
 ```
 
-These lectures are built on `linux` instances through `github actions` so are
-executed using the following [hardware specifications](https://docs.github.com/en/actions/using-github-hosted-runners/about-github-hosted-runners#supported-runners-and-hardware-resources)
\ No newline at end of file
+(status:machine-details)=
+
+These lectures are built on `linux` instances through `github actions`  and `amazon web services (aws)` to
+enable access to a `gpu`. These lectures are built on a [p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/)
+that has access to `8 vcpu's`, a `V100 NVIDIA Tesla GPU`, and `61 Gb` of memory.
diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md
index c2a9954b4..6bf34e132 100644
--- a/lectures/wealth_dynamics.md
+++ b/lectures/wealth_dynamics.md
@@ -19,6 +19,16 @@ kernelspec:
 
 # Wealth Distribution Dynamics
 
+```{admonition} GPU in use
+:class: warning
+
+This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming.
+
+Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.
+
+Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax#pip-installation-gpu-cuda) for installing JAX with GPU support.
+```
+
 ```{contents} Contents
 :depth: 2
 ```
@@ -30,6 +40,9 @@ In addition to what's in Anaconda, this lecture will need the following librarie
 tags: [hide-output]
 ---
 !pip install quantecon
+# If your machine has CUDA support, please follow the guide in GPU Warning.
+# Otherwise, run the line below:
+!pip install --upgrade "jax[CPU]"
 ```
 
 ## Overview
@@ -78,6 +91,19 @@ import numpy as np
 import quantecon as qe
 from numba import njit, float64, prange
 from numba.experimental import jitclass
+import jax
+import jax.numpy as jnp
+from collections import namedtuple
+from myst_nb import glue
+```
+
+Let's check the backend used by JAX and the devices available
+
+```{code-cell} ipython3
+# Check if JAX is using GPU
+print(f"JAX backend: {jax.devices()[0].platform}")
+# Check the devices available for JAX
+print(jax.devices())
 ```
 
 ## Lorenz Curves and the Gini Coefficient
@@ -250,11 +276,240 @@ their wealth.
 We are using something akin to a fixed savings rate model, while
 acknowledging that low wealth households tend to save very little.
 
-## Implementation
+## Implementation using JAX
+
+Let's define a model to represent the wealth dynamics.
+
+```{code-cell} ipython3
+# NamedTuple Model
+Model = namedtuple("Model", ("w_hat", "s_0", "c_y", "μ_y",
+                             "σ_y", "c_r", "μ_r", "σ_r", "a",
+                             "b", "σ_z", "z_mean", "z_var", "y_mean"))
+```
+
+Here's a function to create the Model with the given parameters
+
+```{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 = np.exp(z_mean + z_var / 2)
+    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 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)
+```
+
+The following function updates one period with the given current wealth and persistent state.
+
+```{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
+
+    rand_key, *subkey = jax.random.split(rand_key, 3)
+    zp = wdy.a * z + wdy.b + wdy.σ_z * jax.random.normal(rand_key, shape=size)
+
+    # Update wealth
+    y = wdy.c_y * jnp.exp(zp) + jnp.exp(wdy.μ_y + wdy.σ_y * jax.random.normal(subkey[0], shape=size))
+    wp = y
+
+    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
+
+# Create the jit function
+update_states_jax = jax.jit(update_states_jax, static_argnums=(2,))
+```
+
+Here’s function to simulate the time series of wealth for individual households using a `for` loop and JAX.
+
+```{code-cell} ipython3
+# Using JAX and for loop
+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.
+
+    * This implementation uses a for loop.
+
+    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)
+
+    w_0 = jax.device_put(w_0).reshape(size)
+
+    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)
+
+# Create the jit function
+wealth_time_series_for_loop_jax = jax.jit(wealth_time_series_for_loop_jax, static_argnums=(1,3,))
+```
+
+Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution using the above function.
+
+```{code-cell} ipython3
+wdy = create_wealth_model() # default model
+ts_length = 200
+size = (1,)
+```
+
+```{code-cell} ipython3
+qe.tic()
+w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready()
+qe.toc()
+```
+
+Running the above function again will be even faster because of JAX's JIT.
+
+```{code-cell} ipython3
+qe.tic()
+# 2nd time is expected to be very fast because of JIT
+w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready()
+qe.toc()
+```
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+ax.plot(w_jax_result)
+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).
+
+```{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.
+
+    * This implementation uses for jax.lax.scan
+
+    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)
+
+    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)
+
+    w_final = jnp.array([w_0])
+
+    # 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
+
+    arrays_last, w_values = jax.lax.scan(update_w_z, arrays, rand_sub_keys)
+    return jnp.concatenate((w_final, w_values))
+
+# Create the jit function
+wealth_time_series_jax = jax.jit(wealth_time_series_jax, static_argnums=(1,3,))
+```
+
+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`.
+
+```{code-cell} ipython3
+wdy = create_wealth_model() # default model
+ts_length = 200
+size = (1,)
+```
+
+```{code-cell} ipython3
+qe.tic()
+w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready()
+glue("wealth_time_series_jax_time_1", qe.toc())
+```
+
+Running the above function again will be even faster because of JAX's JIT.
+
+```{code-cell} ipython3
+qe.tic()
+# 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()
+glue("wealth_time_series_jax_time_2", qe.toc())
+```
+
+```{code-cell} ipython3
+fig, ax = plt.subplots()
+ax.plot(w_jax_result)
+plt.show()
+```
+
+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
+
+    * wdy: NamedTuple Model
+    * w_distribution: array_like, represents current cross-section
+
+    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.
+
+    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
+
+# Create the jit function
+update_cross_section_jax = jax.jit(update_cross_section_jax, static_argnums=(1,3,))
+```
+## Implementation using Numba
 
 Here's some type information to help Numba.
 
 ```{code-cell} ipython3
+
 wealth_dynamics_data = [
     ('w_hat',  float64),    # savings parameter
     ('s_0',    float64),    # savings parameter
@@ -278,6 +533,7 @@ Here's a class that stores instance data and implements methods that update
 the aggregate state and household wealth.
 
 ```{code-cell} ipython3
+
 @jitclass(wealth_dynamics_data)
 class WealthDynamics:
 
@@ -345,6 +601,7 @@ class WealthDynamics:
 Here's function to simulate the time series of wealth for in individual households.
 
 ```{code-cell} ipython3
+
 @njit
 def wealth_time_series(wdy, w_0, n):
     """
@@ -373,6 +630,7 @@ Now here's function to simulate a cross section of households forward in time.
 Note the use of parallelization to speed up computation.
 
 ```{code-cell} ipython3
+
 @njit(parallel=True)
 def update_cross_section(wdy, w_distribution, shift_length=500):
     """
@@ -411,14 +669,32 @@ the implications for the wealth distribution.
 
 ### Time Series
 
-Let's look at the wealth dynamics of an individual household.
+Let's look at the wealth dynamics of an individual household using numba.
 
 ```{code-cell} ipython3
 wdy = WealthDynamics()
 
 ts_length = 200
+```
+
+```{code-cell} ipython3
+qe.tic()
+w = wealth_time_series(wdy, wdy.y_mean, ts_length)
+glue("wealth_time_series_time_1", qe.toc())
+```
+
+```{code-cell} ipython3
+qe.tic()
+# Check the time for 2nd execution
 w = wealth_time_series(wdy, wdy.y_mean, ts_length)
+glue("wealth_time_series_time_2", qe.toc())
+```
+
+Notice the time difference between the `wealth_time_series`: {glue:}`wealth_time_series_time_1` and `wealth_time_series_jax`: {glue:}`wealth_time_series_jax_time_1`
+
+
 
+```{code-cell} ipython3
 fig, ax = plt.subplots()
 ax.plot(w)
 plt.show()
@@ -435,7 +711,26 @@ Let's look at how inequality varies with returns on financial assets.
 The next function generates a cross section and then computes the Lorenz
 curve and Gini coefficient.
 
+Let's first write the function that uses the JAX implementation and then for the numba.
+
+```{code-cell} ipython3
+# Uses JAX
+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)
+    ψ_star = jax.device_get(ψ_star)  # get back to numpy form
+    return qe.gini_coefficient(ψ_star), qe.lorenz_curve(ψ_star)
+```
+
+The following function uses the numba implementation
+
 ```{code-cell} ipython3
+# Uses numba
 def generate_lorenz_and_gini(wdy, num_households=100_000, T=500):
     """
     Generate the Lorenz curve data and gini coefficient corresponding to a
@@ -459,6 +754,28 @@ This is unavoidable because we are executing a CPU intensive task.
 In fact the code, which is JIT compiled and parallelized, runs extremely fast relative to the number of computations.
 
 ```{code-cell} ipython3
+%%time
+
+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')
+ax.legend(loc="upper left")
+plt.show()
+```
+
+Now let's try to run the same code snippet but using the numba version.
+
+```{code-cell} ipython3
+%%time
+
 fig, ax = plt.subplots()
 μ_r_vals = (0.0, 0.025, 0.05)
 gini_vals = []
@@ -480,15 +797,6 @@ We will look at this again via the Gini coefficient immediately below, but
 first consider the following image of our system resources when the code above
 is executing:
 
-(htop_again)=
-```{figure} /_static/lecture_specific/wealth_dynamics/htop_again.png
-:scale: 80
-```
-
-Notice how effectively Numba has implemented multithreading for this routine:
-all 8 CPUs on our workstation are running at maximum capacity (even though
-four of them are virtual).
-
 Since the code is both efficiently JIT compiled and fully parallelized, it's
 close to impossible to make this sequence of tasks run faster without changing
 hardware.
@@ -509,7 +817,31 @@ rise.
 Let's finish this section by investigating what happens when we change the
 volatility term $\sigma_r$ in financial returns.
 
+Firstly, using JAX, we have
+
 ```{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)
+
+ax.plot(f_vals, f_vals, label='equality')
+ax.legend(loc="upper left")
+plt.show()
+```
+
+Using numba, we get,
+
+```{code-cell} ipython3
+%%time
+
 fig, ax = plt.subplots()
 σ_r_vals = (0.35, 0.45, 0.52)
 gini_vals = []
@@ -625,9 +957,25 @@ z_0 = wdy.z_mean
 
 First let's generate the distribution:
 
+Using the JAX implementation
+
 ```{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)
+ψ_star = jax.device_get(ψ_star) # get back numpy form
+```
+
+Using the numba implementation
+
+```{code-cell} ipython3
+num_households = 250_000
+T = 500  # how far to shift forward in time
+wdy = WealthDynamics()
 ψ_0 = np.full(num_households, wdy.y_mean)
 z_0 = wdy.z_mean
 
@@ -648,4 +996,4 @@ plt.show()
 ```
 
 ```{solution-end}
-```
\ No newline at end of file
+```