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

TL;DR

lrux is a JAX-based software package that accelerates low-rank determinant and Pfaffian updates, significantly improving the efficiency of quantum Monte Carlo algorithms on modern hardware.

Contribution

The paper introduces lrux, a novel JAX-compatible library that reduces update costs from cubic to quadratic time and supports advanced strategies for high-performance quantum computations.

Findings

Achieves up to 1000x speedup on GPUs for large matrices

Supports both determinant and Pfaffian updates with delayed strategies

Integrates seamlessly with JAX for automatic differentiation and compilation

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.