Track synthesis, materials and catalysis as it is published.
Search papers, labs, and topics across Lattice.
The authors develop `modax`, a JAX-compatible library of GPU-accelerated explicit and linearly implicit ODE solvers specifically architected for massive ensembles of low-dimensional (<200D), uncoupled dynamical systems. Standard high-level vectorized frameworks (like pure JAX) suffer massive performance penalties when ensemble trajectories diverge or become stiff, bottlenecking tasks like Bayesian parameter estimation and continuous-time scientific modeling. By pairing a low-level thread-based execution model with sparsity-exploiting data structures, the framework achieves a 30x speedup over `diffrax` with Tsit5 on chaotic Lorenz systems and over 700x with implicit Rodas5P on stiff Robertson systems.
Divergent ODE trajectories that typically cripple GPU vectorization in high-level frameworks run up to 700x faster when mapped directly to thread-level Rosenbrock-Wanner kernels with sparse matrix support.
Solving large ensembles of small, independent ordinary differential equations (ODEs) is a common task in computational science and engineering, arising for example in Bayesian parameter estimation, Monte Carlo uncertainty quantification, and the integration of uncoupled physical systems such as the linearised Einstein-Boltzmann equations of cosmology. Modern graphics processing units (GPUs) with thousands of cores are well suited to such parallel workloads, but realising this potential requires careful attention to GPU-specific architectural constraints. We present modax, a Python library of GPU-accelerated ODE solvers specifically optimised for solving large ensembles of low-dimensional, independent problems. Our solvers are compatible with the JAX ecosystem but are written in a low-level, thread-based programming paradigm better suited to ensembles of highly divergent trajectories than their pure-JAX counterparts. Benchmarks on the Lorenz, Robertson and van der Pol lattice systems compare modax against diffrax, DiffEqGPU$.$jl and torchdiffeq across ensemble size, dimensionality, and trajectory divergence. Our explicit Tsit5 solver achieves 30x speed-up over diffrax on Lorenz systems, and our implicit Rodas5P solver achieves over 700x speed-up on Robertson systems, demonstrating the superiority of linearly implicit Rosenbrock-Wanner methods on GPUs. Our implicit solver scales well to higher dimensions thanks to the use of sparsity-exploiting techniques and data structures. We conclude by applying modax to three examples from the field of cosmology - Bayesian parameter estimation from primordial abundances, Monte Carlo uncertainty quantification of the global 21cm signal, and the computation of uncoupled Fourier modes in the primordial power spectrum - and demonstrate significant efficiency improvements in our ability to model physical phenomena described by low-dimensional (<200D) ODEs.