[inventory_dynamics] Update code to JAX and latest style guide

Convert the lecture from Numba (jitclass / @jit(parallel=True)) to JAX:

- Replace the `Firm` jitclass with a `NamedTuple` and pure update functions.
- Drive the time loop with `jax.lax.fori_loop`, which reads like the Python
  for-loops it replaces, and generate a fresh per-period shock with
  `random.fold_in(key, t)`. Each path therefore takes a single key, so the
  calling code (and every `jax.vmap`) needs only a one-dimensional array of
  keys instead of a 2-D key grid.
- Use the modern typed-key API (`random.key`) rather than `random.PRNGKey`.
- Vectorize the Monte Carlo draw loops with `jax.vmap` (including the
  50,000-draw marginal-distribution cell, which as a Python loop over a
  jitted call ran ~28s on GPU; vmapped it is ~0.3s).
- Use `sample.min()/.max()` in `plot_kde` instead of Python `min()/max()`,
  which iterate a JAX array element-by-element (~4s per call).
- Add a note that `X.at[t].set(...)` is compiled to an in-place update under
  `jax.jit`, so the functional style costs no per-period allocation.

Fix a correctness bug in the restock-frequency exercise: the previous
`simulate_firm_restocks` captured the global `x_init` (= 50 by that point in
the lecture) rather than taking it as an argument, so `compute_freq`'s
`x_init=70` was silently ignored. The exercise asks for firms starting at
X_0 = 70; the code answered for X_0 = 50, reporting ~0.96 instead of the
correct ~0.45. Passing `x_init` explicitly through the vmapped call fixes it.

Also: give the two exercise-solution helpers distinct names so they no longer
shadow each other, and minor typo, terminology, and sentence-case title fixes.

Whole-lecture execution drops from ~57s to ~5s on an RTX 4080.

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 3 people