[career] Convert to JAX using nested vmap (supersedes #617)

[jstac]

Converts career.md to JAX. This supersedes #617 by @HumphreyYang and builds on it — the NamedTuple refactor, the vectorized Bellman operator and the vmap-over-while_loop approach to the first passage times all come from that PR. Humphrey is credited as co-author on the commit.

Same treatment as #984 for jv.md, and same reason for a separate branch: reviewing #617 turned into a wider discussion about when a lecture should use JAX, and starting from main was simpler than stacking further requests onto that branch.

Why JAX belongs here

Unlike some conversions, this one has a genuinely parallel workload today rather than a promise about larger grids: {ref}career_ex2 runs 25,000 independent first-passage simulations, which is exactly what vmap plus a GPU is for.

Approach

The Bellman equation is $Tv(\theta,\epsilon) = \max{I, II, III}$. Those three options become a scalar function of one state, written in the order the text presents them:

def _B(v, cw, i, j):
    stay_put = cw.θ[i] + cw.ϵ[j] + cw.β * v[i, j]                         # I
    new_job  = cw.θ[i] + cw.G_mean + cw.β * v[i, :] @ cw.G_probs          # II
    new_life = cw.G_mean + cw.F_mean + cw.β * cw.F_probs @ v @ cw.G_probs # III
    return jnp.array([stay_put, new_job, new_life])

Two vmaps evaluate it at every state, and then

T          = jnp.max(B(v, cw), axis=-1)
get_greedy = jnp.argmax(B(v, cw), axis=-1) + 1

Because both reduce over the same array, the ~20 lines of grid-search code that T and get_greedy used to duplicate disappear, and the 1/2/3 action encoding becomes explicit rather than buried in an if/elif chain.

gen_probs is deleted

The helper that computed the beta-binomial pmf by hand is bit-for-bit identical to BetaBinomial(n, a, b).pdf() — which the lecture already imports and already uses for F_probs and G_probs. Verified across (a,b) = (0.5,0.5), (1,1), (2,5), (100,100).

ax.plot(range(n + 1), BetaBinomial(n, a, b).pdf(), '-o', label=ab_label)

This also removes the from scipy.special import binom, beta import.

Correction. In my review of #617 I asked for this function to be rewritten in log-space JAX, to avoid a float32 NaN at a = b = 100. That was the wrong call — the function shouldn't exist at all. Apologies for the detour, @HumphreyYang.

Other changes

• 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.
• Exercise 1 simulates paths with jax.lax.scan, which also disposes of the broken quantecon.jr.draw() hint rather than patching it.
• Exercise 2's while_loop gets the iteration bound it was missing.
• Adds {include} _admonition/gpu.md; no !pip install jax and no CPU pin. Both CI and publish run on g4dn.2xlarge with jax[cuda13], and on a GPU box the passage-time simulation is ~87× faster than it is pinned to the CPU.

Validation

Checked against the Numba implementation on main:

check	result
max abs(v_numba - v_new)	1.8e-04 (relative 9.0e-07)
greedy policy	identical at all 2500 grid cells
median first passage time, β=0.95	7 (text says ~7)
median first passage time, β=0.99	14 (text says ~14)

All fourteen code cells execute on a GPU (RTX 4080). I checked the five figures against the reference images: the pmf plot, the value function surface, the policy contour with its new life / new job / stay put regions, the Exercise 1 sample paths (one worker settling, one holding a career while churning jobs), and the Exercise 3 contour with its enlarged stay-put region. Worth confirming against the preview build once CI runs.

Related

• [jv] Convert to JAX using nested vmap (supersedes #618) #984 — the same treatment for jv.md
• QuantEcon/QuantEcon.manual#113 — the JAX guidance, which now documents this scalar-kernel-plus-nested-vmap pattern

🤖 Generated with Claude Code

[github-actions]

📖 Netlify Preview Ready!

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

Commit: 1c09778

📚 Changed Lectures

• career

---

Build Info

• Workflow: Build Project [using jupyter-book]