Rewrite the Non-Conjugate Priors lecture (#913)
* Rewrite the Non-Conjugate Priors lecture
Complete rewrite teaching the same material more cleanly:
- Remove the two NamedTuple classes (BayesianInference,
BayesianInferencePlot) and the string-dispatch indirection in favor
of small functions: one reusable binomial_model(prior, k, n), a
run_nuts helper, and a plot_prior_posterior helper.
- Reorganize into: (1) MCMC reproduces the conjugate beta posterior
(validated against the analytical result from prob_meaning, with
arviz R-hat/trace diagnostics); (2) non-conjugate priors worked
through one at a time (uniform incl. a boundary-exclusion example,
truncated log-normal, truncated Laplace); (3) variational inference
via an AutoNormal autoguide, compared against the NUTS posterior.
- Continuity with prob_meaning: same coin-flip DGP with true θ=0.4,
and a deliberately small sample so the prior visibly matters.
- Switch plotting to matplotlib + arviz (drop seaborn); keep the full
ELBO derivation.
Verified end-to-end via jupytext export + headless execution.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Run MCMC chains vectorized instead of setting host device count
set_host_device_count(4) only helps CPU parallel chains; on the GPU
build it is inert and the default parallel method falls back to running
chains sequentially with a warning. Drop it and use
chain_method="vectorized", which runs all four chains on a single
device — efficient on one GPU and portable to CPU users, while still
yielding multiple chains for the R-hat diagnostic.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Explain what a NumPyro model represents for first-time readers
NumPyro's style is idiosyncratic for someone new to it. Add a conceptual
on-ramp before the model: a model is a *declaration* of the generative
story (not a computation, returns nothing, never called directly) that
an inference engine traces; the obs keyword decides whether a sample
site is latent or observed (the likelihood); the string site names are
the engine's handles. Also add a short note on JAX PRNG keys, explaining
why data uses NumPy's generator while NUTS uses random.PRNGKey.
Prose only; code cells unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Use jax.random.key instead of the legacy PRNGKey
random.PRNGKey is the legacy key constructor; JAX now recommends
random.key, which returns a typed key. NumPyro accepts it throughout
(MCMC, SVI, sample_posterior). Switch all four uses. Not deprecated yet,
but this future-proofs the lecture.
Verified end-to-end with jupytext export + headless execution.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Make run_nuts model-agnostic
Previously run_nuts hard-coded binomial_model as a global and took a
prior argument that duplicated the model's own prior argument (inviting
confusion about whether they could differ — they couldn't, since the
prior was forwarded). Pass the model explicitly and forward *args to it,
so there is no hidden global, the name is honest, and the prior is
supplied exactly once. Call sites become run_nuts(binomial_model, ...).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Tie down Y notation in the ELBO derivation
The lecture denotes the data by k throughout, but the ELBO derivation
uses generic Y. Add a clause defining Y as the observed data (the count
k) so the switch doesn't trip readers.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Fix prior-density plot to respect the prior's support
numpyro's dist.Uniform.log_prob returns its constant value everywhere,
ignoring [low, high], so the restrictive-uniform example plotted a flat
prior across all of [0, 1] instead of a box on [0.5, 0.95]. Mask the
plotted density to prior.support(grid). Sampling was already correct
(NUTS respects the support); this only affects the prior curve.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 file changed
Lines changed: 272 additions & 1043 deletions
