$ show probml/dynamax
Dynamax is a JAX-native library for probabilistic state space models (SSMs), providing inference (filtering, smoothing, posterior sampling) and learning (EM, SGD, blocked Gibbs) for hidden Markov models (HMMs), linear Gaussian SSMs (LGSSMs), nonlinear Gaussian SSMs (NLGSSMs), generalized Gaussian SSMs (GGSSMs), and switching linear dynamical systems (SLDS). The architecture separates low-level pure functions (e.g., `hmm_filter`, `lgssm_filter`) from high-level object-oriented model classes (e.g., `GaussianHMM`, `LinearGaussianSSM`), all built on JAX for automatic differentiation, vectorization, and GPU/TPU acceleration. The library includes advanced algorithms such as parallel associative-scan filtering/smoothing, information-form Kalman filters, Rao-Blackwellized particle filters, and conditional-moments Gaussian filters (EKF, UKF, GHKF). It is designed for research use, with a focus on numerical correctness and performance, and is accompanied by extensive documentation, tutorials, and a JOSS publication.
The LLM identified 5 files worth highlighting for reuse:
Provides a `ParameterProperties` system for managing trainable and constrained parameters in JAX, with transformations to unconstrained space and log-det Jacobian computation. This pattern is reusable for any JAX-based probabilistic model requiring parameter constraints and gradient-based inference.
Implements custom Bayesian matrix distributions (InverseWishart, MatrixNormalInverseWishart, NormalInverseWishart, NormalInverseGamma) with analytic moments and posterior updates. These non-trivial distributions can be extracted for use in other probabilistic programming projects.
Defines abstract components for HMMs (initial state, transitions, emissions) and a composable `HMM` base class. This modular design allows users to define custom emission distributions and combine them with standard components, and could be extracted as a standalone HMM framework.
Implements parallel (associative scan) filtering and smoothing for linear Gaussian SSMs, achieving O(log T) time complexity on GPU. This non-trivial algorithm implementation can be reused in other JAX projects requiring parallel state-space inference.
Provides a JAX-native k-means implementation with k-means++ initialization, multiple restarts, and numerical stability guards. It is self-contained and can be extracted for clustering tasks in JAX.