[career] Convert to JAX using nested vmap (#985)
Rewrites the career choice lecture to use JAX on the GPU, replacing the
CareerWorkerProblem class, operator_factory, and the nested prange loops.
The three options in the Bellman equation are written as a scalar function _B
of one state, returning them in the order they appear in the text. Two
applications of jax.vmap then evaluate it at every state, and T and get_greedy
become the max and the argmax of that same array -- which removes the
duplicated grid-search code the two functions previously carried, and makes
the 1/2/3 action encoding explicit rather than buried in a branch.
Other changes:
- Model primitives move to a NamedTuple with a factory function.
- solve_model uses a bounded jax.lax.while_loop and returns the iteration
count and final error so callers can check convergence.
- Deletes gen_probs. It computed the beta-binomial pmf by hand, and is
bit-for-bit identical to BetaBinomial(n, a, b).pdf() from quantecon, which
the lecture already imports and uses for F_probs and G_probs. This also
drops the scipy.special import.
- Exercise 1 simulates paths with jax.lax.scan.
- Exercise 2 writes one first-passage simulation as a bounded
jax.lax.while_loop and runs 25,000 of them at once under jax.vmap. This is
the part of the lecture that genuinely benefits from a GPU.
- Adds the shared GPU admonition.
Validated against the Numba implementation on main: value functions agree to
1.8e-04 (relative 9.0e-07), and the greedy policy is identical at all 2500
grid cells. The median first passage time is 7 at the default parameters and
14 at β=0.99, matching the values stated in the text. All fourteen code cells
execute on an RTX 4080 and the five figures reproduce the reference images.
Supersedes #617.
Co-authored-by: Humphrey Yang <u6474961@anu.edu.au>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 file changed
Lines changed: 204 additions & 186 deletions
