lrux: Fast low-rank updates of determinants and Pfaffians in JAX

Read full text
Authors: Ao Chen, Christopher Roth
Cited by: 0

Abstract

We present lrux, a JAX-based software package for fast low-rank updates of determinants and Pfaffians, targeting the dominant computational bottleneck in various quantum Monte Carlo (QMC) algorithms. The package implements efficient low-rank updates that reduce the cost of successive wavefunction evaluations from $\mathcal{O}(n^3)$ to $\mathcal{O}(n^2k)$ when the update rank $k$ is smaller than the dimension $n$ of matrices. Both determinant and Pfaffian updates are supported, together with delayed-update strategies that trade floating-point operations for reduced memory traffic on modern accelerators. lrux natively integrates with JAX transformations such as JIT compilation, vectorization, and automatic differentiation, and supports both real and complex data types. Benchmarks on GPUs demonstrate up to $1000\times$ speedup at large matrix sizes. lrux enables scalable, high-performance evaluation of antisymmetric wavefunctions and is designed as a drop-in component for a wide range of QMC workflows. lrux is available at https://github.com/ChenAo-Phys/lrux.

Fields of Study & Subjects

Quantum Computing Algorithms and Architecture Parallel Computing and Optimization Techniques Matrix Theory and Algorithms

Publication details

Journal / source
arXiv (Cornell University)
Publication date
5 February 2026
Document type
Preprint