Skip to content

Fix parameter order in solve_model function calls in ifp_advanced #762

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mmcky wants to merge 2 commits into from
+8 −8
Closed

Fix parameter order in solve_model function calls in ifp_advanced #762

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ca4439b
Fix parameter order in solve_model function calls in ifp_advanced
mmcky Dec 9, 2025
68ff93d
Merge branch 'main' into fix-ifp-advanced-parameter-order
mmcky Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions lectures/ifp_advanced.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,15 +487,15 @@ a_init = σ_init.copy()
Let's generate an approximation solution with JAX:

```{code-cell} ipython3
a_star, σ_star = solve_model(ifp, a_init, σ_init)
σ_star, a_star = solve_model(ifp, σ_init, a_init)

@mmcky mmcky Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed — this is correct, and the PR is being closed without merging.

Tracing the labels through: K(a_in, c_in, ifp) takes the asset grid first (it is the x-coords of jnp.interp) and returns (a_out, c_out). solve_model then calls it as c_out, a_out = K(c_in, a_in, ifp) — swapped on both the arguments and the unpacking. That double swap means solve_model's first parameter is what actually reaches K's a_in slot, and its first return value is K's a_out. So the real contract of the function as written is solve_model(ifp, a_init, c_init) -> (a_grid, c_policy) — exactly the ordering on main, and the opposite of what the signature and docstring claim.

Two notes on the impact, for the record. The argument reordering in this PR is a harmless no-op at every call site, because the two arrays passed are literally identical there (a_init = σ_init.copy() at :484, c_init = a_init at :644/:736/:813). The return unpacking is the real regression: c_vec ends up holding the asset grid and a_vec the consumption policy, which flow into compute_asset_stationary(c_vec, a_vec, …) → simulate_household → jnp.interp(a, a_vec[:, z], c_vec[:, z]) at :534 with x and y reversed. It does not raise; it silently produces wrong asset distributions and Gini coefficients.

The genuine defect you surfaced is the misleading names inside solve_model itself, which is a readability problem in a lecture whose source students are meant to read. That will be fixed in a separate PR against latest main, along the lines of your second suggestion — un-swapping the body to a_out, c_out = K(a_in, c_in, ifp) so the parameter names and return order are truthful, with no behavioural change.

```

Let's try it again with a timer.

```{code-cell} python3
with qe.Timer(precision=8):
a_star, σ_star = solve_model(ifp, a_init, σ_init)
a_star.block_until_ready()
σ_star, a_star = solve_model(ifp, σ_init, a_init)
σ_star.block_until_ready()
```

## Simulation
Expand Down Expand Up @@ -642,7 +642,7 @@ s_grid = ifp.s_grid
n_z = len(ifp.P)
a_init = s_grid[:, None] * jnp.ones(n_z)
c_init = a_init
a_vec, c_vec = solve_model(ifp, a_init, c_init)
c_vec, a_vec = solve_model(ifp, c_init, a_init)
Comment thread
mmcky marked this conversation as resolved.
assets = compute_asset_stationary(c_vec, a_vec, ifp, num_households=200_000)

# Compute Gini coefficient for the plot
Expand Down Expand Up @@ -734,8 +734,8 @@ for a_r in a_r_vals:
n_z_temp = len(ifp_temp.P)
a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)
c_init_temp = a_init_temp
a_vec_temp, c_vec_temp = solve_model(
ifp_temp, a_init_temp, c_init_temp
c_vec_temp, a_vec_temp = solve_model(
ifp_temp, c_init_temp, a_init_temp
)
Comment thread
mmcky marked this conversation as resolved.

# Simulate households
Expand Down Expand Up @@ -811,8 +811,8 @@ for a_y in a_y_vals:
n_z_temp = len(ifp_temp.P)
a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)
c_init_temp = a_init_temp
a_vec_temp, c_vec_temp = solve_model(
ifp_temp, a_init_temp, c_init_temp
c_vec_temp, a_vec_temp = solve_model(
ifp_temp, c_init_temp, a_init_temp
)

# Simulate households
Expand Down
Loading