Architecture

gaussx is organised as a layered stack on top of lineax and matfree. Each layer is usable on its own, and each layer builds on the ones beneath it (with one exception, the dispatch core, described under Import structure). You can enter wherever your problem lives: grab a single primitive, build a structured operator, swap a solver strategy, or call a finished recipe.

The stack at a glance

flowchart TB
    subgraph L3["Layer 3 · Recipes"]
        direction LR
        GP["Gaussian processes<br/><small>conditioning · ELBOs · whitening · LOVE · OILMM</small>"]
        SSM["State-space models<br/><small>SDE kernels · Kalman · RTS · SpInGP · CVI</small>"]
        QUAD["Quadrature<br/><small>Gauss-Hermite · unscented · Taylor · MC</small>"]
        INF["Inference and ensembles<br/><small>BLR · natural gradient · EnKF</small>"]
    end

    subgraph L2["Layer 2 · Distributions and sugar"]
        direction LR
        DIST["MultivariateNormal<br/>MultivariateNormalPrecision"]
        SUGAR["log_prob · entropy · KL<br/>conditional · Joseph · project"]
        EF["Exponential family<br/><small>natural ↔ expectation · Fisher</small>"]
    end

    subgraph L15["Layer 1.5 · Strategies and preconditioners"]
        direction LR
        STRAT["DenseSolver · CGSolver · BBMM<br/>MINRES · LSMR · SLQ logdet"]
        PRE["Jacobi · Nyström<br/>PartialCholesky · Operator"]
        FRONT["linear_solve<br/><small>the front door</small>"]
    end

    subgraph L1["Layer 1 · Operators"]
        direction LR
        OPS["Kronecker · KroneckerSum · BlockDiag<br/>BlockTriDiag · LowRankUpdate · Toeplitz<br/>Interpolated · Masked · Circulant"]
        TAGS["Structural tags<br/><small>_tags.py</small>"]
    end

    subgraph L0["Layer 0 · Primitives"]
        PRIM["solve · logdet · cholesky · diag · trace<br/>sqrt · inv · eig · svd · root_decomposition"]
        LINALG["Linear-algebra utilities<br/><small>Woodbury · Schur · safe_cholesky · tridiagonal</small>"]
    end

    subgraph EXT["Foundations"]
        direction LR
        LINEAX["lineax<br/><small>solvers · AbstractLinearOperator · tags</small>"]
        MATFREE["matfree<br/><small>Lanczos · SLQ · Hutchinson</small>"]
        EQXDEP["equinox · jaxtyping · einx"]
    end

    L3 --> L2
    L2 --> L15
    L15 --> L1
    L1 --> L0
    L0 --> EXT

Layer	Name	Owns	Entry point
0	Primitives	Pure functions with structural dispatch	solve, logdet, cholesky, …
1	Operators	Structured matrices + tags	Kronecker, BlockDiag, …
1.5	Strategies	How a solve/logdet is computed	DenseSolver, CGSolver, …
2	Distributions	Gaussians over operators, sugar ops, exp-family	MultivariateNormal, …
3	Recipes	Domain sequences wiring layers 0–2	GP · SSM · quadrature · inference

---

Layer 0 --- Primitives

Pure functions that match the equations in papers. Every one takes a lineax.AbstractLinearOperator and returns arrays or operators:

x = gaussx.solve(A, b)  # solve A x = b
ld = gaussx.logdet(A)  # log|det(A)|
L = gaussx.cholesky(A)  # A = L L^T   (lazy, structure-preserving)
d = gaussx.diag(A)  # diagonal entries
t = gaussx.trace(A)  # tr(A)
S = gaussx.sqrt(A)  # S with S S = A
A_inv = gaussx.inv(A)  # lazy A^{-1}

Two properties are worth internalising:

• Lazy over dense. cholesky, sqrt, and inv return operators. The Cholesky of a Kronecker is a Kronecker of Cholesky factors; nothing is materialized until you call .as_matrix().
• The expensive surprise is flagged. cholesky(SumOfKroneckers) is the one path where a structured operator you would reasonably expect to stay structured must materialize the dense covariance, so it raises a DenseFallbackWarning and points you at sqrt(...) / sum_of_kroneckers_sample(...) instead. Ordinary unstructured operators densify quietly --- there was never a fast path to lose.

How dispatch actually works

No registry, no metaclass, no plugin system. Each primitive is a chain of isinstance checks in a single readable function, ending in a dense fallback:

flowchart TD
    A["gaussx.solve(A, b)"] --> ID{"Identity?"}
    ID -- yes --> RID["return b"]
    ID -- no --> IV{"InverseOperator?"}
    IV -- yes --> RIV["one matvec with the original"]
    IV -- no --> DG{"Diagonal / Diagonalized?"}
    DG -- yes --> RDG["divide by the eigenvalues  — O(n) / transform pair"]
    DG -- no --> MK{"Masked with a capacitance?"}
    MK -- yes --> RMK["capacitance solve"]
    MK -- no --> BD{"BlockDiag / Kronecker?"}
    BD -- yes --> RBD["per block / Roth's column lemma"]
    BD -- no --> LR{"LowRankUpdate?"}
    LR -- yes --> RLR["Woodbury identity"]
    LR -- no --> KS{"SumOfKroneckers / KroneckerSum(Sqrt) /<br/>SpectralFunction?"}
    KS -- yes --> RKS["joint eigenbasis"]
    KS -- no --> BT{"Block-tridiagonal / bidiagonal /<br/>SparseOperator?"}
    BT -- yes --> RBT["block-banded substitution / CG"]
    BT -- no --> WR{"Wrapper?<br/><small>Tagged · Mul · Div · Neg · Composed</small>"}
    WR -- yes --> RWR["unwrap, recurse,<br/>fix up the scalar"]
    WR -- no --> AD{"Add of two Kroneckers?"}
    AD -- yes --> RAD["sum-of-Kroneckers eigen reduction"]
    AD -- no --> FB["dense / iterative fallback<br/><small>lineax solver chosen from tags</small>"]

    style FB stroke-dasharray: 5 5

logdet, cholesky, diag, trace, sqrt, and inv follow the same shape with their own fast paths. Coverage differs per primitive --- sqrt, for instance, has no block-tridiagonal path and falls back there. In the table below, dense means the operator has no branch in that primitive and takes the fallback in the last row, and lazy means inv wraps it in an InverseOperator whose matvec is a solve. tests/test_docs_dispatch_table.py checks every row against the isinstance chains in src/gaussx/_primitives/, so a new branch without a row, or a dense / lazy cell that gains a branch, fails the test suite.

Operator	solve	logdet	cholesky	diag / trace	sqrt	inv
Identity / Diagonal	O(n)	O(n)	O(n)	O(n)	O(n)	O(n)
BlockDiag	per block	sum of logdets	per block	per block	per block	per block
Kronecker	Roth's lemma	scaled sum	per factor	per factor	per factor	per factor
KroneckerSum	joint eigenbasis	\(\sum \log(\lambda_i + \mu_j)\)	dense	from the factors' diagonals / traces	KroneckerSumSqrt	lazy
KroneckerSumSqrt	joint eigenbasis	dense	dense	dense	dense	lazy
SumOfKroneckers	two terms: whiten + per-factor eigh; else dense	two terms: eigenvalue sum; else dense	dense (warns)	per term	SumOfKroneckersSqrt	lazy
SpectralFunction	per-axis eigenbasis	\(\sum \log\lvert f(\lambda)\rvert\)	dense	diag: per-axis eigenbasis; trace: dense	dense	lazy
LowRankUpdate	Woodbury	determinant lemma	dense	base + update	orthonormal \(U\) with scaled-identity base: closed form; else dense	square base: Woodbury
BlockTriDiag	symmetric: block-banded; else dense	symmetric: block Cholesky; else dense	block Cholesky	per block	dense	lazy
LowerBlockTriDiag / UpperBlockTriDiag	block forward / back substitution	sum over the diagonal blocks	dense	per block	dense	lazy
DiagonalizedOperator / circulant	transform pair	\(\sum \log\lvert\lambda\rvert\)	dense (warns; use sqrt)	trace: \(\sum \lambda\); diag: \(\bar\lambda\) for the FFT pair, else dense (warns)	eigenvalues	eigenvalues
SparseOperator	large and PSD: CG; else dense	large and PSD: SLQ estimate; else dense	sparse Cholesky	diag: stored diagonal; trace: dense	dense	lazy
MaskedOperator	with a capacitance: capacitance solve; else dense	dense	dense	dense	dense	lazy
Toeplitz	dense	dense	dense	\(c_0\) (constant diagonal)	dense	lazy
InterpolatedOperator	dense	dense	dense	dense	dense	lazy
InverseOperator (from inv)	one matvec with \(A\)	dense	dense	dense	dense	returns \(A\)
Tagged	unwrap + recurse	unwrap + recurse	unwrap + recurse	unwrap + recurse	unwrap + recurse	lazy
Mul / Div / Neg	unwrap + recurse, fix up the scalar	unwrap + recurse, add \(\pm n \log\lvert c\rvert\)	\(\sqrt{c}\) folded into the factor; -A raises	unwrap + recurse, scale	\(\sqrt{c}\) folded into the root; -A raises	unwrap + recurse, scale by \(1/c\)
Composed	\(B^{-1}(A^{-1} b)\)	square factors: sum of logdets; else dense	dense	dense	dense	square factors: \(B^{-1} A^{-1}\); else lazy
Add	two Kronecker terms: eigen reduction; else dense	two Kronecker terms: eigenvalue sum; else dense	dense	sum of the two parts	dense	lazy
Everything else	lineax solver chosen from the tags	slogdet	jax.scipy Cholesky	dense	dense eigh	lazy InverseOperator

Several of the dense cells have an opt-in matrix-free escape hatch backed by matfree:

• diag(A, stochastic=True) and trace(A, stochastic=True) switch to Hutchinson / XTrace probing. trace_and_diag shares a single probe pass between both estimates.
• sqrt(A, lanczos_order=k) returns a lazy SqrtOperator that routes its matvecs through Lanczos instead of an eigendecomposition.

For logdet, the matrix-free route is a Layer 1.5 concern rather than a primitive flag: pass SLQLogdet() (or a ComposedSolver) as the solver=.

---

Layer 1 --- Operators

Structured operators extending lineax.AbstractLinearOperator. Each is an equinox.Module --- an immutable PyTree, safe under jit, grad, and vmap --- supporting mv, as_matrix, transpose, and carrying structural tags.

flowchart LR
    ALO["lineax.AbstractLinearOperator"]

    ALO --> P["Products and sums"]
    ALO --> LRK["Low-rank"]
    ALO --> B["Banded and structured"]
    ALO --> K["Interpolated and masked"]
    ALO --> LZ["Lazy algebra"]

    P --> P1["Kronecker"]
    P --> P2["KroneckerSum"]
    P --> P3["SumOfKroneckers"]
    P --> P4["BlockDiag"]

    LRK --> L1["LowRankUpdate"]
    LRK --> L2["SVDLowRankUpdate"]

    B --> B1["BlockTriDiag"]
    B --> B2["Lower/UpperBlockTriDiag"]
    B --> B3["Toeplitz"]
    B --> B4["ToeplitzCholesky"]

    K --> K1["InterpolatedOperator"]
    K --> K2["MaskedOperator"]

    LZ --> Z1["sum_operator"]
    LZ --> Z2["scaled_operator"]
    LZ --> Z3["product_operator"]

Operator	Represents	Why it is worth it
Kronecker(A, B, ...)	\(A \otimes B \otimes \cdots\)	\(O(\sum_i n_i^3)\) instead of \(O((\prod_i n_i)^3)\); matvec is Roth's column lemma via einx
KroneckerSum(A, B)	\(A \oplus B = A \otimes I + I \otimes B\)	Diagonalises in the joint eigenbasis, eigenvalues \(\lambda_i + \mu_j\)
SumOfKroneckers	\(\sum_k A_k \otimes B_k\)	Structured Cholesky/sqrt for separable-plus-noise covariances. Not the same as KroneckerSum; was named SumKronecker, which survives as a deprecated alias
BlockDiag(A, B, ...)	\(\mathrm{diag}(A, B, \ldots)\)	Every primitive splits per block
BlockTriDiag(D, A)	Symmetric block-tridiagonal precision	\(O(N d^3)\) --- the precision structure of Markovian GPs
LowRankUpdate(L, U, d, V)	\(L + U\,\mathrm{diag}(d)\,V^\top\)	Woodbury solves, determinant-lemma logdets; pass orthonormal=True for SVD/Nyström factors
SVDLowRankUpdate	Deprecated --- use LowRankUpdate(..., orthonormal=True)	Kept as a subclass so isinstance checks keep working; warns on construction
Toeplitz	Symmetric Toeplitz	\(O(n \log n)\) matvecs and sampling via FFT circulant embedding
InterpolatedOperator	\(W K_{uu} W^\top\) (KISS-GP style)	Grid interpolation weights, sparse in effect
MaskedOperator	Row/column sub-selection of a base operator	Irregular domains, held-out indices

Arithmetic composes with lineax's own operators, and the primitives unwrap the composition automatically:

K = gaussx.Kronecker(A, B)
perturbed = K + 0.1 * lx.IdentityLinearOperator(K.in_structure())
gaussx.solve(perturbed, y)  # dispatches through the sum

Kernel operators live in kernellib

KernelOperator, ImplicitKernelOperator and ImplicitCrossKernelOperator moved to kernellib in gaussx 0.2.0, with Nyström / RFF operators, HSIC / MMD, Falkon and EigenPro. They are lineax operators, so every primitive and strategy here works on them unchanged --- as long as their structural tags are set: tags are claims, not inferences.

Structural tags

Tags mark structure and properties; the is_* predicates are what the primitives consult. gaussx re-exports lineax's property tags so user code needs only one import.

Tag	Source	Consulted by
symmetric_tag	lineax	solve, logdet, solver selection
positive_semidefinite_tag	lineax	cholesky, sqrt, CG, SLQ
negative_semidefinite_tag	lineax	linear_solve negation path
diagonal_tag	lineax	all primitives (O(n) paths)
lower_triangular_tag / upper_triangular_tag	lineax	triangular solves
unit_diagonal_tag / tridiagonal_tag	lineax	specialised solves
kronecker_tag	gaussx	all primitives (per factor)
kronecker_sum_tag	gaussx	eigenbasis paths
block_diagonal_tag	gaussx	all primitives (per block)
block_tridiagonal_tag	gaussx	banded Cholesky
low_rank_tag	gaussx	solve (Woodbury), logdet (determinant lemma)

Query with is_kronecker, is_low_rank, is_positive_semidefinite, … to inspect an operator without knowing its concrete type. Full list in the Operators & Tags reference.

---

Layer 1.5 --- Strategies & preconditioners

A strategy bundles a solve and a logdet algorithm behind one object. Anything in gaussx that takes a solver= keyword takes one of these, and solver=None means "use structural dispatch". inv_quad_logdet follows the same rule: with strategy=None it takes the exact structural solve and logdet whenever the operator has them, and only falls back to BBMMSolver's shared CG pass (a stochastic log-determinant) for operators that do not.

flowchart TB
    CALL["caller with solver=..."] --> DISP{"solver is None?"}
    DISP -- yes --> PRIMS["Layer 0 primitives<br/><small>structural dispatch</small>"]
    DISP -- no --> STRAT["AbstractSolverStrategy"]

    STRAT --> DENSE["DenseSolver<br/><small>structural + dense</small>"]
    STRAT --> AUTO["AutoSolver<br/><small>picks by type + size</small>"]
    STRAT --> ITER["CGSolver · PreconditionedCGSolver<br/>MINRESSolver · LSMRSolver · BBMMSolver"]
    STRAT --> COMP["ComposedSolver<br/><small>solve from one, logdet from another</small>"]

    COMP --> ITER
    COMP --> LOGDET["DenseLogdet · SLQLogdet<br/>IndefiniteSLQLogdet"]

    ITER -->|"preconditioner="| PRECOND["AbstractPreconditioner"]
    PRECOND --> PJ["JacobiPreconditioner"]
    PRECOND --> PN["NystromPreconditioner"]
    PRECOND --> PP["PartialCholeskyPreconditioner"]
    PRECOND --> PO["OperatorPreconditioner<br/><small>bring your own M⁻¹</small>"]

    AUTO --> DENSE
    AUTO --> ITER

Strategy	Algorithm	Best for
DenseSolver	Structural dispatch, Cholesky for PSD	Small/medium, structured
AutoSolver	Structured → dense; small → dense; large PSD → CG	When you would rather not choose
CGSolver	Conjugate gradients (+ optional preconditioner)	Large PSD, matrix-free
PreconditionedCGSolver	CG with a required preconditioner	Ill-conditioned kernel matrices
MINRESSolver	MINRES	Symmetric indefinite
LSMRSolver	LSMR	Least squares / rectangular
BBMMSolver	Batched blocked matrix-matrix CG	Many right-hand sides at once
ComposedSolver	Any solve + any logdet	The standard large-kernel recipe: CG solve, SLQ logdet
DenseLogdet	Eigendecomposition	Exactness
SLQLogdet / IndefiniteSLQLogdet	Stochastic Lanczos quadrature	Large PSD / symmetric-indefinite estimates

AutoSolver's rule is deliberately boring and readable: structured operators go to DenseSolver (structural dispatch is already the fast path), dense operators below size_threshold (default 1000) go to DenseSolver, large PSD operators go to CGSolver, everything else falls back to DenseSolver. "Structured" is every operator with an exact structural logdet (including through TaggedLinearOperator, scalar multiples and negation), so a large BlockTriDiag or KroneckerSum stays exact. The large-PSD branch's logdet is a stochastic, fixed-seed SLQ estimate; pass ComposedSolver(CGSolver(), DenseLogdet()) for an exact one.

The front door

linear_solve is the high-level entry point. It accepts a lineax operator or a bare (matvec, in_structure) pair, normalises negative-definite systems, picks an iterative solver from the operator's tags, and threads a preconditioner through. as_linear_operator wraps a raw matvec callable into a tagged FunctionLinearOperator for matrix-free workflows.

AbstractPreconditioner.as_operator(A) receives the system operator, so static preconditioners (Jacobi, or an externally supplied approximate inverse) can ignore it while data-dependent ones (partial Cholesky) build their factor lazily at solve time. This is the seam through which PDE packages inject their own approximate inverses --- a spectral solve or a multigrid V-cycle is wrapped in OperatorPreconditioner and passed in, so gaussx never imports the packages that build them.

---

Layer 2 --- Distributions & sugar

flowchart LR
    subgraph DISTS["Distributions (NumPyro-compatible)"]
        MVN["MultivariateNormal<br/><small>carries Σ as an operator</small>"]
        MVNP["MultivariateNormalPrecision<br/><small>carries Λ = Σ⁻¹ directly</small>"]
    end

    subgraph OPSUGAR["Sugar operations"]
        LP["gaussian_log_prob · gaussian_entropy<br/>quadratic_form · add_jitter"]
        KL["kl_standard_normal<br/>gaussian_kl"]
        COND["conditional · project<br/>joseph_update"]
    end

    subgraph EFAM["Exponential family"]
        NAT["η₁ = Λμ,&nbsp; η₂ = −½Λ"]
        CONV["natural ↔ mean/cov ↔ expectation"]
        FISH["log_partition · fisher_info<br/>sufficient_stats · kl_divergence"]
    end

    MVN --> LP
    MVNP --> LP
    LP --> P0["Layer 0: solve · logdet · cholesky"]
    KL --> P0
    COND --> P0
    NAT --> CONV --> FISH --> P0

• MultivariateNormal(loc, cov_operator, solver=...) takes any covariance operator and any solver strategy. The distribution owns the math; the strategy owns the numerics.
• MultivariateNormalPrecision carries \(\Lambda = \Sigma^{-1}\) directly --- the natural home for natural-parameter guides, where materializing \(\Sigma\) would be wasted work.
• Both require numpyro, which is an optional dependency. The rest of gaussx imports fine without it (see the guarded import at the bottom of src/gaussx/__init__.py).
• Sugar ops evaluate \(\log\mathcal{N}(x \mid \mu, \Sigma)\), entropy, quadratic forms, KLs, conditioning, and the numerically stable Joseph-form covariance update --- all through structured solve + logdet.
• Exponential family provides the conversions (mean/cov ↔ natural ↔ expectation), log-partition, Fisher information, and sufficient statistics that natural-gradient and EP updates are built from.

---

Layer 3 --- Recipes

Layer 3 is thin wiring: domain-specific sequences of Layer 0–2 operations. It is where the library stops being linear algebra and starts being someone's method. Five families, each with its own API page.

flowchart TB
    subgraph R["Layer 3"]
        direction TB
        A["Gaussian processes<br/><small>_gp/</small>"]
        B["State-space models<br/><small>_ssm/</small>"]
        C["Quadrature and moments<br/><small>_quadrature/</small>"]
        D["Inference and ensembles<br/><small>_inference/</small>"]
    end

    A -->|"conditioning · ELBO · whitening"| A1["sparse_conditional · collapsed_elbo<br/>unwhiten · love_variance · oilmm_project"]
    B -->|"filtering · smoothing"| B1["kalman_filter · rts_smoother<br/>parallel_* · spingp_* · cvi_update_sites"]
    C -->|"expectations under a Gaussian"| C1["GaussHermite · Unscented · Taylor · MC<br/>ep_tilted_moments · uncertain_gp_predict"]
    D -->|"posterior updates"| D1["blr_full_update · damped_natural_update<br/>ensemble_kalman_gain · gaspari_cohn"]

Gaussian processes (API) --- posterior conditioning, prediction caches, Matheron pathwise sampling, whitened parameterizations, Gaussian and MC ELBOs, the collapsed (Titsias) bound, LOVE predictive variances, closed-form leave-one-out CV, Kronecker GP marginal likelihood, and OILMM multi-output projections.

State-space models (API) --- stationary 1-D kernels with rational spectral densities have exact SDE representations, turning \(O(N^3)\) GP inference into \(O(N d^3)\) Kalman filtering. Ships the SDE kernel zoo (Matérn, periodic, quasi-periodic, cosine, constant, plus sum/product composition), the sequential filter and RTS smoother, their \(O(\log N)\) parallel associative-scan counterparts (optionally PSD-projected), steady-state (infinite-horizon) filters via the discrete algebraic Riccati equation, SpInGP, and CVI site machinery for non-conjugate likelihoods.

Quadrature & moment matching (API) --- one AbstractIntegrator interface over Gauss-Hermite, unscented/cubature, Taylor, and Monte Carlo rules. Everything needing an expectation (expected log-likelihoods, EP tilted moments, uncertain-input GP prediction, \(\Psi\) statistics) takes an integrator argument, so swapping the rule never touches the model.

Bayesian inference & ensembles (API) --- conjugate BLR updates (full and diagonal), Newton and damped natural-gradient steps, Gauss-Newton curvature, Riemannian PSD correction, ensemble covariances and Kalman gain, ETKF transforms, Gaspari-Cohn localization, and the RTPP/RTPS inflation family.

Two API constraints worth knowing up front:

• kronecker_posterior_predictive(...) needs exact test prior diagonals via K_test_diag_factors= for predictive variances.
• ssm_to_naturals(A, Q, mu_0, P_0) takes the transition noise Q of shape (N-1, d, d) and P_0 separately, as MarkovGaussian does. The older stacked layout (Q[0] == P_0) is deprecated until 0.5.0; with it, an inconsistent initial covariance raises an EquinoxRuntimeError (via equinox.error_if), eagerly and at run time under jax.jit / jax.vmap.

---

Package layout

Every subpackage is private (_-prefixed); the public API is re-exported through src/gaussx/__init__.py, which is the single source of truth for what is supported.

src/gaussx/
├── __init__.py             # Public API (`gaussx.__all__`)
├── _tags.py                # Structural tags + is_* predicates
├── _einx.py                # einx wrappers: rearrange, reduce, repeat, einsum
├── _solve_frontend.py      # linear_solve + as_linear_operator
├── _testing.py             # Test utilities (random PD matrices, assertions)
│
├── _primitives/            # Layer 0 — solve, logdet, cholesky, diag, trace,
│                           #   sqrt, inv, eig, svd, root, frobenius, submatrix
├── _linalg/                # Layer 0 — Woodbury, Schur, safe_cholesky,
│                           #   tridiagonal, Lyapunov, symmetrize
│
├── _operators/             # Layer 1 — Kronecker, KroneckerSum, SumOfKroneckers,
│                           #   BlockDiag, BlockTriDiag, LowRankUpdate, SVD
│                           #   low-rank, Toeplitz, interpolated/masked operators,
│                           #   KISS-GP grids, lazy algebra, capacitance
├── _sparse/                # Layer 1 — symbolic + numeric sparse Cholesky for
│                           #   SparseOperator
├── _gmrf/                  # Layer 1 — GMRF precision builders (RW, AR(1),
│                           #   Besag/BYM2, SPDE Matérn)
│
├── _strategies/            # Layer 1.5 — Dense, Auto, CG, PreconditionedCG,
│                           #   MINRES, LSMR, BBMM, Composed, SLQ logdets
├── _preconditioners/       # Layer 1.5 — Jacobi, Nyström, PartialCholesky,
│                           #   Operator (bring-your-own M⁻¹)
│
├── _distributions/         # Layer 2 — MultivariateNormal(+Precision),
│                           #   log-prob/entropy/KL, conditional, Joseph, project
├── _expfam/                # Layer 2 — natural ↔ expectation, Fisher, partition
│
├── _gp/                    # Layer 3 — conditioning, ELBOs, whitening, LOVE,
│                           #   LOO, Matheron, OILMM, Kronecker GP, caches
├── _ssm/                   # Layer 3 — SDE kernels, Kalman/RTS (sequential,
│                           #   parallel, sqrt, infinite-horizon), SpInGP, CVI
├── _quadrature/            # Layer 3 — integrators, likelihoods, expectations,
│                           #   Ψ-statistics, uncertain-input GP prediction, ADF
├── _inference/             # Layer 3 — BLR, natural gradient, EnKF, localization
│
├── _sketching/             # Outside the stack — random subspace embeddings
└── _randomized/            # Outside the stack — randomized SVD / eigh,
                            #   Nyström, RP-Cholesky

Layer 3 lives in named subpackages (_gp/, _ssm/, …) rather than a single _recipes/ directory --- the families grew large enough to deserve their own namespaces, and each maps one-to-one onto an API reference page.

Import structure

Layers 3 and 2 import only from below. Layers 0, 1 and 1.5 form one dispatch core instead: the primitives isinstance-dispatch on the operator classes (Layer 0 imports Layer 1), some operators call primitives for their own factorisations (SumOfKroneckers → cholesky / sqrt / solve, MaskedOperator → solve), and the primitives fall back to strategy objects (AutoSolver, CGSolver, SLQLogdet) that themselves call solve / logdet. _sparse joins the cycle through the sparse Cholesky path.

gaussx imports its own modules at module scope except on the edges that close one of these cycles. Each such import sits inside the function that needs it, with a # lazy import, cycle: ... comment naming the module chain that imports back, and tests/test_lazy_imports.py keeps the list of those edges explicit, so a new one has to be added deliberately.

---

Dependencies

Package	Role	Required
jax / jaxlib	Array backend	Yes
equinox	Module system, PyTrees	Yes
lineax	Base operators, solvers, property tags	Yes
matfree	Lanczos, SLQ, Hutchinson	Yes
jaxtyping	Array shape annotations	Yes
einx	Tensor reshaping/contraction	Yes
numpyro	MultivariateNormal distributions	Optional (gaussx[numpyro])

---

Where gaussx sits in the ecosystem

gaussx is deliberately a floor, not a building. It owns structured operators and the primitives over them; everything domain-specific is somebody else's library.

flowchart TB
    subgraph UP["Built on"]
        LINEAX["lineax"]
        MATFREE["matfree"]
        EQX["equinox"]
    end

    GX["<b>gaussx</b><br/><small>structured operators · primitives · Gaussians</small>"]

    subgraph PROB["Probabilistic modelling"]
        KL["kernellib<br/><small>kernels, kernel operators, kernel methods</small>"]
        PYROX["pyrox-gp<br/><small>GP models, kernels, guides on NumPyro</small>"]
        FILTERAX["filterax<br/><small>ensemble filtering</small>"]
        OPTAXB["optax_bayes<br/><small>natural-gradient optimizers</small>"]
    end

    subgraph PDE["PDE / geoscience"]
        FVX["finitevolX<br/><small>finite-volume ops, Arakawa C-grids</small>"]
        SDX["spectraldiffx<br/><small>pseudospectral discretization</small>"]
    end

    LINEAX --> GX
    MATFREE --> GX
    EQX --> GX

    GX -->|"LowRankUpdate · solvers · trace_product"| KL
    KL -->|"kernels · kernel operators"| PYROX
    GX -->|"operators · solvers · distributions"| PYROX
    GX -.->|"ensemble covariance · Kalman gain"| FILTERAX
    GX -.->|"Fisher · natural params"| OPTAXB
    GX -->|"CG · Nyström precond · tridiagonal"| FVX
    GX -.->|"capacitance / Woodbury correction<br/>(planned)"| SDX

    click LINEAX "https://github.com/patrick-kidger/lineax" _blank
    click MATFREE "https://github.com/pnkraemer/matfree" _blank
    click EQX "https://github.com/patrick-kidger/equinox" _blank
    click KL "https://github.com/jejjohnson/kernellib" _blank
    click PYROX "https://github.com/jejjohnson/pyrox" _blank
    click FILTERAX "https://github.com/jejjohnson/filterax" _blank
    click OPTAXB "https://github.com/jejjohnson/optax_bayes" _blank
    click FVX "https://github.com/jejjohnson/finitevolX" _blank
    click SDX "https://github.com/jejjohnson/spectraldiffx" _blank

pyrox-gp is the largest consumer. It is the modelling shell gaussx deliberately refuses to be: kernels with hyperparameter priors, NumPyro sample sites, guides, likelihoods, sparse and Markov GP models. It reaches into gaussx for more than forty public symbols --- Kronecker, BlockDiag, the solver strategies, MultivariateNormal(Precision), base_conditional, variational_elbo_gaussian, whitened_svgp_predict, kalman_filter, rts_smoother, ep_tilted_moments, the OILMM projections, and the natural-parameter conversions.

kernellib owns everything with a kernel in it: kernel functions and objects, the kernel operators (KernelOperator, ImplicitKernelOperator, ImplicitCrossKernelOperator), Nyström / RFF low-rank operators, HSIC / CKA / MMD, Falkon and EigenPro. They moved out of gaussx in 0.2.0. kernellib builds on gaussx's LowRankUpdate, solver strategies and trace_product; gaussx never imports kernellib.

finitevolX takes the raw solver substrate: its tridiagonal solves are gaussx.tridiagonal_solve / tridiagonal_solve_batched, and its Nyström preconditioner wraps gaussx.NystromPreconditioner around a gaussx.as_linear_operator view of the (PSD-probed) finite-volume operator.

spectraldiffx is the planned next adopter: it supplies the base spectral solve and the boundary indices, and gaussx returns the reusable capacitance operator that corrects it (gaussx.CapacitanceSolver, already implemented on the gaussx side). That correction is the Woodbury identity --- the same code path that solves a low-rank-plus-diagonal covariance. Likewise filterax and optax_bayes are the intended homes for the ensemble-filtering and natural-gradient-optimizer layers built on the Layer 3 primitives gaussx already ships.

That overlap is not a coincidence. Under the SPDE view, elliptic differential operators and Gaussian precision operators are the same objects, so one tested CG loop, one Nyström preconditioner, and one Woodbury correction can serve a GP marginal likelihood and a masked Poisson solve alike. The Unified Solvers design note works through where that boundary is drawn, and what gaussx will never accept (grids, coordinates, boundary conditions, spectral transforms, multigrid V-cycles --- those stay upstairs).

For the philosophy behind these boundaries, see Vision. For the symbol-level reference, see the API docs.