Unify K operator signature in cake_eating_egm_jax with cake_eating_egm

Update the Coleman-Reffett operator K and solver functions in the JAX
implementation to match the signature from the NumPy version:

- K now takes (c_in, x_in, model) and returns (c_out, x_out)
- solve_model_time_iter now takes (c_init, x_init) and returns (c, x)
- Applied same changes to K_crra and solve_model_crra in exercises

The efficient JAX implementation is fully preserved:
- Vectorization with vmap
- JIT compilation
- Use of jax.lax.while_loop

Tested successfully with maximum deviation of 1.43e-06 from analytical
solution and execution time of ~0.009 seconds.

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

Co-Authored-By: Claude <noreply@anthropic.com>

Author: jstac

lectures/cake_eating_egm_jax.md

@@ -116,7 +116,11 @@ Here's the Coleman-Reffett operator using EGM.
 The key JAX feature here is `vmap`, which vectorizes the computation over the grid points.
 
 ```{code-cell} python3
-def K(σ_array: jnp.ndarray, model: Model) -> jnp.ndarray:
+def K(
+        c_in: jnp.ndarray,  # Consumption values on the endogenous grid
+        x_in: jnp.ndarray,  # Current endogenous grid
+        model: Model        # Model specification
+    ):
     """
     The Coleman-Reffett operator using EGM
 
@@ -126,11 +130,8 @@ def K(σ_array: jnp.ndarray, model: Model) -> jnp.ndarray:
     β, α = model.β, model.α
     grid, shocks = model.grid, model.shocks
 
-    # Determine endogenous grid
-    x = grid + σ_array  # x_i = k_i + c_i
-
     # Linear interpolation of policy using endogenous grid
-    σ = lambda x_val: jnp.interp(x_val, x, σ_array)
+    σ = lambda x_val: jnp.interp(x_val, x_in, c_in)
 
     # Define function to compute consumption at a single grid point
     def compute_c(k):
@@ -139,9 +140,12 @@ def K(σ_array: jnp.ndarray, model: Model) -> jnp.ndarray:
 
     # Vectorize over grid using vmap
     compute_c_vectorized = jax.vmap(compute_c)
-    c = compute_c_vectorized(grid)
+    c_out = compute_c_vectorized(grid)
+
+    # Determine corresponding endogenous grid
+    x_out = grid + c_out  # x_i = k_i + c_i
 
-    return c
+    return c_out, x_out
 ```
 
 We define utility and production functions globally.
@@ -171,47 +175,47 @@ The solver uses JAX's `jax.lax.while_loop` for the iteration and is JIT-compiled
 ```{code-cell} python3
 @jax.jit
 def solve_model_time_iter(model: Model,
-                          σ_init: jnp.ndarray,
+                          c_init: jnp.ndarray,
+                          x_init: jnp.ndarray,
                           tol: float = 1e-5,
-                          max_iter: int = 1000) -> jnp.ndarray:
+                          max_iter: int = 1000):
     """
     Solve the model using time iteration with EGM.
     """
 
     def condition(loop_state):
-        i, σ, error = loop_state
+        i, c, x, error = loop_state
         return (error > tol) & (i < max_iter)
 
     def body(loop_state):
-        i, σ, error = loop_state
-        σ_new = K(σ, model)
-        error = jnp.max(jnp.abs(σ_new - σ))
-        return i + 1, σ_new, error
+        i, c, x, error = loop_state
+        c_new, x_new = K(c, x, model)
+        error = jnp.max(jnp.abs(c_new - c))
+        return i + 1, c_new, x_new, error
 
     # Initialize loop state
-    initial_state = (0, σ_init, tol + 1)
+    initial_state = (0, c_init, x_init, tol + 1)
 
     # Run the loop
-    i, σ, error = jax.lax.while_loop(condition, body, initial_state)
+    i, c, x, error = jax.lax.while_loop(condition, body, initial_state)
 
-    return σ
+    return c, x
 ```
 
 We solve the model starting from an initial guess.
 
 ```{code-cell} python3
-σ_init = jnp.copy(grid)
-σ = solve_model_time_iter(model, σ_init)
+c_init = jnp.copy(grid)
+x_init = grid + c_init
+c, x = solve_model_time_iter(model, c_init, x_init)
 ```
 
 Let's plot the resulting policy against the analytical solution.
 
 ```{code-cell} python3
-x = grid + σ  # x_i = k_i + c_i
-
 fig, ax = plt.subplots()
 
-ax.plot(x, σ, lw=2,
+ax.plot(x, c, lw=2,
         alpha=0.8, label='approximate policy function')
 
 ax.plot(x, σ_star(x, model.α, model.β), 'k--',
@@ -224,15 +228,16 @@ plt.show()
 The fit is very good.
 
 ```{code-cell} python3
-max_dev = jnp.max(jnp.abs(σ - σ_star(x, model.α, model.β)))
+max_dev = jnp.max(jnp.abs(c - σ_star(x, model.α, model.β)))
 print(f"Maximum absolute deviation: {max_dev:.7}")
 ```
 
 The JAX implementation is very fast thanks to JIT compilation and vectorization.
 
 ```{code-cell} python3
 with qe.Timer(precision=8):
-    σ = solve_model_time_iter(model, σ_init).block_until_ready()
+    c, x = solve_model_time_iter(model, c_init, x_init)
+    jax.block_until_ready(c)
 ```
 
 This speed comes from:
@@ -282,19 +287,21 @@ def u_prime_inv_crra(x, γ):
 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:
+def K_crra(
+        c_in: jnp.ndarray,  # Consumption values on the endogenous grid
+        x_in: jnp.ndarray,  # Current endogenous grid
+        model: Model,       # Model specification
+        γ: float            # CRRA parameter
+    ):
     """
     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)
+    σ = lambda x_val: jnp.interp(x_val, x_in, c_in)
 
     # Define function to compute consumption at a single grid point
     def compute_c(k):
@@ -303,55 +310,63 @@ def K_crra(σ_array: jnp.ndarray, model: Model, γ: float) -> jnp.ndarray:
 
     # Vectorize over grid using vmap
     compute_c_vectorized = jax.vmap(compute_c)
-    c = compute_c_vectorized(grid)
+    c_out = compute_c_vectorized(grid)
+
+    # Determine corresponding endogenous grid
+    x_out = grid + c_out
 
-    return c
+    return c_out, x_out
 ```
 
 We also need a solver that uses this operator.
 
 ```{code-cell} python3
 @jax.jit
 def solve_model_crra(model: Model,
-                     σ_init: jnp.ndarray,
+                     c_init: jnp.ndarray,
+                     x_init: jnp.ndarray,
                      γ: float,
                      tol: float = 1e-5,
-                     max_iter: int = 1000) -> jnp.ndarray:
+                     max_iter: int = 1000):
     """
     Solve the model using time iteration with EGM and CRRA utility.
     """
 
     def condition(loop_state):
-        i, σ, error = loop_state
+        i, c, x, 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
+        i, c, x, error = loop_state
+        c_new, x_new = K_crra(c, x, model, γ)
+        error = jnp.max(jnp.abs(c_new - c))
+        return i + 1, c_new, x_new, error
 
     # Initialize loop state
-    initial_state = (0, σ_init, tol + 1)
+    initial_state = (0, c_init, x_init, tol + 1)
 
     # Run the loop
-    i, σ, error = jax.lax.while_loop(condition, body, initial_state)
+    i, c, x, error = jax.lax.while_loop(condition, body, initial_state)
 
-    return σ
+    return c, x
 ```
 
 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 = {}
+endogenous_grids = {}
 
 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
+    c_init = jnp.copy(model_crra.grid)
+    x_init = model_crra.grid + c_init
+    c_gamma, x_gamma = solve_model_crra(model_crra, c_init, x_init, γ)
+    jax.block_until_ready(c_gamma)
+    policies[γ] = c_gamma
+    endogenous_grids[γ] = x_gamma
     print(f"Solved for γ = {γ}")
 ```
 
@@ -361,7 +376,7 @@ Plot the policies on their endogenous grids.
 fig, ax = plt.subplots()
 
 for γ in γ_values:
-    x = model_crra.grid + policies[γ]
+    x = endogenous_grids[γ]
     if γ == 1.0:
         ax.plot(x, policies[γ], 'k-', linewidth=2,
                 label=f'γ = {γ:.2f} (log utility)', alpha=0.8)