[career] Convert to JAX and content checks

[HumphreyYang]

This PR updates career by converting the code to JAX.

Some changes are done to vectorize the for loops in addition to some old change rules in the previous lectures.

For the exercises, draws from $F$ and $G$ are generated with inverse transform sampling (jax.random.uniform + jnp.searchsorted) rather than qe.random.draw.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-617--sunny-cactus-210e3e.netlify.app (6232e18)

📚 Changed Lecture Pages: career

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-617--sunny-cactus-210e3e.netlify.app (1ededf4)

📚 Changed Lecture Pages: career

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-617--sunny-cactus-210e3e.netlify.app (5842d1e)

📚 Changed Lecture Pages: career

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-617--sunny-cactus-210e3e.netlify.app (8707aab)

📚 Changed Lecture Pages: career

[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-617--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-617--sunny-cactus-210e3e.netlify.app/career.html

---

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

[jstac]

Thanks @HumphreyYang — this is a nice piece of work, and a good example of a lecture where JAX genuinely earns its place. Apologies it has sat for a while.

I checked the numerics locally against the original Numba implementation, and the core of the conversion is solid:

• The Bellman operator vectorization is correct and a real improvement in clarity. Collapsing the nested prange loops into v1/v2/v3 array expressions makes the three options I/II/III far easier to read against the maths than the indexed version. The jnp.meshgrid(..., indexing='ij') broadcasting checks out: v_star agrees with the original to 6.0e-05, converging in 213 iterations against the original's 212, and the greedy policy is identical on all 2500 grid cells.
• vmap over 25,000 while_loops for the passage times is exactly the right tool, and much nicer than the @jit(parallel=True) + prange version. It reproduces median = 7.0, matching the text.
• The NamedTuple + factory function refactor is precisely the pattern we want, and reads better than the old class.
• I also checked precision, since JAX defaults to float32: it's fine here. Against float64 the solution converges at iteration 212 vs 211 and reaches an exact fixed point, so there's no need for jax_enable_x64. Worth knowing rather than assuming.

A few things I'd like to sort out before merging, then some smaller points that are entirely for your consideration.

Requests

1. quantecon.jr.draw() in the Exercise 1 hint doesn't exist.

-To generate the draws from the distributions $F$ and $G$, use `quantecon.random.draw()`.
+To generate the draws from the distributions $F$ and $G$, use `quantecon.jr.draw()`.

This looks like a random → jr find-and-replace that caught the prose. qe.random.draw exists; qe.jr doesn't. Related: import quantecon as qe was dropped and the solution now uses draw_from_cdf with jnp.searchsorted instead — so the hint should describe that (or point at jax.random), and the PR description's note about leaving qe.random.draw in place is out of date.

2. gen_probs needs another look.

def gen_probs(n, a, b):
    probs = jnp.zeros(n+1)      # unused
    k_vals = jnp.arange(n+1)    # unused
    probs = jnp.array([binom(n, k) * beta(k + a, n - k + b) / beta(a, b) for k in range(n+1)])
    return probs

The first two lines are dead, and the comprehension is still the original Python loop calling SciPy — jnp.array only wraps the result at the end, so nothing actually moved to JAX.

Two ways out, and I'd lean towards the first:

Keep it in NumPy. This is a plotting cell with no performance requirement, and scipy.special.binom/beta are already vectorized, so the whole function is two lines (verified allclose with the current output):

def gen_probs(n, a, b):
    k = np.arange(n + 1)
    return binom(n, k) * beta(k + a, n - k + b) / beta(a, b)

Or do it properly in JAX, if you'd rather the lecture be uniform. Two catches there, both worth knowing:

• jax.scipy.special has beta and betaln, but no binom — you have to build it from gammaln.
• The naive version returns NaN in float32, and this figure plots a = b = 100, so that curve would silently disappear. Working in log space fixes it:

from jax.scipy.special import betaln, gammaln

def gen_probs(n, a, b):
    k = jnp.arange(n + 1)
    log_p = (gammaln(n + 1.) - gammaln(k + 1.) - gammaln(n - k + 1.)
             + betaln(k + a, n - k + b) - betaln(a, b))
    return jnp.exp(log_p)

I checked this against SciPy across (a, b) = (0.5, 0.5), (1, 1) and (100, 100): it agrees to 3.6e-06 in float32, whereas the direct jax.scipy.special.beta route gives NaN at (100, 100).

3. Could you add a line explaining the CPU pin?

jax.config.update('jax_platform_name', 'cpu')   # Set JAX to use CPU

A reader will reasonably wonder why a JAX lecture disables the accelerator. If the reason is that the 50×50 grid is small enough that dispatch overhead dominates, that's a genuinely instructive point and worth a sentence. It's also worth a sanity check that the pin is still needed — the old code was @jit(parallel=True) with prange, so it was already using every core, and with the GPU off the comparison is closer than it looks.

For consideration only

None of these are blockers.

• solve_model now returns silently on non-convergence. The original printed "Failed to converge!" if max_iter was hit. Worth keeping some signal.
• passage_time's lax.while_loop has no iteration bound. It terminates fine for this model, but inside a vmap there's no way to interrupt it if a future parameterization ever made "stay put" unreachable. A t < max_t guard is cheap insurance.
• Tie-breaking changed subtly. The original used strict > and fell through to action 3; jnp.argmax returns the first maximum, so it falls through to action 1. I verified all 2500 cells agree at the default parameterization, so nothing is wrong today — but it's a silent behavioural difference that might be worth a comment.
• ϵ → ε. The maths uses \epsilon, which renders as ϵ, so the code and the text no longer use the same glyph. Minor, but we do try to keep code close to the notation.
• jr.PRNGKey → jax.random.key is the current recommended interface (typed key arrays). We're standardising on it.
• !pip install quantecon jax — we generally avoid installing jax at the top of a lecture, since it can pull jax[cpu] rather than the configured build. Though given the CPU pin above, this one may be moot; happy to be told otherwise.
• gen_path's key=None fallback is unreachable — every call site passes a key.
• Small performance note, purely FYI: a vmapped while_loop runs until every batch element finishes, so with a max passage time of 56 it executes ~6.7× more body steps than the sum of the actual path lengths. Irrelevant at this scale, but useful to know when batching heavy-tailed stopping times.

Context

We've just consolidated our JAX guidance into a single page — QuantEcon/QuantEcon.manual#113 — covering when a lecture should convert, an idiomatic style guide, and a Numba migration section. Several points above (the while_loop bound, jax.random.key, device placement, testing precision rather than assuming it) are drawn from it. Comments very welcome there, especially if anything in it doesn't match how you'd actually write these conversions.

Thanks again — the Bellman rewrite in particular is a clear improvement on what was there.

[jstac]

many thanks @HumphreyYang , above comment written by claude under my guidance. if you can attend to the above i'll try to get this merged quickly

[jstac]

Thanks again @HumphreyYang — closing this in favour of #985, but as with #618 I want to be clear that this PR did the substantive work.

Reviewing it opened up a broader question about when a lecture should use JAX at all, and working through that ended up reshaping our guidance rather than just this file. By the end there were enough changes that a fresh branch off main was simpler than continuing to add requests here. That's mechanics, not a judgement on what you wrote.

Your work is the starting point for #985 and much of it survives:

• the NamedTuple + factory function refactor,
• dropping operator_factory in favour of plain top-level functions,
• vectorizing the Bellman operator, and
• writing the first passage times as a lax.while_loop under vmap — which turns out to be the part of this lecture that most justifies using JAX at all.

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

The main structural change is writing the three options as a scalar function of one state, in the order they appear in the text, then applying vmap twice:

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])

T and get_greedy then reduce over the same array with jnp.max and jnp.argmax, which removes their duplicated grid-search code and makes the 1/2/3 action encoding explicit instead of leaving it to an if/elif chain.

Three corrections to my review above, since they shouldn't stand.

gen_probs — I sent you the wrong way. I asked you to rewrite it in log-space JAX to avoid a float32 NaN at a = b = 100. The function shouldn't exist: it 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. I verified this across (a,b) = (0.5,0.5), (1,1), (2,5) and (100,100). #985 deletes the helper and the scipy.special import with it. Sorry for the detour — I should have looked for the duplication before proposing a fix.

The CPU pin — I asked the wrong question. I asked you to justify jax.config.update('jax_platform_name', 'cpu'). Having measured it, the answer is that it should simply go: on a GPU box the passage-time simulation in Exercise 2 is ~87× faster without it, and in jv.md the value function iteration is ~40× faster. This repo's ci.yml and publish.yml both run on g4dn.2xlarge with jax[cuda13], so the hardware is already there and paid for. I also speculated that "with the GPU off the comparison is closer than it looks" — I had no GPU to hand at the time and shouldn't have said it.

Tie-breaking. I flagged that jnp.argmax returns the first maximum whereas the original's strict > fell through to option 3, and noted it was latent rather than live. That still holds — the policies agree at all 2500 cells — but #985 makes the encoding explicit rather than leaving the difference undocumented.

For what it's worth, #985 validates cleanly against the Numba implementation on main: value functions agree to a relative 9.0e-07, the greedy policy is identical at every one of the 2500 grid cells, and the median first passage times come out at 7 and 14, matching the values quoted in the text.

Please do review #985, and push back if you think the nested-vmap version reads worse than what you had. Thanks again — this was a helpful place to start from.