gabp-sparse-inv

Differentiable selected inversion for sparse block matrices in PyTorch.

For a sparse symmetric positive-definite block matrix A, the package computes the blocks of A⁻¹ that lie on the sparsity pattern of A, or on its filled (chordal) pattern, without forming the dense inverse. In a Gaussian Markov random field with precision A, these are the marginal covariances of every node and the covariances of neighbouring nodes. Every kernel gives first-order gradients with respect to the input blocks, and the functional kernels also give second derivatives. Linear solves, log-determinants and Gaussian samples come from the same block factorization, and a non-symmetric variant covers structurally symmetric patterns eliminated in a fixed order.

Relation to other tools

The algorithms are classical: block elimination followed by the Takahashi and Erisman–Tinney recurrences for the selected inverse, and reverse-mode differentiation of those recurrences. For SPD matrices the derivative of the selected inverse is, up to sign, the Hessian of the sparse log-determinant barrier, for which Andersen, Dahl and Vandenberghe (2013) give factorization-cost algorithms implemented in CHOMPACK. This package implements them for PyTorch: tensors and autograd, leading batch dimensions, dense b × b blocks, analytic backwards that keep no autograd tape, and tests against dense references.

It is a reference implementation written with per-block Python loops, so on a single large matrix it is much slower than compiled sparse-direct libraries such as CHOLMOD, CHOMPACK or PEXSI. Use those when you need speed outside autograd. Use this package when you need selected inverse entries and their gradients inside a PyTorch computation; Choosing a route gives measured sizes. Leading batch dimensions let many matrices with one sparsity pattern share each block operation.

Install

pip install gabp-sparse-inv

0.4.1 corrects several cases where 0.4.0 returned wrong numbers without raising; the changelog lists what can behave differently.

Examples

The blocks of the inverse of a block-tridiagonal SPD matrix:

import torch
from gabp_sparse_inv import random_spd_chain, selected_inverse_chain

A = random_spd_chain(num_blocks=8, block_size=4, seed=0, diag_load=2.0)
G_diag, G_lower = selected_inverse_chain(A.diag, A.lower, check=True)

dense_inverse = torch.linalg.inv(A.to_dense())
torch.testing.assert_close(G_diag[0], dense_inverse[:4, :4])

A general sparse pattern takes node blocks plus a lower-triangular edge list that names each off-diagonal block once, as (i, j) with i > j. The output follows the filled pattern, and gradients reach every stored input block. An explicit elimination order skips the default minimum-degree search, and batched=True runs one batched block operation per elimination level:

from gabp_sparse_inv import (
    elimination_order_nested_dissection,
    grid_edges,
    random_spd_graph,
    selinv_junction,
)

A = random_spd_graph(12, grid_edges(3, 4), 2, seed=0)
order = elimination_order_nested_dissection(3, 4)
diag = A.diag.clone().requires_grad_(True)
edge_val = A.edge_val.clone().requires_grad_(True)

G_diag, filled_edges, G_lower = selinv_junction(diag, A.edge_index, edge_val, order, batched=True)
total_variance = G_diag.diagonal(dim1=-2, dim2=-1).sum()
total_variance.backward()          # fills diag.grad and edge_val.grad

An edge list that holds both directions of every edge, as in PyTorch Geometric, raises ValueError; keep one direction:

both = torch.cat([A.edge_index, A.edge_index.flip(0)], dim=1)
both_val = torch.cat([A.edge_val, A.edge_val.mT])      # A[j, i] = A[i, j]ᵀ
keep = both[0] > both[1]
G_diag, _, _ = selinv_junction(A.diag, both[:, keep], both_val[keep], order)

Symbolic analysis runs once per pattern and is reused when the values change. One numeric factor serves several queries, differentiably:

from gabp_sparse_inv import analyze_junction

analysis = analyze_junction(A.num_nodes, A.edge_index, order)
factor = analysis.factor(A.diag, A.edge_val, check=True, batched=True)

logdet = factor.logdet()
x = factor.solve(torch.ones(A.num_nodes, 2, dtype=A.diag.dtype))
variances = factor.marginal_diag()
samples = factor.sample(4, generator=torch.Generator().manual_seed(0))

for shift in (0.0, 0.5):           # new values, same pattern: no new symbolic work
    d = (A.diag + shift * torch.eye(2, dtype=A.diag.dtype)).requires_grad_(True)
    G_diag = analysis.factor(d, A.edge_val, batched=True).selected_inverse()[0]
    G_diag.diagonal(dim1=-2, dim2=-1).sum().backward()

The marginal likelihood of a grid GMRF observed with Gaussian noise, differentiated with respect to its hyperparameters:

from gabp_sparse_inv import grid_gmrf_precision, junction_marginal_log_likelihood

kappa = torch.tensor(1.0, dtype=torch.float64, requires_grad=True)
a = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
sigma2 = torch.tensor(0.1, dtype=torch.float64, requires_grad=True)    # noise variance
y = torch.randn(16, 1, dtype=torch.float64, generator=torch.Generator().manual_seed(0))

diag, edge_index, edge = grid_gmrf_precision(4, 4, kappa, a)
log_evidence = junction_marginal_log_likelihood(diag, edge_index, edge, y, sigma2)
log_evidence.backward()            # fills kappa.grad, a.grad and sigma2.grad

Choosing a route

The SPD selected inverse on a general pattern can be differentiated four ways. All four give the same values and first-order gradients.

Without order, each functional call and each analyze_junction call first computes a greedy minimum-degree order in Python with O(n²) node scans. On a 10,000-node path that takes 13.5 s against 0.03 s with an explicit order. For grids, elimination_order_nested_dissection(rows, cols) gives a good order; for other patterns, compute an order once and reuse it, or reuse the analysis.

Time for the value and gradient of the total marginal variance Σ tr(G_vv) on a 2-D grid with scalar nodes (b = 1), nested-dissection order, float64. Measured on one laptop CPU (AMD Ryzen 7 7730U, 8 cores, 8 PyTorch threads, Windows 11, PyTorch 2.14.1); each entry is the median of three runs after a warm-up, or a single run where marked with *. Other hardware and block sizes will move the crossovers.

The dense baseline assembles the matrix with one differentiable scatter, factors it with torch.linalg.cholesky and differentiates cholesky_inverse with autograd; it needs O(n²) memory. The analytic backward avoids the autograd tape but runs its reverse sweep as a per-node Python loop, so with small blocks it is slower than the batched route. The reused-analysis column excludes the symbolic analysis, which is paid once per pattern: 0.03 s, 0.12 s and 0.23 s with the nested-dissection order at these sizes, against 0.14 s, 0.72 s and 1.6 s for the default minimum-degree order.

Contracts and limits

• The SPD kernels require SPD input. A pivot that is not positive definite raises torch.linalg.LinAlgError naming the node. Singular PSD, indefinite and complex-Hermitian matrices are not supported.
• Every call checks the structure of its input: integer indices in range, i > j with each block listed once, a parent array that forms one tree, and block shapes and counts. check=True adds the value checks: finite blocks and symmetric diagonal blocks.
• The non-symmetric kernels eliminate in a fixed order without pivoting, so the condition number of A alone does not control their accuracy: [[1e-18, 1], [2, 1]] has condition number 2.6 and loses every digit when node 0 is eliminated first. They measure element growth locally, against the largest input entries of each block row and column, and raise torch.linalg.LinAlgError above eps^(-1/2) for solves and eps^(-1/4) for selected inverses, whose errors grow faster. The test is a heuristic, not a stability guarantee: an accepted elimination can lose up to about half the working digits, more when A is ill conditioned. On refusal, choose an order that eliminates well-conditioned pivots first, compute in float64 (compute_dtype=torch.float64), or use a pivoted LU. The derivations describe the criterion and a randomized check of it.
• In the tested regimes the forward error of the SPD kernels follows κ(A)·u; there is no proved structured bound for the full algorithm.
• The analytic backwards (selinv_tree, selinv_junction_analytic, selinv_nonsym_junction_analytic) are first order and raise when asked for a differentiable gradient (create_graph=True, as in Hessian-vector products). The functional kernels (selected_inverse_tree, selinv_junction, selinv_nonsym_junction, the solves, log-determinants and likelihoods) support second derivatives.
• Computation runs in float32 or float64. Float16 and bfloat16 work as storage formats with compute_dtype=torch.float32; without it the CPU backend raises, because LAPACK has no half-precision Cholesky.
• Results stay on the device of the input values; index tensors stay on the CPU. Only CPU execution is tested. tests/test_cuda.py checks the main entry points on a CUDA device against the CPU results, but it has not been run for this release.
• A right-hand side is [..., n, b] (vectors) or [..., n, b, k] (columns). A shape that fits both is read as vectors only when it equals diag.shape[:-1], as in torch.linalg.solve; rhs.unsqueeze(-1) selects columns.
• junction_posterior_marginal_variances returns one variance per coordinate: [..., n] for scalar nodes and [..., n, b] for block nodes.
• On trees the cost is O(n b³) time and O(n b²) storage. On general patterns it depends on the elimination order: with w_v neighbours of node v eliminated after it in the filled graph, numeric storage is proportional to Σ(1 + w_v) b² and work to Σ(1 + w_v²) b³.

Testing

pip install -e ".[dev]"
pytest

Each kernel is checked against dense inverses in double precision, on-pattern residual identities, finite differences, gradcheck and gradgradcheck, and the other kernels on shared special cases. CI runs the suite on Linux, Windows and macOS, with the minimum supported PyTorch, and against the installed wheel. CI has no GPU.

Documentation

• Derivations: the recurrences, derivative conventions, costs, numerical boundaries and prior work.
• Architecture: module map and conventions for contributors.
• Changelog and contributing guide.
• Research archive: the research code, experiments and manuscripts that main carried before 0.4.0.

Citation and license

Citation metadata is in CITATION.cff. If the selected inverse matters to your work, please also cite the algorithms: Takahashi, Fagan and Chen (1973), Erisman and Tinney (1975), Lin et al. (2011), and Andersen, Dahl and Vandenberghe (2013). The software is MIT licensed; see LICENSE.