Revise JAX intro lecture and add autodiff lecture

[jstac]

Summary

• Bug fixes: Fixed coefficient mismatch (0.1 * x**2 vs x**2) in NumPy/JAX function comparison; updated deprecated jax.random.PRNGKey to jax.random.key throughout
• New figures: Added code-generated PRNG key splitting tree diagram and JIT compilation pipeline diagram
• New section: Added vmap section with mean/median example showing why Python loops are inefficient with JAX
• Autodiff preview: Reworked the gradients section as a brief preview with forward reference to the new autodiff lecture
• New lecture: Added "Adventures with Autodiff" (adapted from lecture-jax), covering finite differences vs symbolic vs autodiff, gradient descent with Barzilai-Borwein, and OLS regression
• Housekeeping: Consolidated all imports into initial cell, added (jax_intro)= reference label, updated _toc.yml

Test plan

• CI build passes (executes all notebooks)
• Verify PRNG key splitting tree figure renders correctly
• Verify JIT pipeline figure renders correctly
• Verify autodiff lecture cross-references resolve
• Check vmap example output is clear

🤖 Generated with Claude Code

[jstac]

Detailed changelog

jax_intro.md

Bug fixes:

• Fixed coefficient mismatch in NumPy vs JAX function comparison (0.1 * x**2 vs x**2 — the two versions of f were computing different functions)
• Updated jax.random.PRNGKey → jax.random.key (5 occurrences) — PRNGKey is the legacy API

New content:

• Added code-generated PRNG key splitting tree diagram illustrating how split produces independent keys
• Added JIT compilation pipeline diagram (Python function → trace → XLA → execution)
• Added new vmap section with a mean/median example showing why Python loops are inefficient with JAX, and how transformations compose (jit(vmap(...)))
• Reworked "Gradients" section into "Automatic differentiation: a preview" with forward reference to the new autodiff lecture

Housekeeping:

• Consolidated all imports into the initial cell (was scattered across 4 locations)
• Added (jax_intro)= reference label for cross-referencing
• Rephrased transition text after import consolidation

autodiff.md (new lecture)

• Adapted from lecture-jax/lectures/autodiff.md
• Covers: finite differences vs symbolic calculus vs autodiff, differentiating through control flow and interpolation, gradient descent with Barzilai-Borwein, OLS regression example, polynomial regression exercise
• Updated jax.random.PRNGKey → jax.random.key
• Moved sympy import to top, dropped unused diff import
• Removed nvidia-smi cell (GPU admonition covers this)
• Added pip install cell and intro referencing the jax_intro lecture

numpy_vs_numba_vs_jax.md

Bug fixes:

• Fixed variable name z_max_numpy → z_max_numba for the Numba result
• Fixed print label "JAX vmap v1" → "v2" in the vmap v2 section
• Fixed garbled em dash (--— → ---)
• Changed np.meshgrid → jnp.meshgrid when operating on JAX arrays (avoids silent host-to-device transfer)
• Qualified "significantly faster due to GPU acceleration" → "significantly faster, especially on a GPU"

New content:

• Added "Overall recommendations" synthesis section covering: JAX for vectorized work, Numba for sequential loops, lax.scan differentiability advantage, and a rule of thumb for choosing between them

Housekeeping:

• Consolidated all imports to top (numba, lax, partial were mid-lecture)
• Replaced cm.jet with cm.viridis (perceptually uniform, colorblind-friendly)

_toc.yml

• Added autodiff after numpy_vs_numba_vs_jax in the "High Performance Computing" section

[jstac]

please also review this and merge when ready @mmcky .

please check that the figure showing the jax.jit compilation process in jax_intro came out okay.

Once these changes and those in #512 are merged, please make live