Skip to main content
Michael Ellis
← Back to open source projects

Python · Open source

smcjax

Sequential Monte Carlo and Particle Filtering in JAX

Particle filters, state-space simulation, and sequential inference diagnostics implemented in JAX.

View on GitHubPackage

Overview

I built smcjax from the sequential inference work in my master’s thesis. The library includes bootstrap, auxiliary, and Liu-West particle filters along with simulation and diagnostics for state-space models.

The filters use jax.lax.scan and support JIT compilation and vectorization. Model-specific sampling and observation functions stay separate from the filtering code, so the same algorithms can be used with different models.

Technical highlights

  • Bootstrap, auxiliary, and Liu-West particle filters
  • Four resampling schemes through BlackJAX
  • Conditional resampling based on effective sample size
  • Weighted summaries, particle diversity, CRPS, and evidence diagnostics
  • JIT- and vmap-compatible typed APIs

Topics

Sequential Monte CarloParticle filteringHidden Markov modelsBayesian inferenceJAX