Revise JAX intro lecture and add autodiff lecture (#513)

* Revise JAX intro lecture and add autodiff lecture

- Fix coefficient mismatch in NumPy vs JAX function comparison
- Update jax.random.PRNGKey to jax.random.key throughout
- Add code-generated figures (PRNG key splitting tree, JIT pipeline)
- Add vmap section with examples and transformation composition
- Rework gradients section as autodiff preview with forward reference
- Add autodiff lecture (adapted from lecture-jax) to TOC
- Consolidate all imports into initial cell

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Improve vmap example and consolidate imports

- Replace artificial sum-of-squares vmap example with mean/median statistics
- Add explanation of why Python loops are inefficient with JAX
- Move all imports to initial cell per lecture conventions

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix JIT pipeline figure: make compact, drop italic labels

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Revise NumPy vs Numba vs JAX lecture

- Fix variable name typo (z_max_numpy → z_max_numba)
- Fix vmap v2 print label
- Fix garbled em dash
- Consolidate imports to top of lecture
- Use jnp.meshgrid instead of np.meshgrid for JAX arrays
- Replace cm.jet with cm.viridis
- Qualify JAX speed claim re GPU
- Add overall recommendations section synthesizing trade-offs

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Remove redundant MyST labels that clash with doc targets

Remove (jax_intro)= from jax_intro.md and (autodiff)= from autodiff.md.
These labels duplicate the automatic :std:doc: targets created from the
filenames, causing myst.xref_ambiguous warnings that fail CI with -W.

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Matt McKay <mmcky@users.noreply.github.com>

Author: 3 people

lectures/_toc.yml

@@ -25,6 +25,7 @@ parts:
   - file: numba
   - file: jax_intro
   - file: numpy_vs_numba_vs_jax
+  - file: autodiff
 - caption: Working with Data
   numbered: true
   chapters: