[jv] Convert to JAX and content checks

[HumphreyYang]

This PR updates jv by converting the code to JAX.

Standard convertion similar to preceding lectures.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-618--sunny-cactus-210e3e.netlify.app (fc998f6)

📚 Changed Lecture Pages: jv

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-618--sunny-cactus-210e3e.netlify.app (731446f)

📚 Changed Lecture Pages: jv

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-618--sunny-cactus-210e3e.netlify.app (d24a951)

📚 Changed Lecture Pages: jv

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-618--sunny-cactus-210e3e.netlify.app (9ffac19)

📚 Changed Lecture Pages: jv

[jstac]

🤖 Status note for a future session — from a maintainer investigation on 2026-07-08 into why open-PR previews 404. Context only, not instructions.

Netlify preview: https://pr-618--sunny-cactus-210e3e.netlify.app/ currently returns 404.

Why previews are down (repo-wide findings)

1. This branch is stale — 178 commits behind main. A preview build compiles the whole site from this branch. This branch's lectures/house_auction.md still has unpinned !pip install prettytable, which now breaks on a wcwidth incompatibility. main fixed this on 2026-06-28 by pinning prettytable<3.18 (#939). This alone fails any rebuild of this branch until it's updated to main.

2. The arviz failure was a red herring — do NOT pin arviz or rewrite plotting. A 2026-07-07 rebuild also failed in ar1_bayes/ar1_turningpts with an arviz_plots figsize ValueError. That was a transient bug in an intermediate arviz-plots 1.x release, already fixed in arviz 1.2.0. Verified locally on a clean latest-stack venv: the real az.plot_trace(trace) cell (pymc + numpyro InferenceData) runs green. The lectures use only 1.x-compatible arviz APIs (plot_trace, summary, from_numpyro, compare).

Recommended first step for this PR

Update this branch to main (merge or rebase — pulls in #939 plus ~178 other commits), then let CI rebuild. On today's latest libraries the site builds clean, so the preview should return. house_auction is the known blocker; updating also picks up other since-merged fixes — rebuild and address any remaining per-lecture failures. Verify with:

curl -sI https://pr-618--sunny-cactus-210e3e.netlify.app/jv.html

---

This PR touches: jv.md. Last CI build: success@2025-09-18. Branch: 178 commits behind main as of 2026-07-08.

[jstac]

Thanks @HumphreyYang — another solid conversion, and the core of it is genuinely better than what it replaces. Apologies for the delay on this one too.

I checked the numerics against the original Numba implementation (holding f_rvs fixed across both so the comparison is apples-to-apples):

• The Bellman operator is correct and much more readable. Replacing the quadruple-nested prange/for s/for ϕ structure with vmap over a flattened (s, ϕ) grid is a real improvement — the grid search now reads as one idea rather than four levels of indentation. Both versions converge in 205 iterations with max |v_numba − v_jax| = 6.3e-06 (relative 5.2e-07).
• jnp.where(constraint_satisfied, value, -jnp.inf) is better than the original's -1 sentinel. The old code initialised max_val = -1 and scored infeasible pairs at -1, which would have silently returned the sentinel had any feasible value fallen below it. -inf removes that.
• The NamedTuple + factory refactor is the pattern we want, and dropping operator_factory in favour of plain top-level functions is a clear simplification.
• I also double-checked the jnp.meshgrid(search_grid, search_grid) call, since it uses the default indexing='xy' rather than 'ij' — because both axes are the same grid, the set of (s, ϕ) pairs is identical either way, so there's no bug here. Worth a comment though, as it's the kind of thing that bites when the two grids differ.

A few things I'd like to resolve before merge, then some optional points.

Requests

1. The 45-degree diagram loop has become dramatically slower.

for x in plot_grid:
    for i in range(jv.mc_size):
        key, subkey = jr.split(key)
        b = 1 if jr.uniform(subkey) < π(s(x)) else 0
        u = f_rvs[i]
        y = h(x, b, u)
        ax.plot(x, y, 'go', alpha=0.25)

Every one of the 5,000 iterations now does a jr.split, a jr.uniform, two jnp.interp calls and a g(...) as separate device dispatches, plus a host sync to concretise the JAX scalar for the Python if. Timed on this cell's actual dimensions (grid_size=25, mc_size=50, so 100 × 50):

version	time
this PR	5.36 s
original NumPy	0.042 s
vectorized JAX	0.001 s

So it's ~128× slower than the code it replaces. This is the "JAX inside a Python loop" pattern — outside jit there's no compiler to fuse anything, so you pay full dispatch overhead per element.

Since h(x, b, u) = (1-b)·g + b·max(g, u) is just a where, the whole thing vectorizes cleanly:

key = jr.key(0)
b = jr.uniform(key, (plot_grid_size, jv.mc_size)) < π(s(plot_grid))[:, None]
gx = g(jv, plot_grid, ϕ(plot_grid))[:, None]
y = jnp.where(b, jnp.maximum(gx, f_rvs[None, :]), gx)

ax.plot(jnp.repeat(plot_grid, jv.mc_size), y.ravel(), 'go', alpha=0.25)

That also replaces 5,000 individual ax.plot calls with one, which is a large part of the win. (Note max(...) in the current h is the Python builtin on JAX scalars — jnp.maximum is what you want.)

2. The exercise statement still refers to the removed constructor.

The non-executed {code-block} in the Exercise 1 statement (jv_ex1) was left as:

jv = JVWorker(grid_size=25, mc_size=50)

JVWorker is now a NamedTuple requiring all ten fields, so a reader copying this gets a TypeError. It should be create_jv_worker(grid_size=25, mc_size=50), matching the solution below it. Because it's a code-block rather than a code-cell, CI won't catch it.

3. T and get_greedy are now near-identical.

The two functions share about twenty lines verbatim — the search_grid, the objective closure, the meshgrid/stack, and the vmap — differing only in whether they return values[max_idx] or s_phi_pairs[max_idx]. That duplication was somewhat forced by the old operator_factory structure, but there's nothing forcing it now:

@jax.jit
def _search(jv, x, v):
    """Grid search over feasible (s, ϕ); returns best value and best pair."""
    search_grid = jnp.linspace(jv.ɛ, 1, 15)
    s_vals, ϕ_vals = jnp.meshgrid(search_grid, search_grid)
    pairs = jnp.stack([s_vals.ravel(), ϕ_vals.ravel()], axis=1)

    def objective(s_ϕ):
        s, ϕ = s_ϕ
        value = state_action_values(jv, s_ϕ, x, v)
        return jnp.where(s + ϕ <= 1.0, value, -jnp.inf)

    values = jax.vmap(objective)(pairs)
    i = jnp.argmax(values)
    return values[i], pairs[i]

@jax.jit
def T(jv, v):
    return jax.vmap(lambda x: _search(jv, x, v)[0])(jv.x_grid)

@jax.jit
def get_greedy(jv, v):
    policies = jax.vmap(lambda x: _search(jv, x, v)[1])(jv.x_grid)
    return policies[:, 0], policies[:, 1]

Worth doing in a lecture people read closely.

For consideration only

• solve_model no longer signals non-convergence. It's now @jax.jit with a lax.while_loop and discards both iterations and final_error, so hitting max_iter returns silently. The original printed "Failed to converge!". Returning the iteration count and error alongside v_final, and checking outside the jit, would keep that.
• jax.config.update('jax_platform_name', 'cpu') — same question as on [career] Convert to JAX and content checks #617: worth a sentence saying why, since a reader will wonder why a JAX lecture turns off the accelerator.
• !pip install jax at the top — we generally avoid this, since it can pull jax[cpu] rather than the configured build. Probably moot given the CPU pin, but flagging for consistency.
• key=jr.PRNGKey(0) as a default argument. Two things: jax.random.key is the current recommended interface, and a key in a default argument is evaluated once at definition. It's harmless here since keys are immutable, but key=None with key = jr.key(0) if key is None else key is the more conventional shape. Note this also changes behaviour from the original, where every JVWorker() drew fresh f_rvs from global NumPy state — now repeated calls share draws. That's arguably an improvement, but it will shift the figures relative to the live site.
• π is no longer a model parameter. It was π=np.sqrt in the constructor and is now a hardcoded top-level jnp.sqrt. Simpler, and nothing in the lecture varies it — but the text does describe π as part of the model specification, so it may be worth a word either way.
• ax.grid() was dropped from the policy plots, so those figures will differ from the live site. Looks incidental rather than intended.
• Minor: state_action_values takes s_phi as a packed array and immediately unpacks it — passing s and ϕ separately would read better, with the packing confined to the vmap.

Context

We've consolidated our JAX guidance into a single page — QuantEcon/QuantEcon.manual#113. The Python-loop point in request 1 is written up there explicitly, along with device placement, jax.random.key, and bounding while_loops. Comments welcome, particularly if anything in it doesn't match how you'd write these in practice.

Thanks again — the vmap'd grid search is a nice piece of work and the value function checks out exactly.

[jstac]

Thanks again @HumphreyYang — closing this in favour of #984, but I want to be clear that this PR did most of the useful work.

Reviewing it opened up a broader question about when a lecture should use JAX at all, and that discussion ended up reshaping our guidance rather than just this one file. By the time we'd worked through it there were enough changes to make a fresh branch off main simpler than continuing to add requests here — that's purely about mechanics, not about the quality of what you did.

Your work is the starting point for #984 and a lot of it survives intact:

• the NamedTuple + factory function refactor,
• replacing operator_factory with plain top-level functions,
• vectorizing the grid search with vmap, and
• using jnp.where for the feasibility constraint instead of the old -1 sentinel.

You're credited as co-author on the commit.

The main structural change is writing the Bellman right-hand side as a scalar function of one state and one action pair, then applying vmap three times to stand in for the triple loop — following the pattern in opt_savings_2. That has the nice side effect of letting T and get_greedy both reduce over the same array, which removes the duplication I raised in request 3 rather than factoring it into a helper.

Two corrections to things I said above, since they shouldn't stand uncorrected:

The CPU pin. I asked you to justify jax.config.update('jax_platform_name', 'cpu'). I since measured it, and the honest answer is that it should just be removed: on a GPU box the value function iteration is ~40× faster without it, and in career.md the passage-time simulation is ~87× faster. This repo's CI and publish workflows both run on g4dn.2xlarge with jax[cuda13], so the hardware is already there. #984 drops the pin and adds the GPU admonition.

My "~128× slower" figure for the 45 degree diagram. That measurement excluded the 5,000 individual ax.plot calls, which both your version and main make, so the real cell-to-cell difference is smaller than I implied. The underlying point stands — JAX ops dispatched one at a time inside a Python loop are the wrong shape, and it gets worse on a GPU because each element pays a host round trip — but the number was overstated and I should have said what it excluded. #984 sidesteps it by drawing all the realizations in one vectorized call and plotting them with a single ax.plot.

Please do review #984 — and push back if you think the nested-vmap version is harder to read than what you had. Thanks for the work here, it was genuinely a helpful place to start from.

[jstac]

Many thanks for kicking this off @HumphreyYang , much appreciated. I went for a nested vmap, admittedly a taste consideration. In the final version I'm going to emphasize that jax doesn't actually buy us speed gains for such small grids, but it scales much better than numpy.