[jv] Convert to JAX using nested vmap (supersedes #618)

[jstac]

Converts jv.md to JAX. This supersedes #618 by @HumphreyYang, and builds directly on it — the NamedTuple refactor, the vmap'd grid search and the jnp.where handling of the feasibility constraint all come from that PR. Humphrey is credited as co-author on the commit.

Why a separate PR

Reviewing #618 turned into a wider discussion about when a lecture should use JAX at all, so rather than pile more requests onto that branch it seemed cleaner to start from main and apply what came out of it. #618 can be closed in favour of this.

Approach

The Bellman right-hand side is written as a scalar function of one state and one action pair, so it reads like the equation it implements:

def _B(v, jv, x, s, ϕ):
    v_func = lambda z: jnp.interp(z, jv.x_grid, v)
    gxϕ = g(jv, x, ϕ)
    integral = jnp.mean(v_func(jnp.maximum(gxϕ, jv.f_rvs)))
    q = π(s) * integral + (1 - π(s)) * v_func(gxϕ)
    return jnp.where(s + ϕ <= 1, x * (1 - s - ϕ) + jv.β * q, -jnp.inf)

Three applications of jax.vmap then vectorize it over ϕ, s and x in turn, standing in for the triple loop:

# The argument order of _B is    (v,    jv,   x,    s,    ϕ)
_B_ϕ   = jax.vmap(_B,    in_axes=(None, None, None, None, 0))     # over ϕ
_B_sϕ  = jax.vmap(_B_ϕ,  in_axes=(None, None, None, 0,    None))  # then over s
_B_xsϕ = jax.vmap(_B_sϕ, in_axes=(None, None, 0,    None, None))  # then over x

This follows the pattern in opt_savings_2. The advantage over writing one big broadcast expression is that the reader never has to decode [:, None] placement to recover the maths — and T and get_greedy become jnp.max and jnp.argmax over the same array B, which removes the ~20 lines of grid-search code the two functions otherwise duplicate.

Other changes

• Model primitives move to a NamedTuple with a factory function; both action grids are stored alongside the state grid rather than rebuilt inside T.
• solve_model uses a bounded jax.lax.while_loop and returns the iteration count and final error, so non-convergence is visible rather than silent.
• The 45 degree diagram draws all realizations in one vectorized call instead of looping over states and draws, and plots them with a single ax.plot rather than 5,000.
• Exercise 2 evaluates w*(ϕ) on the grid instead of via a list comprehension.
• Adds {include} _admonition/gpu.md.
• No !pip install jax — CI already installs jax[cuda13], and no CPU pin, so the lecture uses the GPU the runner provides.

A note on why JAX here

The lecture now says this explicitly, because it seemed worth being honest about:

The grids here are small, and this model would also run perfectly well in NumPy.

We use JAX because the code is almost as readable as the NumPy equivalent, while scaling far better — to finer grids, or to richer versions of the model with additional state variables, where the same code will make full use of a GPU.

Validation

Checked against the Numba implementation on main, using identical draws from f so the two are comparable:

check	result
iterations to converge	205 both
max abs(v_numba - v_new)	7.5e-06 (relative 6.2e-07)
s and ϕ policies	identical grid point at all 50 states

Also checked float32 against float64: the policy grid indices are unchanged and Monte Carlo error at mc_size=100 dominates the difference, so the default precision is kept and jax_enable_x64 is not set.

All twelve code cells execute on a GPU (RTX 4080), and the three figures reproduce the behaviour the prose describes — s high then dropping to zero, ϕ rising then declining, the 45 degree diagram random below ≈0.2 and deterministic above it converging near 1, and w*(ϕ) peaking at ≈0.6. Worth confirming against the preview build once CI runs.

Related

Our JAX guidance has been consolidated into a single page in QuantEcon/QuantEcon.manual#113, which is where the conventions used here are written down.

🤖 Generated with Claude Code

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-984--sunny-cactus-210e3e.netlify.app

Commit: 2065a56

📚 Changed Lectures

• jv

---

Build Info

• Workflow: Build Project [using jupyter-book]

[jstac]

Many thanks @HumphreyYang , and sorry to supercede your PR. my bad for being slow with a review.

[jstac]

@HumphreyYang @mmcky @kp992 @xuanguang-li

This is a JAX conversion, initially kicked off by @HumphreyYang . We have a lot of open PRs and it will probably help the monorepo move if we have fewer, so I'll go ahead and merge. The style is compliant with the manual.

Note the nested vmap style used to define the Bellman operator, which @kp992 initiated in a different setting a few years ago. I prefer this over heavily vectorized code, since it's closer to the maths. That preference is now noted in the manual in a recently merged PR.

[kp992]

+1 @jstac. Nested vmap helps to directly relate the maths and also directly relates how we would write it with simple for loop in python. So it helps the readers to understand this easily.

[jstac]

Thanks @kp992 , it's now in the manual: https://manual.quantecon.org/styleguide/jax.html#scalar-kernel-plus-nested-vmap

[xuanguang-li]

Hi @jstac,

This form of the value function looks great! I think the single-state argument version is much easier to understand because we don't have to imagine how the vector is stacked.

It also reminds me of the vectorized functions in the basic Python lectures. I remember that some functions can be vectorized automatically. I'm not sure whether the same is true for JAX functions, but if so, the recursive vmap calls might be further simplified for some arguements.