API Reference

gaussx is structured linear algebra, Gaussian distributions, and exponential-family primitives for JAX, built on lineax, Equinox, and matfree. The reference is organised by the package's layered architecture rather than dumped as one flat page:

Section	Layer	What's inside
Primitives	0	Pure functions with structural dispatch — solve, logdet, cholesky, trace, diag, sqrt, inv, eig, svd, root decompositions
Operators & Tags	1	Kronecker, BlockDiag, LowRankUpdate, Toeplitz, block-tridiagonal and kernel operators, plus the structural tags that drive dispatch
Solvers & Preconditioners	1.5	Solver strategy objects (DenseSolver, CGSolver, BBMMSolver, SLQ logdets), the linear_solve front door, and preconditioners
Linear-Algebra Utilities	—	safe_cholesky, symmetrize, Woodbury and Schur identities, matrix-RHS solves, tridiagonal solves
Distributions & Exponential Family	2	MultivariateNormal / MultivariateNormalPrecision, Gaussian sugar ops, KL divergences, natural-parameter conversions
Gaussian Processes	3	Conditioning, whitening, prediction caches, Matheron updates, ELBOs, LOVE / LOO, OILMM projections
Kernels & Approximations	3	Nyström and RFF operators, EigenPro preconditioning, HSIC / MMD, grid + interpolation helpers
Quadrature & Moment Matching	3	Integrators (Gauss-Hermite, unscented, Taylor, MC), likelihoods, kernel expectations, uncertain-input GP prediction
State-Space Models & Kalman	3	SDE kernels, Kalman filter / RTS smoother (sequential, parallel, infinite-horizon), SpInGP, CVI sites
Bayesian Inference & Ensembles	3	Bayesian linear regression, Newton / natural-gradient updates, ensemble Kalman primitives (localization, inflation, ETKF)

Conventions

A few patterns hold across the whole package:

• Operators are lineax operators. Every structured matrix extends lineax.AbstractLinearOperator and is an immutable equinox.Module pytree — safe under jit / grad / vmap. Dense matrices enter the system as lx.MatrixLinearOperator(A, tags).

• Tags drive dispatch. Primitives inspect operator structure (Kronecker, block-diagonal, low-rank, …) and properties (lineax.positive_semidefinite_tag, symmetric, triangular) to pick the cheapest algorithm. Tag your operators — an untagged dense PSD matrix falls back to LU where a tagged one gets Cholesky.

• solver=None means structural dispatch. Functions that accept an optional solver:AbstractSolverStrategy use the structure-aware default when it is None; pass CGSolver(), BBMMSolver(), or a ComposedSolver to override the numerical path without touching the math.

• Lazy over dense. Primitives like inv, sqrt, and cholesky return operators, not arrays, wherever structure allows; nothing is materialized until .as_matrix() is called. When a structured path has to densify, a DenseFallbackWarning is emitted.

• Pure functions. Outside the operator classes everything is a pure function: arrays and operators in, arrays and operators out. PRNG keys are explicit arguments for every stochastic routine.

Every public class and function carries a Google-style docstring with shapes in jaxtyping notation; tensor contraction and reshaping inside the package go through einx.