---
title: 'lrux: JAX Low-Rank Matrix Updates'
url: https://www.emergentmind.com/topics/lrux
type: topic
---

# lrux: JAX Low-Rank Matrix Updates

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 [2602.05255].

## 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 [2602.05255].

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 [2602.05255].

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
\[
A_0 \rightarrow A_1 \rightarrow \dots \rightarrow A_t,
\qquad
A_t = A_{t-1} + v_t u_t^T,
\]
with \(u_t,v_t \in \mathbb{C}^{n\times k}\). The matrix determinant lemma yields
\[
\det(A_t)=\det(A_{t-1})\,\det\big(I_k + u_t^T A_{t-1}^{-1} v_t\big).
\]
With
\[
R_t = I_k + u_t^T A_{t-1}^{-1} v_t,
\]
the determinant ratio is
\[
r_t = \frac{\det A_t}{\det A_{t-1}} = \det R_t.
\]
If \(A_{t-1}^{-1}\) is already available, constructing \(R_t\) costs \(\mathcal{O}(n^2 k)\), and the inverse update follows from Sherman–Morrison–Woodbury:
\[
A_t^{-1}
=
A_{t-1}^{-1}
-
A_{t-1}^{-1} v_t R_t^{-1} u_t^T A_{t-1}^{-1}.
\]
The stated conditions are that \(A_{t-1}\) must be invertible and that \(R_t\) must be invertible; if \(R_t\) is singular, then \(A_t\) is singular [2602.05255].

For Pfaffians, the update is expressed in skew-symmetric low-rank form,
\[
A_t - A_{t-1} = -u_t J u_t^T,
\]
where \(u_t\) is \(n\times 2k\) and
\[
J=
\begin{pmatrix}
0 & I_k\\
-I_k & 0
\end{pmatrix},
\qquad
J^{-1}=-J.
\]
Using the Pfaffian low-rank identity,
\[
\frac{\mathrm{pf}(A + BCB^T)}{\mathrm{pf}(A)}
=
\frac{\mathrm{pf}(C^{-1}+B^T A^{-1} B)}{\mathrm{pf}(C^{-1})},
\]
and setting \(B=u_t\), \(C=-J\), lrux obtains
\[
R_t = J + u_t^T A_{t-1}^{-1} u_t,
\]
with ratio
\[
r_t = \frac{\mathrm{pf} A_t}{\mathrm{pf} A_{t-1}} = \frac{\mathrm{pf} R_t}{\mathrm{pf} J}.
\]
The inverse update takes the Woodbury-type form
\[
A_t^{-1}
=
A_{t-1}^{-1}
+
\left(A_{t-1}^{-1}u_t\right) R_t^{-1} \left(A_{t-1}^{-1}u_t\right)^T.
\]
In the common rank-\(k=1\) case, the expression simplifies further to antisymmetrized outer products of \(A_{t-1}^{-1}x\) and \(A_{t-1}^{-1}y\), scaled by the scalar Pfaffian ratio [2602.05255].

The resulting complexity reduction is summarized below.

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

If initialization costs \(\mathcal{O}(n^3)\) once and one then performs \(T\) rank-\(k\) updates, the total cost becomes
\[
\mathcal{O}(n^3) + T\,\mathcal{O}(n^2 k),
\]
which is effectively \(\mathcal{O}(n^2 k)\) per step for long Markov chains with modest \(k\) [2602.05255].

## 3. JAX-native implementation and exposed functionality

lrux implements determinant LRUs by maintaining \(A^{-1}\) and \(\det A\). For each update \(A\mapsto A+vu^T\), it computes \(R=I_k+u^T A^{-1}v\), forms the ratio \(r=\det R\), updates the determinant by multiplication with \(r\), and optionally updates the inverse through the Sherman–Morrison–Woodbury expression. The usage pattern shown in the paper is
```python
r = lrux.det_lru(Ainv, u, v)  # ratio only
r, Ainv = lrux.det_lru(Ainv, u, v, return_update=True)
```
Initial inverses are computed via `jax.numpy.linalg.inv`, and determinants of small \(k\times k\) matrices are evaluated through JAX linear algebra [2602.05255].

Pfaffian LRUs follow the analogous pattern. The library maintains \(A^{-1}\) and \(\mathrm{pf} A\), computes
\[
R = J + u^T A^{-1}u,
\qquad
r=\mathrm{pf}(R)/\mathrm{pf}(J),
\]
updates the Pfaffian multiplicatively, and optionally updates the inverse. The corresponding interface is
```python
r = lrux.pf_lru(Ainv, u)  # ratio only
r, Ainv = lrux.pf_lru(Ainv, u, return_update=True)
```
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 [2602.05255].

A distinctive implementation detail is the handling of gradients for Pfaffians. The paper derives
\[
\frac{\partial\, \mathrm{pf} A'}{\partial A_{ij}}
=
\frac{1}{2}\,\mathrm{pf}(A')\left(A'^{-1}\right)_{ji},
\]
with \(A'=(A-A^T)/2\) 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 [2602.05255].

The library is explicitly organized around JAX transformations. All core routines are JIT-friendly; the examples use
```python
det_lru_fn = jax.jit(lrux.det_lru, static_argnums=3, donate_argnums=0)
```
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 [2602.05255].

## 4. Delayed-update algorithms and accelerator-oriented optimization

The direct inverse update in both determinant and Pfaffian LRUs involves outer products between \(n\times k\) and \(k\times n\) factors. Although these operations are only \(\mathcal{O}(n^2 k)\) in floating-point count, the paper emphasizes that they are memory-bandwidth bound when \(k\) is small, especially on GPUs, because large \(n^2\) 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 [2602.05255].

For determinants, with
\[
a_t = A_{t-1}^{-1} v_t,
\qquad
b_t = A_{t-1}^{-T} u_t R_t^{-T},
\]
the delayed representation is
\[
A_\tau^{-1} = A_0^{-1} - \sum_{t=1}^{\tau} a_t b_t^T.
\]
The ratio matrix can then be computed as
\[
R_\tau
=
I_k + u_\tau^T A_0^{-1} v_\tau
-
\sum_{t=1}^{\tau-1}
\big(u_\tau^T a_t\big)\big(b_t^T v_\tau\big),
\]
and the factors \(a_\tau\) and \(b_\tau\) follow corresponding recurrences. The reported cost is \(\mathcal{O}(n^2 k + \tau n k^2)\) per step with memory \(\mathcal{O}(n^2 + \tau n k)\). To prevent unbounded growth in \(\tau\), the implementation uses a maximum delay \(T\); when \(\tau=T\), the full inverse is reconstructed, reset as the new base inverse, and the accumulated factors are cleared [2602.05255].

For Pfaffians, the delayed structure is analogous. Defining
\[
a_t = A_{t-1}^{-1}u_t,
\]
the inverse is represented as
\[
A_\tau^{-1}
=
A_0^{-1}
+
\sum_{t=1}^{\tau} a_t R_t^{-1} a_t^T.
\]
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:
\[
\mathcal{O}(n^2 k + \tau n k^2)
\quad\text{per step,}\qquad
\mathcal{O}(n^2 + \tau n k)
\quad\text{memory.}
\]
The paper recommends \(T \approx n/(20k)\) as a heuristic for Pfaffians and notes that, in the determinant example, \(T\lesssim n/k\) preserves overall \(\mathcal{O}(n^2 k)\) scaling while reducing memory traffic substantially [2602.05255].

The user-facing interface exposes these delayed schemes through carrier objects:
```python
carrier = lrux.init_det_carrier(A, max_delay)
r, carrier = lrux.det_lru_delayed(carrier, u, v, return_update=True, current_delay=τ)
```
and
```python
carrier = lrux.init_pf_carrier(A, max_delay)
r, carrier = lrux.pf_lru_delayed(carrier, u, return_update=True, current_delay=τ)
```
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 \(A\), computes \(\det A\) or \(\mathrm{pf} A\) together with \(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 [2602.05255].

The package supports several common update patterns. A single row update is written as
\[
A_1-A_0 = e_i u^T = vu^T,
\]
with \(v=e_i\) one-hot and \(u\) dense. A single column update is
\[
A_1-A_0 = v e_j^T = vu^T.
\]
Multiple rows or simultaneous row-and-column changes are likewise expressed as \(vu^T\), with \(v\) built from concatenated one-hot vectors and \(u\) 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 \(-uJu^T\) form [2602.05255].

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^2)\) to \(\mathcal{O}(n)\) in some operations by replacing expressions like `u @ M @ v_dense` with index-based access such as `u @ M[:, v_onehot]` [2602.05255].

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 [2602.05255].

## 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 \(k=1\) and varying matrix size \(n\). The results show \(\mathcal{O}(n^3)\) scaling for direct computation and \(\mathcal{O}(n^2)\) scaling for both LRU variants at large \(n\). At \(n=1024\), the reported speedups are \(\sim 200\times\) for determinants and \(\sim 1000\times\) for Pfaffians relative to direct computation. In delayed-update tests with 16,384 determinants or Pfaffians in parallel, matrix size \(n=128\), and \(128\) consecutive updates, total time was \(20\%-40\%\) lower than with direct inverse updates; the optimum delay on the reported system was \(T=16\) for determinants and \(T=4\) for Pfaffians [2602.05255].

The paper places strong emphasis on numerical stability. LRUs require \(A_{t-1}^{-1}\) and small \(k\times k\) matrices \(R_t\) to remain well-conditioned. If \(\det A_{t-1}\approx 0\) or \(\det R_t\approx 0\), 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 [2602.05255].

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 \(k\) 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 [2602.05255].

Source: https://www.emergentmind.com/topics/lrux