I am delighted to announce the release of “dynamax”, an open source library for dynamic state space models in JAX. It supports inference and learning in HMMs, Linear Gaussian SSMs, as well as non-linear and non-Gaussian SSMs. See https://github.com/probml/dynamax for the code.

For state estimation, we support various efficient, deterministic methods, such as forwards-backwards, Kalman filtering, sigma point filtering, etc. For parameter estimation, we support MLE via EM or SGD (using optax), and Bayes via Gibbs or (collapsed) HMC (using blackjax).

Dynamax subsumes existing libraries, such as hmmlearn, pykalman, filterpy, etc.
We are also working on various extensions to this library, such as for time-series forecasting, which we will release in the future.

Big thanks to my colleagues @scottlinderman, @petergchang, @gileshd, karalleyna, mxinglongli, et al. and support from