State Space Models Python Packages

dynamax

A Python package for probabilistic state space modeling with JAX

deeptab

DeepTab is a Python package that simplifies tabular deep learning by providing a suite of models for regression, classification, and distributional regression tasks. It includes models such as Mambular, TabM, FT-Transformer, TabulaRNN, TabTransformer, and tabular ResNets.

mambular

DeepTab is a Python package that simplifies tabular deep learning by providing a suite of models for regression, classification, and distributional regression tasks. It includes models such as Mambular, TabM, FT-Transformer, TabulaRNN, TabTransformer, and tabular ResNets.

bayesnewton

Bayes-Newton—A Gaussian process library in JAX, with a unifying view of approximate Bayesian inference as variants of Newton's method.

mlx-recurrence

A plug-in framework for linear-recurrence Metal kernels on Apple Silicon (SSD, GLA, RG-LRU)

identifiability-diagnostic

An O(T) pre-training geometric diagnostic system that detects latent regime unidentifiability via singular value subspace collapse in Switching State-Space Models (S-SSMs) before model training begins.

elume

Agentic memory with mental-model architecture. Deterministic, replay-safe substrate for cognitive-architecture agents: LinOSS temporal encoding, Hopfield attractor memory, mental-model and metacognitive primitives, strategy evolution, pluggable MemEvolve benchmark adapter. Contribution: byte-identical replay, immutable lineage, provider abstraction, composable operators.

rts-smoother

High-performance Rauch–Tung–Striebel Kalman smoother for the local level model

linoss-dynamics

NumPy runtime helpers for LinOSS-style oscillator dynamics — sequence stepping, closed-form damped oscillator, Bayesian filtering, MLE fit. See Roadmap issue for vision.

recurrentlens

Mechanistic interpretability for State-Space Models: SAEs, feature visualization, and a Hub registry for Mamba/Mamba-2.

discretax

State Space Models in Jax

smcjax

Sequential Monte Carlo and particle filtering in JAX

parallelwatch

🧠 Real-time multivariate temporal correlation and anomaly detection primitive. Leverages structured state-space sequence modeling (SSMs) to map system failures across massive telemetry streams without O(M^2) cross-attention overhead.

linax

State Space Models in Jax

lrnnx

A unified PyTorch library for Linear RNNs