[imp_sample] Convert code to JAX and check stylesheet compliance

[HumphreyYang]

This is one of the lectures that have the longest runtime.

I made a few big changes that makes the code 4 x faster than Numba optimized version.

Updating in progress...

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-620--sunny-cactus-210e3e.netlify.app (eba931b)

📚 Changed Lecture Pages: imp_sample

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-620--sunny-cactus-210e3e.netlify.app (60826d6)

📚 Changed Lecture Pages: imp_sample

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-620--sunny-cactus-210e3e.netlify.app (e57deef)

📚 Changed Lecture Pages: imp_sample

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-620--sunny-cactus-210e3e.netlify.app (58d1cd2)

📚 Changed Lecture Pages: imp_sample

[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-620--sunny-cactus-210e3e.netlify.app/ currently returns 404.

Why previews are down (repo-wide findings)

1. This branch is stale — 168 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).

Recommended first step for this PR

Update this branch to main (merge or rebase — pulls in #939 plus ~168 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-620--sunny-cactus-210e3e.netlify.app/imp_sample.html

---

This PR touches: imp_sample.md. Last CI build: success@2025-10-17. Branch: 168 commits behind main as of 2026-07-08.

[jstac]

@HumphreyYang thanks for this — the JAX conversion is real work and mostly reads well. I ran the lecture end to end on a GPU box (RTX 4080, JAX 0.10) and it completes in 12 seconds, so the speedup is there. A few things need attention before this can go in, one of them a genuine bug.

What's good

• beta_pdf computed through gammaln in log space is a real improvement over the original gamma(a + b) / (gamma(a) * gamma(b)) — much better behaved in the tails, which matters a lot in a lecture about extreme likelihood ratios.
• fori_loop with the key threaded through the carry, plus vmap over a split key array, is correct key discipline.
• The prose corrections are all right: "peculiar properly" → "property", "despite it's converging" → "its", "replace $E$ by with sample averages".
• This also supersedes the recent removal of @jit(parallel=True) from simulate on main — under JAX the Numba parallel-RNG hazard goes away entirely.

Blocker: the final histogram plots the wrong series

In the last code cell (the $h_3$ block):

results = simulate_multiple_T(...)      # returns (μ_L_p_all, μ_L_q_all), each of shape (n_T, N_simu)
for i, t in enumerate(T_list):
    μ_L_p, μ_L_q = results[i]           # indexes the tuple, then unpacks along the T axis

results[i] selects a method, and unpacking it then splits along the T axis. So each panel compares a method against itself rather than Monte Carlo against importance sampling. Measured values:

	red, labelled "$g$ generating"	blue, labelled "$h_3$ generating"
correct, $T=1$	1.0066 (MC)	1.5436 (IS)
correct, $T=20$	0.2523 (MC)	0.4913 (IS)
as plotted, panel $T=1$	1.0066 (MC @ $T$=1)	0.2523 (MC @ $T$=20)
as plotted, panel $T=20$	1.5436 (IS @ $T$=1)	0.4913 (IS @ $T$=20)

Both legends are therefore wrong, and neither panel shows the comparison the surrounding text describes.

It raises no error only because n_T == 2 == len(results); with three $T$ values it would fail loudly. The $h_2$ block a few cells above gets it right:

μ_L_p_all_h2, μ_L_q_all_h2 = all_results_h2
...
μ_L_p = μ_L_p_all_h2[i]
μ_L_q = μ_L_q_all_h2[i]

The $h_3$ block just needs the same treatment.

Undocumented reduction in sample size

simulate now takes N_samples=1000 and passes it down, where previously it called estimate with its default N=10000. That is a 10x cut in the inner sample, and nothing in the text mentions it. It inflates every variance the lecture reports:

N_samples	MC variance	IS variance
1,000	1.106	0.000522
10,000	6.431	0.000060

I checked that the narrative still holds — at $h_2$, $T=20$ the importance sampling mean is 1.0000 with variance 0.0013, so "the mean is very close to 1 and the variance is small" still reads true. But readers will see numbers an order of magnitude away from the published lecture, and some of the "4x faster" is simply less work being done.

Could you confirm whether the reduction was deliberate? Either restoring N=10000 or noting the change in the text would resolve it — I don't want to guess which you intended.

Style-guide items

1. A figure lost its label. plt.title('real data generating process $g$ and importance distribution $h$') was removed from the a_list / b_list comparison plot without a caption replacing it, so that figure is now unlabelled. main has four plt.title calls; this branch converts one to a caption, deletes this one, and leaves two.

2. One of seven figures has a caption. manual/styleguide/figures.md asks for mystnb metadata on every code-generated figure, and for the matplotlib titles to go. The four histogram figures have neither. The per-panel axs[...].set_title(f'$T$={t}') calls are fine to keep — the manual makes an explicit exception for subplots — but the figures still need captions of their own. Also worth trimming 'Real data generating process $g$ and importance distribution $h$', since the manual suggests roughly 5-6 words.

3. import jax.random as jr and jr.PRNGKey(0). Both are out of step with the manual: the Randomness section of jax.md spells out jax.random.* and asks for jax.random.key rather than the older PRNGKey. The jr alias was in the manual once and was removed deliberately in QuantEcon.manual#76. See also Revert import jax.random as jr in jv and career to spelled-out jax.random #986 — jv and career picked up the same alias and are being reverted. The objection is that np and jnp are conventions a reader arrives already knowing, whereas jr is local to the file, so someone meeting jr.split several hundred lines below the import has to scroll back to decode it.

4. simulate_multiple_T is defined inside a cell that also does all the plotting, which buries the helper. Worth its own cell.

5. key is threaded across roughly eight separate cells via key, subkey = jr.split(key). Re-running any one cell silently changes every result below it, which is awkward in a document readers execute piecemeal.

Merge state

The branch is currently CONFLICTING against main — the only conflict is in simulate, where main dropped @jit(parallel=True). Since this PR rewrites that function entirely, resolving to your side is correct.

Happy to help with any of this if useful, but it's your PR and you know the intent behind the sample-size change better than I do — over to you.

[jstac]

@HumphreyYang Claude authored the above, as a list of suggestions to get this merged. Please let me know if you don't have time and I'll handle it.

[HumphreyYang]

Thanks @jstac, I am working on this now.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-620--sunny-cactus-210e3e.netlify.app

Commit: 4d9967e

📚 Changed Lectures

• imp_sample

---

Build Info

• Workflow: Build Project [using jupyter-book]