Add JAX installation commands to lectures

[mmcky]

Problem

The lecture-jax series does not include JAX installation commands in the lecture notebooks, causing them to fail when run in a standard Anaconda environment or in container-based builds.

What's happening

All 23 lectures that import JAX fail with:

ModuleNotFoundError: No module named 'jax'

Affected lectures:

• aiyagari_jax.md
• arellano.md
• autodiff.md
• hopenhayn.md
• ifp_dl.md
• ifp_egm.md
• inventory_dynamics.md
• inventory_ssd.md
• jax_intro.md
• jax_nn.md
• job_search.md
• kesten_processes.md
• lucas_model.md
• markov_asset.md
• mle.md
• newtons_method.md
• opt_invest.md
• opt_savings_1.md
• opt_savings_2.md
• overborrowing.md
• short_path.md
• status.md
• wealth_dynamics.md

Why this matters

1. QuantEcon standard practice: All QuantEcon lectures should be runnable in a standard Anaconda environment without external setup
2. Container builds fail: The new container-based CI system (https://github.com/QuantEcon/actions) cannot build these lectures because JAX is not declared as a dependency
3. Inconsistent with other lectures: Other QuantEcon lecture series (python-intro, python.myst, python-advanced.myst) include pip install commands for packages not in base Anaconda

Current state

• JAX is not in environment.yml
• JAX is not in _notebook_repo/environment.yml
• Lectures directly import jax without installing it first
• CI workflow manually installs JAX: pip install -U "jax[cuda13]" (only works in this repo's CI, not in other environments)

Solution

Add !pip install jax commands to lectures that use JAX, following QuantEcon standard practice:

Option 1: Install at first use (recommended)
Add to the first code cell of each affected lecture:

!pip install jax

Option 2: Install in jax_intro.md only
Add installation to jax_intro.md and note in other lectures that JAX should already be installed if following the series in order.

References

• Container test failures: https://github.com/QuantEcon/actions/actions/runs/21851356161
• QuantEcon actions repository: https://github.com/QuantEcon/actions
• Example from lecture-python-intro: Uses !pip install for non-Anaconda packages