Background
BlackJAX is an open-source Python library for Bayesian inference built on JAX. In Bayesian statistics you want to compute a posterior distribution — your updated belief about unknown quantities after observing data. For real problems this is almost never available in closed form, so you approximate it by drawing random samples. BlackJAX provides the algorithms that do that sampling: you hand it a log-probability function and get samples back. Because it runs on JAX, the algorithms compile to XLA and run on GPU/TPU, are composable, and support automatic differentiation.
My contributions are concentrated in the Sequential Monte Carlo (SMC) subsystem. Unlike most samplers that work directly on the final target distribution, SMC builds a sequence of intermediate distributions bridging from something easy to sample (the prior) toward the posterior, evolving a population of weighted particles at each step. Think of it as slowly turning up the heat on a distribution until it reaches your target. This makes SMC especially effective for multimodal posteriors, model comparison (it naturally computes the marginal likelihood), and sequential data settings.
I've developed and maintained this subsystem from its early form into a flexible, production-ready module — covering algorithm implementation, adaptive infrastructure, API design, and correctness.
Summary
| Theme | What I built |
|---|---|
| Algorithm implementation | IRMH, Waste-Free SMC, Partial posteriors SMC |
| Adaptive infrastructure | Inner kernel tuning, pretuning, joint tuning |
| API design | Algorithm class removal, RMH/RW separation, per-particle kernel parameters, tempering decoupling |
| Correctness | IRMH asymmetry fix, SMC-MCMC integration tests, ESS multivariate tests |
| Documentation | SMC particle representation, TemperedSMC example, IRMH proposal docs |
Contributions by release
ESS (Effective Sample Size) measures how well a set of weighted particles represents the target distribution — if it drops too low the population has collapsed and the approximation fails. This PR added tests verifying the ESS calculation was correct for posteriors with multiple parameters (multivariate) and multiple variables (multivariable), two cases that had not been systematically covered before.
IRMH is a Metropolis-Hastings variant where proposals are drawn from a fixed distribution rather than making small local moves. As a mutation kernel inside SMC it efficiently diversifies particles after resampling when you have a good global proposal. This PR implemented IRMH as a first-class algorithm in BlackJAX.
Tempered SMC defines its sequence of distributions via a tempering parameter λ that scales the likelihood: p(θ) · p(data|θ)^λ, with λ going from 0 (prior only) to 1 (full posterior). Fixed the worked example in the documentation so users could correctly reproduce the algorithm.
Reorganized the test suite into a proper folder hierarchy, making it easier for contributors to find and add tests for specific algorithms.
Random-walk Metropolis-Hastings (RMH) and Random Walk were previously conflated in the codebase. This PR separated them into two distinct, properly named algorithms and generalized Random Walk to support non-Gaussian jump distributions — making the API clearer and more flexible for users building custom samplers.
MCMC kernels can be used as the mutation step inside SMC to diversify particles. This PR added an integration test for this combination — verifying the two subsystems work correctly together — and fixed bugs uncovered in the process.
In Metropolis-Hastings the acceptance ratio depends on whether the proposal is symmetric (q(x′|x) = q(x|x′)). IRMH proposals are not symmetric, and this was not being correctly accounted for in the acceptance step. This PR fixed the acceptance ratio computation, making IRMH statistically correct.
In JAX the representation of a population of samples matters — whether it's a stacked array, a pytree, etc. This PR made the particle representation contract explicit so users could correctly pass their models into the SMC interface.
When using an MCMC kernel as the mutation step inside SMC, the kernel's hyperparameters
(step size for HMC, proposal variance for RMH, etc.) need to be tuned at each SMC step
to work well with the current intermediate distribution — which changes every step.
This PR implemented blackjax.smc.inner_kernel_tuning, a module that automatically
adapts the mutation kernel's parameters during the SMC run. It became the most-referenced
part of the SMC subsystem in community discussions.
Previously all particles in an SMC population had to use the same mutation kernel hyperparameters. This PR lifted that restriction, allowing each particle to carry its own kernel parameters — enabling more flexible adaptation strategies where different particles can explore with different step sizes or proposal distributions.
The defining change for BlackJAX v1.2.0. The library had a layer of Algorithm
wrapper classes that was adding complexity without clarity. This PR removed them entirely,
simplifying the top-level API and making the library more composable. Flagged as the
major change in the release notes.
Standard SMC discards low-weight particles during resampling. Waste-Free SMC (Dau & Chopin, 2022) instead applies more MCMC steps to recycle low-weight particles rather than dropping them, improving sample efficiency. This PR implemented Waste-Free SMC and made it available for both the adaptive tempered and standard tempered SMC variants.
Two contributions in one PR. First, partial posteriors SMC: the ability to run SMC on subsets of the data, incrementally incorporating more data at each step — useful for sequential Bayesian updating. Second, a significant architectural refactor decoupling the tempering schedule (how λ evolves) from the SMC construction, making it easy to plug in custom tempering strategies without touching the core SMC code.
Pretuning addresses a bootstrapping problem: before the SMC run starts, you need initial kernel parameters, but you don't yet know what distribution you'll be targeting at the first step. This PR implemented a short calibration phase before the main SMC loop to estimate good starting hyperparameters.
Added a test suite for the combination of pretuning with adaptive tempering — where the temperature schedule is chosen automatically based on the ESS at each step — verifying the two features interact correctly.
The culmination of the pretuning and inner kernel tuning work. This PR unified both into a
single coherent interface — blackjax.smc.base.smc with inner_kernel_tuning
and pretuning jointly configured — making the full adaptive SMC pipeline available
without requiring users to assemble the pieces manually.
The thread running through all of it: making SMC in BlackJAX correct, adaptive, and composable
— so that researchers can use it for real problems without hand-tuning the internals.
Paper: BlackJAX: Composable Bayesian inference in JAX (arXiv 2402.10797) ·
All PRs: github.com/users/ciguaran/projects/1