[imp_sample] Convert code to JAX and check stylesheet compliance (#620)

* convert to JAX

* updates

* update pep8

* updates

* [imp_sample] Fix likelihood-ratio figure, a NaN, and style-guide items

The likelihood ratio plot was drawn on a linear axis, where the spike at
the left-hand edge flattens the rest of the curve onto zero -- so the
figure appeared to show the ratio going to zero, contradicting the text
directly beneath it. Switch to a log vertical axis, and reference both
figures with numref.

Fix a NaN: jax.random.beta(key, 0.5, 0.5) returns exactly 1.0 about once
per 2e8 draws, and the T=20 importance sampling arm draws exactly that
many. At the boundary the log-space beta_pdf evaluates 0 * -inf, so one
of the 1000 simulated means came back NaN and was silently dropped by
nanmean. Clip draws off the boundary; the reductions become mean/var.

Compute the Monte Carlo arm once and reuse it, rather than recomputing it
for each importance distribution. It does not depend on the importance
distribution, and recomputing it reported different values for the same
quantity in different figures.

Report a median alongside the mean in the histogram panels. The Monte
Carlo estimator is unbiased in population, so its sample mean at T=20 is
a very noisy statistic (variance > 100; the reported value ranges from
0.37 to 12.99 across seeds) and does not fall monotonically in T. The
median does, and it is what the histograms actually show. Adjust the
surrounding text to describe the collapsing distribution rather than a
growing bias in the mean.

Style-guide items: add the pip install cell and GPU admonition that the
other JAX lectures carry; rename figures fig_x -> fig-x; drop figsize
(14, 10) on the 1x2 grids; use numpy rather than jax.numpy for plotting;
factor the three near-identical histogram blocks into one function; label
the reported dispersion as a variance rather than as sigma-hat.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* [imp_sample] Correct the importance sampling maths and simplify the JAX

Mathematical corrections, all pre-existing:

The Monte Carlo sentence described the wrong estimator. It named the target
as E[L(omega^t)] but then described drawing single omega's and averaging
ell(omega), which estimates the t=1 object rather than the mean of the
product. It now describes drawing sequences, forming the product along each,
and averaging the products -- which is what the code has always done.

The central importance sampling display equated estimators where only
expectations are equal: E-hat^p[.] = E-hat^q[. p/q] is false, since those are
different random variables. Split into a population identity (no hats) and
the estimator it suggests (one hat, on the sample average), matching how the
single-draw case is already handled earlier in the lecture.

The lecture attributed the failure of Monte Carlo to skewness and
undersampling. The sharp statement is that the estimator has infinite
variance: E^g[ell^2] = int f^2/g diverges, because f is uniform while g
vanishes like omega^2 at the origin, so the integrand behaves like
omega^-2. Hence no CLT applies. Added the derivation, and used the same
criterion -- finite variance iff int f^2/h converges -- to make the
lecture's hunches about h1, h2 and h3 precise. For h3 the integral diverges
at both endpoints, so it fails for exactly the reason plain Monte Carlo
does. Correspondingly weakened a claim added in the previous commit, which
said the variance at T=20 "exceeds 100" when it is in fact infinite.

Also: restored a dropped expectation operator; settled the importance
distribution superscript on q, which had appeared as both h and q within
three lines; removed the index collision where t was both the product index
and the path length; replaced undefined E_0 with E; w -> omega in three
displays; and fixed "suppose were", "arbitraily", "at higher values of
distribution", a stray +++ cell break and a missing period.

Code simplification:

Dropped the hand-rolled beta_pdf for jax.scipy.stats.beta.pdf. These agree
to 2e-15 across every parameter pair the lecture uses, and the library
version is the same log-space formula written with xlogy/xlog1py -- the
primitives that define 0*log(0) = 0, and so precisely what the hand-rolled
version lacked at the boundary.

Removed jax.jit from f, g, l and estimate_single_path, none of which are
called from Python. Per the style guide, decorate the outermost function
rather than every helper. Verified that this changes neither the results nor
the runtime, since JAX inlines a nested jit at trace time.

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

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: John Stachurski <john.stachurski@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: 3 people