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

Rewrites the on-the-job search lecture to use JAX on the GPU, replacing the
JVWorker class, operator_factory, and the nested prange/for loops.

The Bellman right-hand side is written as a scalar function _B of one state x
and one action pair (s, ϕ), so it reads like the equation it implements.
Three applications of jax.vmap then vectorize it over ϕ, s and x in turn,
playing the role of a triple loop. T and get_greedy both reduce over the
result, which removes the duplicated grid-search code the two functions
previously carried.

Other changes:

- Model primitives move to a NamedTuple with a factory function, and the two
  action grids are stored alongside the state grid.
- solve_model uses a bounded jax.lax.while_loop and returns the iteration
  count and final error so callers can check convergence.
- The 45 degree diagram in Exercise 1 draws all realizations in one vectorized
  call rather than looping over states and draws, and plots them with a single
  ax.plot rather than 5,000.
- Exercise 2 evaluates w*(ϕ) on the grid rather than via a list comprehension.
- Adds the shared GPU admonition.

Validated against the Numba implementation on main, using identical draws from
f: both converge in 205 iterations, the value functions agree to 7.5e-06
(relative 6.2e-07), and the s and ϕ policies select identical grid points at
every state. Checked float32 against float64 -- the policies are unchanged and
Monte Carlo error at mc_size=100 dominates, so the default precision is kept.
All twelve code cells execute on an RTX 4080, and the three figures reproduce
the behaviour the surrounding prose describes.

Supersedes #618.

Co-authored-by: Humphrey Yang <u6474961@anu.edu.au>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 3 people