thehardwarestore.ai

Accelerated Scan is a focused, high-performance library for solving first-order linear recurrences (associative scans) on GPU. It provides a pure CUDA kernel (`warp`) that uses warp shuffle intrinsics, shared memory, and a hierarchical chunked algorithm to achieve state-of-the-art speed, and a Triton-based implementation (`scalar`) that leverages `tl.associative_scan` for portability. A reference PyTorch implementation (`ref`) implements the classic Blelloch tree scan for validation. The library is designed for integration into state space models and linear RNNs, where the recurrence `x[t] = gate[t] * x[t-1] + token[t]` is central. The codebase is small but exceptionally well-crafted, with clear separation of concerns, comprehensive tests, and a clean API. The CUDA kernel in particular is a standout example of GPU optimization, using warp-level primitives, shared memory, and chunking to handle sequences of arbitrary length (up to 65536) with minimal overhead.

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.