Refine JAX EGM lecture: improve imports and add CRRA exercise

- Moved all imports to top of file (NamedTuple with other imports)
- Removed unused Callable import
- Added block_until_ready() to timing for accurate JAX benchmarking
- Improved error output formatting with print statement
- Added comprehensive CRRA utility exercise demonstrating convergence

Exercise improvements:
- Uses correct CRRA form: u(c) = (c^(1-γ) - 1)/(1-γ) that converges to log
- Focuses on γ values approaching 1 from above (1.05, 1.1, 1.2)
- Plots γ=1 (log case) in black with clear labeling
- Includes explanation of endogenous grid coverage differences
- Shows numerical convergence with maximum deviation metrics

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: jstac

lectures/cake_eating_egm_jax.md

@@ -37,6 +37,8 @@ We'll also use JAX's `vmap` function to fully vectorize the Coleman-Reffett oper
 Let's start with some standard imports:
 
 ```{code-cell} ipython
+from typing import NamedTuple
+
 import matplotlib.pyplot as plt
 import jax
 import jax.numpy as jnp
@@ -80,8 +82,6 @@ The `Model` class stores only the data (grids, shocks, and parameters).
 Utility and production functions will be defined globally to work with JAX's JIT compiler.
 
 ```{code-cell} python3
-from typing import NamedTuple, Callable
-
 class Model(NamedTuple):
     β: float              # discount factor
     μ: float              # shock location parameter
@@ -225,14 +225,14 @@ plt.show()
 The fit is excellent.
 
 ```{code-cell} python3
-jnp.max(jnp.abs(σ - σ_star(x, model.α, model.β)))
+print(f"Maximum absolute deviation: {jnp.max(jnp.abs(σ - σ_star(x, model.α, model.β))):.6e}")
 ```
 
 The JAX implementation is very fast thanks to JIT compilation and vectorization.
 
 ```{code-cell} python3
 with qe.Timer():
-    σ = solve_model_time_iter(model, σ_init)
+    σ = solve_model_time_iter(model, σ_init).block_until_ready()
 ```
 
 This speed comes from:
@@ -241,3 +241,155 @@ This speed comes from:
 * Vectorization via `vmap` in the Coleman-Reffett operator
 * Use of `jax.lax.while_loop` instead of a Python loop
 * Efficient JAX array operations throughout
+
+## Exercises
+
+```{exercise}
+:label: cake_egm_jax_ex1
+
+Solve the stochastic cake eating problem with CRRA utility
+
+$$
+u(c) = \frac{c^{1 - \gamma} - 1}{1 - \gamma}
+$$
+
+Compare the optimal policies for values of $\gamma$ approaching 1 from above (e.g., 1.05, 1.1, 1.2).
+
+Show that as $\gamma \to 1$, the optimal policy converges to the policy obtained with log utility ($\gamma = 1$).
+
+Hint: Use values of $\gamma$ close to 1 to ensure the endogenous grids have similar coverage and make visual comparison easier.
+```
+
+```{solution-start} cake_egm_jax_ex1
+:class: dropdown
+```
+
+We need to create a version of the Coleman-Reffett operator and solver that work with CRRA utility.
+
+The key is to parameterize the utility functions by $\gamma$.
+
+```{code-cell} python3
+def u_crra(c, γ):
+    return (c**(1 - γ) - 1) / (1 - γ)
+
+def u_prime_crra(c, γ):
+    return c**(-γ)
+
+def u_prime_inv_crra(x, γ):
+    return x**(-1/γ)
+```
+
+Now we create a version of the Coleman-Reffett operator that takes $\gamma$ as a parameter.
+
+```{code-cell} python3
+def K_crra(σ_array: jnp.ndarray, model: Model, γ: float) -> jnp.ndarray:
+    """
+    The Coleman-Reffett operator using EGM with CRRA utility
+    """
+    # Simplify names
+    β, α = model.β, model.α
+    grid, shocks = model.grid, model.shocks
+
+    # Determine endogenous grid
+    x = grid + σ_array
+
+    # Linear interpolation of policy using endogenous grid
+    σ = lambda x_val: jnp.interp(x_val, x, σ_array)
+
+    # Define function to compute consumption at a single grid point
+    def compute_c(k):
+        vals = u_prime_crra(σ(f(k, α) * shocks), γ) * f_prime(k, α) * shocks
+        return u_prime_inv_crra(β * jnp.mean(vals), γ)
+
+    # Vectorize over grid using vmap
+    compute_c_vectorized = jax.vmap(compute_c)
+    c = compute_c_vectorized(grid)
+
+    return c
+```
+
+We also need a solver that uses this operator.
+
+```{code-cell} python3
+@jax.jit
+def solve_model_crra(model: Model,
+                     σ_init: jnp.ndarray,
+                     γ: float,
+                     tol: float = 1e-5,
+                     max_iter: int = 1000) -> jnp.ndarray:
+    """
+    Solve the model using time iteration with EGM and CRRA utility.
+    """
+
+    def condition(loop_state):
+        i, σ, error = loop_state
+        return (error > tol) & (i < max_iter)
+
+    def body(loop_state):
+        i, σ, error = loop_state
+        σ_new = K_crra(σ, model, γ)
+        error = jnp.max(jnp.abs(σ_new - σ))
+        return i + 1, σ_new, error
+
+    # Initialize loop state
+    initial_state = (0, σ_init, tol + 1)
+
+    # Run the loop
+    i, σ, error = jax.lax.while_loop(condition, body, initial_state)
+
+    return σ
+```
+
+Now we solve for $\gamma = 1$ (log utility) and values approaching 1 from above.
+
+```{code-cell} python3
+γ_values = [1.0, 1.05, 1.1, 1.2]
+policies = {}
+
+model_crra = create_model(α=α)
+
+for γ in γ_values:
+    σ_init = jnp.copy(model_crra.grid)
+    σ_gamma = solve_model_crra(model_crra, σ_init, γ).block_until_ready()
+    policies[γ] = σ_gamma
+    print(f"Solved for γ = {γ}")
+```
+
+Plot the policies on their endogenous grids.
+
+```{code-cell} python3
+fig, ax = plt.subplots()
+
+for γ in γ_values:
+    x = model_crra.grid + policies[γ]
+    if γ == 1.0:
+        ax.plot(x, policies[γ], 'k-', linewidth=2,
+                label=f'γ = {γ:.2f} (log utility)', alpha=0.8)
+    else:
+        ax.plot(x, policies[γ], label=f'γ = {γ:.2f}', alpha=0.8)
+
+ax.set_xlabel('State x')
+ax.set_ylabel('Consumption σ(x)')
+ax.legend()
+ax.set_title('Optimal policies: CRRA utility approaching log case')
+plt.show()
+```
+
+Since the endogenous grids are similar for $\gamma$ values close to 1, the policies overlap nicely.
+
+Note that the plots for $\gamma > 1$ do not cover the entire x-axis range shown.
+
+This is because the endogenous grid $x = k + \sigma(k)$ depends on the consumption policy, which varies with $\gamma$.
+
+Let's check the maximum deviation between the log utility case ($\gamma = 1.0$) and values approaching from above.
+
+```{code-cell} python3
+for γ in [1.05, 1.1, 1.2]:
+    max_diff = jnp.max(jnp.abs(policies[1.0] - policies[γ]))
+    print(f"Max difference between γ=1.0 and γ={γ}: {max_diff:.6e}")
+```
+
+As expected, the differences decrease as $\gamma$ approaches 1 from above, confirming convergence.
+
+```{solution-end}
+```