JAX for Bayes
summary
Implements a regression-for-classification on a small golf dataset across five model libraries and runs inference with four inference libraries (JAX-backed numpyro, TFP, BlackJAX, etc.), showing prior predictive checks (100 prior draws), MCMC sampling (4,000 draws → ESS ≈ 2,200, R-hat ≈ 1), and diagnostics that catch integration errors. Uses JAX autodiff, vmap and jit to compose model definitions, gradient NUTS samplers, and GPU/TPU acceleration, and demonstrates which library pieces interoperate and when results are practically equivalent.
AI-written summary · about
chapters (16)
- 0:00JAX for Bayesian analysis intro
- 1:06Bayesian workflow overview
- 3:56Bayesian model in TFP
- 8:33Fitting model with JAX NUTS
- 11:47Model validation with RV diagnostics
- 15:15Model with geometric constraints
- 19:08Q&A on Bayesian workflow
- 22:01Probabilistic programming ecosystem
- 26:58Split model definition and compute
- 30:07BlackJAX for MCMC sampling
- 31:09TFP JAX-based samplers
- 33:39Numpyro with JAX
- 37:17PIMC and JAX
- 38:20Optimize log probabilities
- 40:09About JAX development
- 41:18Q&A and closing
related talks
11mIntro to JAX: Accelerating Machine Learning research
31mHigh performance machine learning with JAX
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
23mProbabilistic Programming with Vectorized Programmable Inference
01h 10mJAX: accelerated machine learning research via composable function transformations in Python
02h 24mMagical NumPy with JAX
26mTraining Large Deep Learning Models: PyTorch vs JAX
08mDemo: JAX, Flax and Gemma