lrux: JAX Low-Rank Matrix Updates

• lrux is a JAX-based numerical library designed to perform fast low-rank updates for determinants and Pfaffians, reducing the computational cost in QMC workflows.
• It leverages the matrix determinant lemma and Pfaffian identities to update inverses and ratios in O(n² k) per update, dramatically improving efficiency over naive O(n³) methods.
• The library supports JAX transformations like jit, vmap, and autodiff, and employs delayed update schemes for optimal accelerator-oriented performance in fermionic simulations.

lrux is a JAX-based numerical library for fast low-rank updates of determinants of general $n\times n$ matrices and Pfaffians of $2n\times 2n$ skew-symmetric matrices, with an explicit focus on quantum Monte Carlo (QMC) and fermionic neural quantum states (NQS). Its central purpose is to target the dominant computational bottleneck in workflows where successive Monte Carlo moves alter only a few rows, columns, or entries of a Slater or Pfaffian matrix. In that setting, naive recomputation costs $\mathcal{O}(n^3)$ per evaluation, whereas lrux exploits low-rank structure to reduce the per-update cost to $\mathcal{O}(n^2 k)$ for rank-$k$ changes with $k\ll n$. The package supports both determinant and Pfaffian updates, delayed-update strategies, JAX transformations such as jit, vmap, and autodiff, and both real and complex data types (Chen et al., 5 Feb 2026).

1. Computational role and problem setting

In electronic QMC, including VMC, DMC, and AFQMC, the many-fermion wavefunction is evaluated repeatedly at changing electronic configurations $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$. Antisymmetry under particle exchange is commonly represented either through Slater determinants,

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$

or through Pfaffians for paired states,

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$

Both evaluations are $\mathcal{O}(n^3)$ operations when performed from scratch via LU or Pfaffian factorization (Chen et al., 5 Feb 2026).

The bottleneck becomes acute because QMC executes millions to billions of Metropolis or related moves, and each move typically changes only one electron position, one orbital occupation, or a small subset of rows and columns. The same low-rank structure also appears in more advanced updates and in NQS architectures with backflow transformations. lrux is designed around the observation that these local modifications can be represented as low-rank perturbations and that previously computed inverses and determinant or Pfaffian values can therefore be reused rather than recomputed (Chen et al., 5 Feb 2026).

This design places lrux in the class of numerical kernels for antisymmetric wavefunction evaluation rather than in the class of standalone dense linear algebra packages. A plausible implication is that its significance lies less in introducing new asymptotic identities than in consolidating determinant LRUs, Pfaffian LRUs, delayed updates, and JAX-native differentiability into a single implementation suitable for accelerator-oriented QMC and NQS workflows.

2. Low-rank update formalism for determinants and Pfaffians

For determinants, lrux considers a sequence of matrices

$2n\times 2n$0

with $2n\times 2n$1. The matrix determinant lemma yields

$2n\times 2n$2

With

$2n\times 2n$3

the determinant ratio is

$2n\times 2n$4

If $2n\times 2n$5 is already available, constructing $2n\times 2n$6 costs $2n\times 2n$7, and the inverse update follows from Sherman–Morrison–Woodbury: $2n\times 2n$8 The stated conditions are that $2n\times 2n$9 must be invertible and that $\mathcal{O}(n^3)$0 must be invertible; if $\mathcal{O}(n^3)$1 is singular, then $\mathcal{O}(n^3)$2 is singular (Chen et al., 5 Feb 2026).

For Pfaffians, the update is expressed in skew-symmetric low-rank form,

$\mathcal{O}(n^3)$3

where $\mathcal{O}(n^3)$4 is $\mathcal{O}(n^3)$5 and

$\mathcal{O}(n^3)$6

Using the Pfaffian low-rank identity,

$\mathcal{O}(n^3)$7

and setting $\mathcal{O}(n^3)$8, $\mathcal{O}(n^3)$9, lrux obtains

$\mathcal{O}(n^2 k)$0

with ratio

$\mathcal{O}(n^2 k)$1

The inverse update takes the Woodbury-type form

$\mathcal{O}(n^2 k)$2

In the common rank-$\mathcal{O}(n^2 k)$3 case, the expression simplifies further to antisymmetrized outer products of $\mathcal{O}(n^2 k)$4 and $\mathcal{O}(n^2 k)$5, scaled by the scalar Pfaffian ratio (Chen et al., 5 Feb 2026).

The resulting complexity reduction is summarized below.

Method	Per-update cost	Setting
Naive determinant or Pfaffian recomputation	$\mathcal{O}(n^2 k)$6	Fresh factorization
Determinant LRU	$\mathcal{O}(n^2 k)$7	Ratio + inverse update
Pfaffian LRU	$\mathcal{O}(n^2 k)$8	Ratio + inverse update

If initialization costs $\mathcal{O}(n^2 k)$9 once and one then performs $k$0 rank-$k$1 updates, the total cost becomes

$k$2

which is effectively $k$3 per step for long Markov chains with modest $k$4 (Chen et al., 5 Feb 2026).

3. JAX-native implementation and exposed functionality

lrux implements determinant LRUs by maintaining $k$5 and $k$6. For each update $k$7, it computes $k$8, forms the ratio $k$9, updates the determinant by multiplication with $k\ll n$0, and optionally updates the inverse through the Sherman–Morrison–Woodbury expression. The usage pattern shown in the paper is $2n\times 2n$11 Initial inverses are computed via jax.numpy.linalg.inv, and determinants of small $k\ll n$1 matrices are evaluated through JAX linear algebra (Chen et al., 5 Feb 2026).

Pfaffian LRUs follow the analogous pattern. The library maintains $k\ll n$2 and $k\ll n$3, computes

$k\ll n$4

updates the Pfaffian multiplicatively, and optionally updates the inverse. The corresponding interface is $2n\times 2n$12 The package also provides lrux.pf(A) for Pfaffian evaluation via a skew-symmetric factorization, aimed at stability rather than raw speed, and lrux.slogpf(A) returning a (sign, log_abs) pair to avoid overflow or underflow (Chen et al., 5 Feb 2026).

A distinctive implementation detail is the handling of gradients for Pfaffians. The paper derives

$k\ll n$5

with $k\ll n$6 enforcing skew-symmetry. lrux implements custom JAX JVPs through jax.custom_jvp for pf and slogpf, so that these functions remain differentiable and compatible with autodiff, including complex dtypes (Chen et al., 5 Feb 2026).

The library is explicitly organized around JAX transformations. All core routines are JIT-friendly; the examples use $2n\times 2n$13 with static_argnums for non-array arguments such as return_update and donate_argnums for donating Ainv. The benchmarks also rely on jax.vmap to evaluate thousands of determinants or Pfaffians in parallel, mimicking batched wavefunction evaluations in modern VMC. The package supports real and complex dtypes in both single and double precision, although the paper strongly recommends jax_enable_x64=True for QMC workloads (Chen et al., 5 Feb 2026).

4. Delayed-update algorithms and accelerator-oriented optimization

The direct inverse update in both determinant and Pfaffian LRUs involves outer products between $k\ll n$7 and $k\ll n$8 factors. Although these operations are only $k\ll n$9 in floating-point count, the paper emphasizes that they are memory-bandwidth bound when $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$0 is small, especially on GPUs, because large $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$1 data must be streamed at every step. lrux addresses this through delayed updates, which accumulate low-rank corrections in factored form and periodically flush them into the full inverse (Chen et al., 5 Feb 2026).

For determinants, with

$\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$2

the delayed representation is

$\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$3

The ratio matrix can then be computed as

$\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$4

and the factors $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$5 and $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$6 follow corresponding recurrences. The reported cost is $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$7 per step with memory $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$8. To prevent unbounded growth in $\{\mathbf{x}_1,\dots,\mathbf{x}_n\}$9, the implementation uses a maximum delay $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$0; when $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$1, the full inverse is reconstructed, reset as the new base inverse, and the accumulated factors are cleared (Chen et al., 5 Feb 2026).

For Pfaffians, the delayed structure is analogous. Defining

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$2

the inverse is represented as

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$3

The ratio matrix and new factors are computed using sums over previous delayed terms rather than explicit updates of the full inverse at every step. The same asymptotic cost and memory bounds are quoted: $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$4 The paper recommends $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$5 as a heuristic for Pfaffians and notes that, in the determinant example, $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$6 preserves overall $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$7 scaling while reducing memory traffic substantially (Chen et al., 5 Feb 2026).

The user-facing interface exposes these delayed schemes through carrier objects: $2n\times 2n$14 and $2n\times 2n$15 The carrier stores the base inverse, the accumulated low-rank factors, and bookkeeping state. This suggests that lrux treats delayed updates not as an internal optimization hidden from the caller, but as a first-class stateful abstraction tailored to JAX programs that pass explicit state through jit-compiled and vmap-vectorized functions.

5. QMC and NQS workflow integration

lrux is designed as a drop-in component for QMC codes. The workflow described in the paper begins from a configuration of electron coordinates or occupations, constructs the corresponding Slater or Pfaffian matrix $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$8, computes $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_n) = \det\big[\phi_j(\mathbf{x}_i)\big]_{i,j=1}^n,$$9 or $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$0 together with $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$1, and then processes each Monte Carlo move by expressing the matrix change in low-rank form, computing the wavefunction ratio with lrux.det_lru or lrux.pf_lru, accepting or rejecting the move based on that ratio, and updating the inverse either directly or through delayed updates (Chen et al., 5 Feb 2026).

The package supports several common update patterns. A single row update is written as

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$2

with $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$3 one-hot and $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$4 dense. A single column update is

$$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$5

Multiple rows or simultaneous row-and-column changes are likewise expressed as $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$6, with $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$7 built from concatenated one-hot vectors and $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$8 from stacked dense updates. For Pfaffians, the paper gives explicit skew-symmetry-preserving examples involving one row and one column or two rows and two columns, cast into the $$\Psi(\mathbf{x}_1,\dots,\mathbf{x}_{2n}) = \mathrm{pf}\,A.$$9 form (Chen et al., 5 Feb 2026).

A practical feature is the representation of update factors as dense, one-hot, or hybrid tuples such as (dense_part, one_hot_indices). The paper states that this can reduce complexity from $\mathcal{O}(n^3)$0 to $\mathcal{O}(n^3)$1 in some operations by replacing expressions like u @ M @ v_dense with index-based access such as u @ M[:, v_onehot] (Chen et al., 5 Feb 2026).

The target application set explicitly includes Variational Monte Carlo, Diffusion Monte Carlo, Auxiliary-field QMC, and neural quantum states such as FermiNet, neural backflow, and Pfaffian NQS. The paper further states that low-rank structure often appears in NQS when only a small set of orbitals is updated or when neural backflow has a low-rank Jacobian, and that lrux is already used in Pfaffian NQS implementations cited by the authors. Beyond QMC, the same numerical pattern can apply to lattice fermion simulations, optimization of antisymmetric ansätze in variational algorithms, and certain Gaussian process or kernel methods where determinant updates occur (Chen et al., 5 Feb 2026).

6. Performance, stability, limitations, and availability

The benchmark hardware reported in the paper is an NVIDIA A100-80GB GPU. In one-step scaling tests, 1024 determinants or Pfaffians were evaluated in parallel with vmap, using rank $\mathcal{O}(n^3)$2 and varying matrix size $\mathcal{O}(n^3)$3. The results show $\mathcal{O}(n^3)$4 scaling for direct computation and $\mathcal{O}(n^3)$5 scaling for both LRU variants at large $\mathcal{O}(n^3)$6. At $\mathcal{O}(n^3)$7, the reported speedups are $\mathcal{O}(n^3)$8 for determinants and $\mathcal{O}(n^3)$9 for Pfaffians relative to direct computation. In delayed-update tests with 16,384 determinants or Pfaffians in parallel, matrix size $2n\times 2n$00, and $2n\times 2n$01 consecutive updates, total time was $2n\times 2n$02 lower than with direct inverse updates; the optimum delay on the reported system was $2n\times 2n$03 for determinants and $2n\times 2n$04 for Pfaffians (Chen et al., 5 Feb 2026).

The paper places strong emphasis on numerical stability. LRUs require $2n\times 2n$05 and small $2n\times 2n$06 matrices $2n\times 2n$07 to remain well-conditioned. If $2n\times 2n$08 or $2n\times 2n$09, or under the analogous Pfaffian singularity conditions, floating-point errors can grow substantially. Double precision is therefore strongly recommended, especially for long update sequences in QMC. The paper also recommends periodic refactorization from scratch in order to reset rounding error and prevent drift of the maintained inverse away from the true inverse; in delayed-update mode, each block-ending reconstruction provides a partial refresh of this sort (Chen et al., 5 Feb 2026).

Several limitations are explicit. Matrices must be square, Pfaffian matrices must be skew-symmetric and even-dimensional, and performance gains are most pronounced only when rank $2n\times 2n$10 is relatively small. Delayed-update tuning is hardware- and size-dependent, so benchmark-based parameter selection is advised. The Pfaffian routines prioritize stability rather than ultimate speed. This suggests that lrux is optimized for composability, accelerator use, and differentiable programming rather than for replacing all specialized architecture-specific Pfaffian codes.

The software is open source at https://github.com/ChenAo-Phys/lrux, installable via pip install lrux, and was reported at version 0.1.2 with compatibility jax>=0.4.4. The repository includes unit tests, correctness checks against direct determinant and Pfaffian evaluation, example scripts for basic and delayed LRUs, and benchmark scripts for the reported scaling and performance results (Chen et al., 5 Feb 2026).