[exchangeable] Update lecture to use JAX

[kp992]

Update the lecture to use JAX and removes numba related code. Also fixes minor typos and code styling issues.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-772--sunny-cactus-210e3e.netlify.app (5098c0e)

📚 Changed Lecture Pages: exchangeable

[jstac]

@longye-tian , would you be willing to review this PR?

@mmcky The deploy link seems broken. Do they go stale? after how long? Should @longye-tian trigger a new build with an empty commit?

[jstac]

@longye-tian or any other reviewer: Does adding JAX improve this lecture? Please also make this value judgement.

[longye-tian]

Hi @kp992 @jstac,

I think the JAX conversion is natural in the appendix simulation section, especially around the replacement of the old nested-loop simulation. This seems like a good fit because the simulation has two parts: recursion over time and repetition across independent paths.

It might be helpful to say this explicitly in the text before the code. Currently the lecture says:

To proceed, we create some Python code.

Perhaps add something like:

The simulation has two dimensions: recursion over time and repetition across independent paths. In JAX, we use lax.scan for the time recursion and vmap to apply the path simulator across many independent key sequences.

For the learning_example block, I’m less sure that the JAX conversion adds much. This part changes the beta-density helper from the old Numba/vectorized version and wraps several scalar helpers with jax.jit, but the surrounding computation is still mostly SciPy root finding, SciPy quadrature, NumPy grids, and Matplotlib plotting. So JAX transformations do not seem central to this part of the lecture. I can see the value of removing Numba and using JAX for consistency, but plain NumPy/SciPy may be clearer here.

A more useful exposition improvement might be to split this page-long helper function into the three conceptual pieces shown in the output figure:

1. likelihood ratio plot
2. density / probability-region plot
3. posterior-dynamics arrows

That would let the text introduce and discuss each panel one by one, rather than asking readers to parse a long plotting helper before seeing the three ideas it creates.

A few small code-level updates could also make the JAX appendix clearer:

• Instead of random_seed = int(a * b + T + N), consider adding an explicit seed or key argument. That would make the source of randomness clearer.
• The comment # Generate all random keys upfront appears inside simulate_path, but the keys are actually generated in simulate; perhaps change it to something like # Use one random key for each date.

A few small remaining typos I noticed while reading:

• “less that one” → “less than one”
• “Bayes’ Law make $\pi$ decrease” → “Bayes’ Law makes $\pi$ decrease”
• “The above graphs shows” → “The above graph shows”
• expected_rario → expected_ratio

Best,
Longye

[jstac]

🤖 Status note for a future session — from a maintainer investigation on 2026-07-08 into why open-PR previews 404. Context only, not instructions.

Netlify preview: https://pr-772--sunny-cactus-210e3e.netlify.app/ currently returns 404.

Why previews are down (repo-wide findings)

1. This branch is stale — 93 commits behind main. A preview build compiles the whole site from this branch. This branch's lectures/house_auction.md still has unpinned !pip install prettytable, which now breaks on a wcwidth incompatibility. main fixed this on 2026-06-28 by pinning prettytable<3.18 (#939). This alone fails any rebuild of this branch until it's updated to main.

2. The arviz failure was a red herring — do NOT pin arviz or rewrite plotting. A 2026-07-07 rebuild also failed in ar1_bayes/ar1_turningpts with an arviz_plots figsize ValueError. That was a transient bug in an intermediate arviz-plots 1.x release, already fixed in arviz 1.2.0. Verified locally on a clean latest-stack venv: the real az.plot_trace(trace) cell (pymc + numpyro InferenceData) runs green. The lectures use only 1.x-compatible arviz APIs (plot_trace, summary, from_numpyro, compare).

Note on recent timeline activity

This PR was close/reopened on 2026-07-07 by a maintainer session purely to trigger a rebuild test — not a content change. That rebuild failed on the stale-branch issue above. Apologies for the notification churn.

Recommended first step for this PR

Update this branch to main (merge or rebase — pulls in #939 plus ~93 other commits), then let CI rebuild. On today's latest libraries the site builds clean, so the preview should return. house_auction is the known blocker; updating also picks up other since-merged fixes — rebuild and address any remaining per-lecture failures. Verify with:

curl -sI https://pr-772--sunny-cactus-210e3e.netlify.app/exchangeable.html

---

This PR touches: exchangeable.md. Last CI build: failure@2026-07-07. Branch: 93 commits behind main as of 2026-07-08.

[kp992]

Thanks @longye-tian. I will review your comment and make the changes.