[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: c3491d3

📚 Changed Lectures

• imp_sample

---

Build Info

• Workflow: Build Project [using jupyter-book]

[jstac]

Thanks for working on this @HumphreyYang ! Ready for review?

[HumphreyYang]

Thanks for working on this @HumphreyYang ! Ready for review?

Yes! Please let me know how to improve it further.

[jstac]

@HumphreyYang thanks — everything from the last round is in, and the lecture reads well. Rather than send you another list I've pushed a commit to this branch (2263a7c). Please review it as you would any other change, and revert anything you disagree with.

Runs clean end to end in 34s on an RTX 4080, all seven figures render, and I re-checked every numeric claim in the prose against the new output.

What's in the commit, briefly

• The likelihood ratio figure contradicted the text beneath it. ℓ is drawn on a linear axis starting at ω = 10⁻², so it hits ~2372 at the left edge and is then visually pinned to zero across the remaining 98% of the panel. A reader sees the ratio going to zero, immediately above a sentence saying it "approaches infinity". Switched to a log vertical axis; also added numref references and a clause about the (much slower) divergence as ω → 1, which the log scale makes visible at the right edge.
• A NaN. jax.random.beta(key, 0.5, 0.5) returns exactly 1.0 about once per 2×10⁸ draws, and the h₁, T=20 importance sampling arm draws exactly N_simu × N_samples × T = 1000 × 10000 × 20 = 2×10⁸ of them. At the boundary the log-space beta_pdf evaluates (b-1)·log(1-w) = 0 · -inf = nan for f = Beta(1,1), where the true density is 1. One of the 1000 simulated means came back NaN and was silently dropped by nanmean. Draws are now clipped off the boundary; the reductions are plain mean/var. A full audit over every simulated array is now 0 NaN and 0 inf.
• The Monte Carlo arm is computed once and reused. It doesn't depend on the importance distribution, so recomputing it for h₁/h₂/h₃ was both ~40% of the runtime and a source of inconsistency — the same quantity was reported as 0.408 in one figure and 0.801 in another.
• Style-guide items. Added the !pip install jax cell and the GPU admonition that the other 29 JAX lectures carry; renamed the figures fig_x → fig-x (the repo is 200 hyphens to 5 underscores); dropped figsize=(14, 10) on the 1×2 grids; numpy rather than jax.numpy in plotting code; and folded the three near-identical 25-line histogram blocks into one plot_estimates function. Net −57 lines.

The one that needs your judgement: mean vs median

This is the only change that touches the lecture's argument rather than its code, so I want to lay out the reasoning in full.

How it surfaced. Computing the Monte Carlo arm once reshuffled the RNG keys. With the new keys the reported MC means came out

T = 1     5      10     20
    1.00  0.936  0.793  0.884

which is not monotone, and sits directly beneath "Evidently, the bias increases with increases in $T$."

Why this isn't just an unlucky seed to be swapped out. The Monte Carlo estimator here is unbiased in population — E[L(ω^T)] = 1 exactly, for every T. That is the whole premise of the lecture. So there is no bias in the mean that could grow with T; asymptotically in N_simu the reported μ̂ converges to 1 for every T, and the apparent downward drift would disappear entirely.

What actually degrades as T grows is the shape of the sampling distribution: almost all the mass collapses toward zero while the expectation is sustained by rare, very large outliers. That is exactly what the red histograms show, and it is the real reason importance sampling is needed.

The consequence for the reported number. Because the mean rides on those rare outliers, μ̂ is itself a very noisy statistic at large T — the variance across the 1000 estimates is >100 at T=20. Measuring the reported mean across 15 different seeds (same N_simu=1000, N_samples=10000):

T	min	median	max	median of the 1000 estimates
1	0.9930	0.9978	1.0517	0.981
5	0.9453	0.9609	1.0604	0.841
10	0.7701	0.8602	1.1105	0.586
20	0.3705	0.4697	12.9902	0.190

So at T=20 the number printed on the figure can be anywhere from 0.37 to 13 depending on the seed. The sequence in the current version of the PR — 1.0, 0.989, 0.899, 0.408 — is monotone, but that is a property of that particular seed rather than of the estimator. The median, by contrast, falls steadily and is stable across seeds.

What I did. Each panel now reports a median alongside the mean and variance, and the two sentences after the figure now describe the collapsing distribution rather than a growing bias in the mean:

The simulation exercises above show that the importance sampling estimates stay centered on 1 for every T, while the distribution of the standard Monte Carlo estimates drifts steadily to the left as T increases. […] The Monte Carlo estimator is unbiased in population — its expectation is exactly 1 for every T. What deteriorates as T grows is the shape of its sampling distribution […] The median falls steadily in T, whereas the reported mean μ̂ is itself an unreliable statistic once T is large.

What I deliberately didn't do. The quick fix would have been to hunt for a seed that makes the means look monotone again. That would be cherry-picking a seed to support a claim the statistic doesn't robustly support, so I left it alone.

I recognise this is a change to your exposition and not merely to the code, and you may well prefer a different framing — the underlying point is yours and you may want to make it differently. Reverting to the previous wording is a small edit if you'd rather; the figure change and the prose change are independent of everything else in the commit.

---

Claude wrote this comment and the commit it describes; I've reviewed both. — jstac

[jstac]

@HumphreyYang above is from Claude, i'll review now, you can ignore.

[jstac]

A second commit on this branch (c3491d3), from a closer read of the mathematics. All of the errors below are pre-existing — none were introduced by this PR — but since we're already in here it seemed worth fixing them properly.

Two things the lecture was saying that it didn't mean

The Monte Carlo sentence described the wrong estimator. It named the target as $E[L(\omega^t)]$ and then described the procedure as "repeatedly draw $\omega$ from $g$, calculate the likelihood ratio $\ell(\omega)$ for each draw, then average these over all draws." Averaging individual $\ell(\omega)$ values estimates $E[\ell(\omega)]$ — the $t=1$ case — not the mean of the product. The product never appeared in the procedure described.

The code has always done the right thing: estimate_single_path accumulates L = L * f(w)/g(w) over a fori_loop of T steps, so each call returns one draw of $L(\omega^T)$, and estimate then vmaps that over N independent keys and averages. The prose was describing the inner loop as though it were the outer one. It now describes drawing sequences, forming the product along each, and averaging the products.

The central importance sampling display equated estimators. It asserted

$$\hat{E}^p\left[L(\omega^t)\right] = \hat{E}^p\left[\prod \ell\right] = \hat{E}^q\left[\prod \ell \tfrac{p}{q}\right] = \tfrac{1}{N}\sum\cdots$$

The middle equality is false as written — $\hat{E}^p$ and $\hat{E}^q$ are different random variables. What holds is the population identity $E^p[\cdot] = E^q[\cdot, p/q]$; only the final step, defining the sample average, earns a hat. The lecture already does this correctly for the single-draw case a few paragraphs earlier, so it was contradicting itself. Now split into a population identity followed by the estimator it suggests.

The variance is infinite, not merely large

The lecture explains the failure of Monte Carlo as skewness plus undersampling. That's right in spirit but understates it. The second moment of a single likelihood ratio is

$$E^g\left[\ell(\omega)^2\right] = \int_0^1 \frac{f(\omega)^2}{g(\omega)},d\omega$$

and $f$ is uniform while $g$ vanishes like $\omega^{G_a-1} = \omega^2$ at the origin, so the integrand behaves like $\omega^{-2}$ and the integral diverges. Numerically, $\int_\epsilon^1 f^2/g$ grows exactly like $1/\epsilon$ — 23.7, 237, 2368, 23670 for $\epsilon = 10^{-2} \ldots 10^{-5}$.

So $\ell(\omega)$ has no finite variance, and neither does $L(\omega^t)$ for any $t$. No CLT applies to the sample averages, which is the precise reason they are so erratic — and it explains the seed sensitivity I reported in the previous comment: the $\hat\sigma^2$ of 150 at $T=20$ isn't a large number, it's a sample statistic estimating infinity. I've correspondingly weakened my own sentence from the earlier commit, which said the variance "exceeds 100."

The same criterion sharpens the $h_1$/$h_2$/$h_3$ discussion. Reweighting gives each observation the value $\ell(\omega)g(\omega)/h(\omega) = f(\omega)/h(\omega)$, so the estimator has finite variance exactly when $\int_0^1 f^2/h < \infty$:

	$\int f^2/h$	finite variance?
$g$ (plain Monte Carlo)	diverges at 0	no
$h_1 = Beta(0.5, 0.5)$	1.234	yes
$h_2 = Beta(1, 1.2)$	1.042	yes
$h_3 = Beta(2, 5)$	diverges at both ends	no
$h = f = Beta(1,1)$	1 exactly, variance 0	yes (degenerate)

The lecture's hunch that $h_3$ is poor is exactly right, and its stated reason — almost no mass near 0 and near 1 — is precisely the condition that makes the integral diverge. It just never got sharpened into the variance statement. It now does, along with the observation that $h_3$ fails for the same reason plain Monte Carlo does.

(One small correction there: "$h_3$ has zero mass at values very close to 0" became "almost no mass". $h_3$ is strictly positive on $(0,1)$, which is what makes it a valid importance distribution at all — genuine zero mass would make the estimator biased rather than merely infinite-variance, so the distinction carries the argument.)

Smaller notational fixes

A dropped expectation operator on the right-hand side of one display; the importance distribution superscript appearing as both $h$ and $q$ within three lines, now settled on $q$ with $\omega_{n,i}^q$ defined once; the index collision where $t$ was both the product index and the path length; undefined $E_0$ replaced by $E$ (three places); $w \to \omega$ in three displays; plus "suppose were", "arbitraily", "at higher values of distribution", a stray +++ cell break, and a missing period.

Code

• beta_pdf → jax.scipy.stats.beta.pdf. These agree to $2\times10^{-15}$ across every parameter pair the lecture uses. Worth noting why the library version doesn't have the boundary bug fixed in the previous commit: it's the same log-space formula written with xlogy/xlog1py, the primitives that define $0\cdot\log 0 = 0$. The hand-rolled version was that formula with plain * and jnp.log, which is exactly where the 0 * -inf came from.
• Removed jax.jit from f, g, l and estimate_single_path, none of which are called from Python. The style guide asks for the outermost function to be decorated rather than every helper. Two decorators remain, on estimate and simulate, both genuine Python entry points. Verified this changes neither results nor runtime — JAX inlines a nested jit at trace time, so XLA sees the same graph either way.

Verified end to end on a GPU: 34s, 0 NaN and 0 inf across every simulated array, results unchanged.

One deliberate non-change: the code still uses w for $\omega$ while using Unicode μ elsewhere in the same cells. Renaming would be 35 occurrences and we decided to leave it, but flagging it in case you'd rather it were consistent.

---

Claude wrote this comment and the commit it describes; I've reviewed both. — jstac

[jstac]

@thomassargent30 @HumphreyYang Misc. improvements to the JAX code. Tighted and improved the discussion of importance sampling under alternative h_i. Merging.

[mmcky]

✅ Translation sync completed (zh-cn)

Target repo: QuantEcon/lecture-python.zh-cn
Translation PR: QuantEcon/lecture-python.zh-cn#232
Files synced (1):

• lectures/imp_sample.md