Add ensemble sampling methods to BlackJAX

[williamjameshandley]

This PR adds two affine-invariant ensemble MCMC samplers to BlackJAX: the stretch move sampler and ensemble slice sampling.

Overview

Ensemble samplers use multiple interacting walkers to explore the posterior distribution without requiring gradient information. They are popular in astronomy due to their self-tuning capability. Despite their inability to use gradient information, there is still a benefit to JAX if the probability density is appropriately vectorized over multiple calls.

Stretch Move Sampler

The stretch move is an affine-invariant ensemble sampler introduced by Goodman & Weare (2010) and popularized by the emcee package (Foreman-Mackey et al. 2013).

Key features:

• Affine-invariant: performance independent of linear transformations
• Uses complementary ensemble to propose moves along random directions
• Single tuning parameter a (stretch scale)
• Supports arbitrary nsplits for group splitting (default: 2)
• Optional randomized splitting for improved mixing

Ensemble Slice Sampling

Ensemble slice sampling (ESS) extends the slice sampling framework to ensemble methods, as implemented in the Zeus package (Karamanis et al. 2021).

Key features:

• Combines slice sampling with ensemble proposals
• Three direction modes: differential, random, gaussian
• Adaptive tuning of scale parameter mu via Robbins-Monro
• Built-in convergence diagnostics through expansion/contraction ratios
• Supports arbitrary nsplits for group splitting (default: 2)

API Design

Both samplers follow BlackJAX conventions and expose consistent APIs:

import blackjax

# Stretch move sampler
stretch = blackjax.ensemble.stretch.as_top_level_api(
    logdensity_fn,
    a=2.0,
    randomize_split=True,
    nsplits=2
)

# Ensemble slice sampling
slice_sampler = blackjax.ensemble.slice.as_top_level_api(
    logdensity_fn,
    move="differential",  # or "random", "gaussian"
    mu=1.0,
    tune=True,
    randomize_split=True,
    nsplits=2
)

Both samplers:

• Support PyTree coordinates for structured parameters
• Return EnsembleState with coords, log_probs, and optional blobs
• Provide acceptance diagnostics via EnsembleInfo
• Are fully compatible with JAX transformations (jit, vmap)

Implementation

The ensemble module is organized into three files:

• base.py: Shared utilities for ensemble manipulation (shuffle, split, concatenate, complement operations)
• stretch.py: Stretch move implementation (~250 lines)
• slice.py: Ensemble slice sampling implementation (~700 lines including direction functions and slice procedure)

Extensive refactoring consolidated common patterns into base.py, eliminating ~165 lines of duplication while maintaining clean, readable implementations.

Testing

Comprehensive test coverage with 15 tests across both samplers:

Stretch sampler (4 tests):

• Basic stretch move mechanics
• 2D Gaussian sampling and convergence
• PyTree coordinate support

Slice sampler (11 tests):

• All three direction functions (differential, random, gaussian)
• 1D and 2D Gaussian sampling
• PyTree support for coordinates and directions
• JIT compilation
• Initialization with and without blobs

All tests pass with 100% coverage of core functionality.

References

• Goodman & Weare (2010). "Ensemble samplers with affine invariance." Communications in Applied Mathematics and Computational Science, 5(1), 65-80. DOI: 10.2140/camcos.2010.5.65
• Foreman-Mackey et al. (2013). "emcee: The MCMC Hammer." PASP, 125(925), 306. arXiv: 1202.3665
• Karamanis et al. (2021). "zeus: A Python implementation of Ensemble Slice Sampling for efficient Bayesian parameter inference." arXiv: 2105.03468

Code references:

• emcee: Stretch move reference implementation
• Zeus: Ensemble slice sampling reference implementation

[junpenglao]

Wow this looks great! Thanks for working on it!

[junpenglao]

Please add a test case in https://github.com/blackjax-devs/blackjax/blob/main/tests/mcmc/test_sampling.py

[junpenglao]

Nit: unnest and rewrite it as a regular if...else... block

[codecov]

Codecov Report

❌ Patch coverage is 96.83908% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 98.44%. Comparing base (56df032) to head (6f8054d).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
blackjax/ensemble/slice.py	97.95%	4 Missing ⚠️
blackjax/ensemble/stretch.py	94.87%	4 Missing ⚠️
blackjax/ensemble/base.py	95.52%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #797      +/-   ##
==========================================
- Coverage   98.45%   98.44%   -0.02%     
==========================================
  Files          66       72       +6     
  Lines        3242     3719     +477     
==========================================
+ Hits         3192     3661     +469     
- Misses         50       58       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.