Linear-Algebra Utilities

Numerically careful building blocks shared by the higher layers: robust factorizations, classical matrix identities in operator form, and batched / matrix-RHS solve helpers.

Robust factorization & hygiene

safe_cholesky retries with geometrically growing diagonal jitter inside a statically-bounded jax.lax.fori_loop (JIT-compatible and reverse-mode differentiable) when a matrix is not numerically positive-definite; symmetrize removes the floating-point asymmetry that accumulates in covariance updates.

Structured linear algebra and Gaussian primitives for JAX.

safe_cholesky(operator: lx.AbstractLinearOperator, *, initial_jitter: float = 1e-08, max_jitter: float = 0.01, max_retries: int = 5, growth_factor: float = 10.0) -> Float[Array, 'N N']

Cholesky decomposition with adaptive jitter for near-singular matrices.

The first attempt routes through gaussx._primitives.cholesky, so structured operators (DiagonalLinearOperator, BlockDiag, Kronecker, BlockTriDiag) keep their structure on the happy path. If the result contains NaNs (the matrix is not numerically positive-definite), retries with geometrically increasing diagonal jitter: cholesky(A + eps * I) where eps starts at initial_jitter and grows by growth_factor each retry, up to max_jitter. Jittering in general destroys structure, so retries operate on the dense matrix form.

The retry loop is a jax.lax.fori_loop with static bounds (max_retries is a Python int), so the function is both JIT-compatible and reverse-mode differentiable — lax.while_loop would forbid the latter. Each iteration is a lax.cond that becomes a no-op once the factor is clean, so a matrix that succeeds on attempt k still pays for exactly k factorisations.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A lineax linear operator whose Cholesky factor is required. Must be square and positive-definite.	required
initial_jitter	float	Starting jitter magnitude added to the diagonal.	1e-08
max_jitter	float	Upper bound on jitter (clamped after growth).	0.01
max_retries	int	Maximum number of jittered retries after the initial attempt. Must be a static Python int — it sets the loop's trip count, so it cannot be a traced value (e.g. a jitted function argument).	5
growth_factor	float	Multiplicative factor applied to jitter each retry.	10.0

Returns:

Name	Type	Description
	Float[Array, 'N N']	Lower-triangular Cholesky factor as a dense array.
	Float[Array, 'N N']	If all attempts fail the result will contain NaNs — this is
intentional	Float[Array, 'N N']	JAX cannot raise exceptions inside jit-traced
	Float[Array, 'N N']	code, so callers should check for NaNs when robustness matters.

Source code in src/gaussx/_linalg/_safe_cholesky.py

def safe_cholesky(
    operator: lx.AbstractLinearOperator,
    *,
    initial_jitter: float = 1e-8,
    max_jitter: float = 1e-2,
    max_retries: int = 5,
    growth_factor: float = 10.0,
) -> Float[Array, "N N"]:
    """Cholesky decomposition with adaptive jitter for near-singular matrices.

    The first attempt routes through `gaussx._primitives.cholesky`, so
    structured operators (``DiagonalLinearOperator``, ``BlockDiag``,
    ``Kronecker``, ``BlockTriDiag``) keep their structure on the happy path.
    If the result contains NaNs (the matrix is not numerically
    positive-definite), retries with geometrically increasing diagonal jitter:
    ``cholesky(A + eps * I)`` where *eps* starts at ``initial_jitter`` and
    grows by ``growth_factor`` each retry, up to ``max_jitter``. Jittering in
    general destroys structure, so retries operate on the dense matrix form.

    The retry loop is a ``jax.lax.fori_loop`` with static bounds
    (``max_retries`` is a Python int), so the function is both
    JIT-compatible and reverse-mode differentiable — ``lax.while_loop``
    would forbid the latter. Each iteration is a ``lax.cond`` that becomes
    a no-op once the factor is clean, so a matrix that succeeds on attempt
    *k* still pays for exactly *k* factorisations.

    Args:
        operator: A lineax linear operator whose Cholesky factor is
            required. Must be square and positive-definite.
        initial_jitter: Starting jitter magnitude added to the diagonal.
        max_jitter: Upper bound on jitter (clamped after growth).
        max_retries: Maximum number of jittered retries after the initial
            attempt. Must be a static Python int — it sets the loop's trip
            count, so it cannot be a traced value (e.g. a jitted function
            argument).
        growth_factor: Multiplicative factor applied to jitter each retry.

    Returns:
        Lower-triangular Cholesky factor as a dense array.
        If all attempts fail the result will contain NaNs — this is
        intentional: JAX cannot raise exceptions inside ``jit``-traced
        code, so callers should check for NaNs when robustness matters.
    """
    # A traced max_retries would silently turn the fori_loop into a
    # dynamic-bound loop, which reverse-mode AD forbids — fail loudly
    # instead (int() on a tracer raises ConcretizationTypeError).
    max_retries = int(max_retries)

    # Structured first attempt — preserves operator structure where possible.
    L0 = _chol_matrix(operator)
    has_nan0 = jnp.any(jnp.isnan(L0))

    # State: (L, jitter, still_bad). The jitter must be a concrete array so
    # both cond branches carry identical avals.
    init_state = (L0, jnp.asarray(initial_jitter, dtype=L0.dtype), has_nan0)

    def _retry(state):
        _, eps, _ = state
        # The dense operator is built only on the (rare) retry branch, so the
        # happy path of a structured operator never materialises it (gh-365).
        A = operator.as_matrix()
        eye = jnp.eye(A.shape[0], dtype=A.dtype)
        jittered = lx.MatrixLinearOperator(A + eps * eye, lx.positive_semidefinite_tag)
        L = _chol_matrix(jittered)
        has_nan = jnp.any(jnp.isnan(L))
        next_eps = jnp.minimum(eps * growth_factor, max_jitter)
        return (L, next_eps, has_nan)

    def _body(_, state):
        _, _, still_bad = state
        return jax.lax.cond(still_bad, _retry, lambda s: s, state)

    L_final, _, _ = jax.lax.fori_loop(0, max_retries, _body, init_state)
    return L_final

symmetrize(mat: Float[Array, '... n n']) -> Float[Array, '... n n']

Symmetrize the trailing two axes: 0.5 * (X + X^T).

Eliminates residual floating-point asymmetry after covariance updates. Works on batched stacks (..., n, n).

Parameters:

Name	Type	Description	Default
mat	Float[Array, '... n n']	Square matrix or batched stack of square matrices.	required

Returns:

Type	Description
Float[Array, '... n n']	The symmetric part of mat.

Source code in src/gaussx/_linalg/_symmetrize.py

def symmetrize(mat: Float[Array, "... n n"]) -> Float[Array, "... n n"]:
    """Symmetrize the trailing two axes: ``0.5 * (X + X^T)``.

    Eliminates residual floating-point asymmetry after covariance
    updates. Works on batched stacks ``(..., n, n)``.

    Args:
        mat: Square matrix or batched stack of square matrices.

    Returns:
        The symmetric part of ``mat``.
    """
    return 0.5 * (mat + jnp.swapaxes(mat, -1, -2))

Matrix identities

Woodbury, Schur complements, and conditional (Schur-complement) variances — the identities behind every Gaussian conditioning step, exposed directly so recipes never re-derive them.

Legacy conditional_variance signature

The pre-#152 three-positional form conditional_variance(base_diag, A_X, S_u) is still accepted until gaussx 0.7.0 — it is detected when the second positional argument is a lineax.AbstractLinearOperator — but it emits a DeprecationWarning and skips the \(K_{XZ}\)-based Schur subtraction, treating its first argument as an already-computed Schur diagonal. Call conditional_variance(K_XX_diag, K_XZ, A_X, S_u=S_u) instead.

Structured linear algebra and Gaussian primitives for JAX.

woodbury_solve(base: lx.AbstractLinearOperator, U: Float[Array, 'N k'], D: Float[Array, ' k'], b: Float[Array, ' N']) -> Float[Array, ' N']

Standalone Woodbury identity solve: (L + U diag(D) U^T)^{-1} b.

Convenience function for cases where the user has the components but doesn't want to construct a LowRankUpdate operator.

Uses the identity:

(L + U D U^T)^{-1} b = L^{-1}b - L^{-1}U C^{-1} U^T L^{-1}b

where C = D^{-1} + U^T L^{-1} U is the (k, k) capacitance matrix.

Parameters:

Name	Type	Description	Default
base	AbstractLinearOperator	Base operator L, shape (N, N).	required
U	Float[Array, 'N k']	Low-rank factor, shape (N, k).	required
D	Float[Array, ' k']	Diagonal scaling, shape (k,).	required
b	Float[Array, ' N']	Right-hand side, shape (N,).	required

Returns:

Type	Description
Float[Array, ' N']	Solution x, shape (N,).

Source code in src/gaussx/_linalg/_woodbury.py

def woodbury_solve(
    base: lx.AbstractLinearOperator,
    U: Float[Array, "N k"],
    D: Float[Array, " k"],
    b: Float[Array, " N"],
) -> Float[Array, " N"]:
    """Standalone Woodbury identity solve: ``(L + U diag(D) U^T)^{-1} b``.

    Convenience function for cases where the user has the components
    but doesn't want to construct a ``LowRankUpdate`` operator.

    Uses the identity:

        (L + U D U^T)^{-1} b = L^{-1}b - L^{-1}U C^{-1} U^T L^{-1}b

    where ``C = D^{-1} + U^T L^{-1} U`` is the ``(k, k)`` capacitance
    matrix.

    Args:
        base: Base operator L, shape ``(N, N)``.
        U: Low-rank factor, shape ``(N, k)``.
        D: Diagonal scaling, shape ``(k,)``.
        b: Right-hand side, shape ``(N,)``.

    Returns:
        Solution x, shape ``(N,)``.
    """

    op = LowRankUpdate(base, U, D)
    return solve(op, b)

schur_complement(K_XX: lx.AbstractLinearOperator, K_XZ: Float[Array, 'N M'], K_ZZ: lx.AbstractLinearOperator) -> LowRankUpdate

Schur complement: K_XX - K_XZ @ K_ZZ^{-1} @ K_ZX.

Central to GP conditional distributions. Returns a LowRankUpdate operator so that downstream operations (solve, logdet) can exploit the low-rank structure.

Parameters:

Name	Type	Description	Default
K_XX	AbstractLinearOperator	Prior covariance, shape (N, N).	required
K_XZ	Float[Array, 'N M']	Cross-covariance, shape (N, M).	required
K_ZZ	AbstractLinearOperator	Inducing covariance, shape (M, M).	required

Returns:

Type	Description
LowRankUpdate	A LowRankUpdate representing K_XX - K_XZ K_ZZ^{-1} K_ZX.

Source code in src/gaussx/_linalg/_schur.py

def schur_complement(
    K_XX: lx.AbstractLinearOperator,
    K_XZ: Float[Array, "N M"],
    K_ZZ: lx.AbstractLinearOperator,
) -> LowRankUpdate:
    """Schur complement: ``K_XX - K_XZ @ K_ZZ^{-1} @ K_ZX``.

    Central to GP conditional distributions. Returns a
    ``LowRankUpdate`` operator so that downstream operations
    (solve, logdet) can exploit the low-rank structure.

    Args:
        K_XX: Prior covariance, shape ``(N, N)``.
        K_XZ: Cross-covariance, shape ``(N, M)``.
        K_ZZ: Inducing covariance, shape ``(M, M)``.

    Returns:
        A ``LowRankUpdate`` representing K_XX - K_XZ K_ZZ^{-1} K_ZX.
    """
    # Solve K_ZZ @ W_j = K_XZ[j, :]^T for each row j of K_XZ
    # W = K_ZZ^{-1} K_XZ^T has shape (M, N)
    # Then Schur = K_XX - K_XZ @ K_ZZ^{-1} @ K_XZ^T
    _N, M = K_XZ.shape

    # Solve K_ZZ w_j = k_xz_j for each row of K_XZ
    # W^T = vmap(solve(K_ZZ, ·))(K_XZ)  =>  (N, M)
    # lazy import, cycle: _linalg._linalg -> _linalg._schur
    from gaussx._linalg._linalg import solve_rows

    W_T = solve_rows(K_ZZ, K_XZ)  # (N, M)
    W = W_T.T  # (M, N)

    # Represent as K_XX + K_XZ @ (-I_M) @ W
    # = K_XX - K_XZ @ K_ZZ^{-1} @ K_XZ^T
    # LowRankUpdate: base + U @ diag(d) @ V^T
    # U = K_XZ (N, M), d = -ones(M), V = W^T (N, M)
    d = -jnp.ones(M, dtype=W.dtype)
    return LowRankUpdate(K_XX, K_XZ, d, W.T)

conditional_variance(K_XX_diag: Float[Array, ' N'], K_XZ: Float[Array, 'N M'] | lx.AbstractLinearOperator | None = None, A_X: Float[Array, 'N M'] | None = None, S_u: lx.AbstractLinearOperator | None = None) -> Float[Array, ' N']

Predictive variance: Schur complement diagonal plus optional variational correction.

Computes the diagonal of the conditional covariance:

diag(K_XX - A_X K_XZ^T) + diag(A_X S_u A_X^T)

where A_X = K_XZ K_ZZ^{-1} is the projection matrix. Negative base variances are clamped to zero.

Without S_u this is the exact GP predictive variance (diagonal of the Schur complement). With S_u it is the sparse GP variational predictive variance that adds a variational covariance correction.

Parameters:

Name	Type	Description	Default
K_XX_diag	Float[Array, ' N']	Prior diagonal variances, shape (N,).	required
K_XZ	Float[Array, 'N M'] | AbstractLinearOperator | None	Cross-covariance matrix, shape (N, M).	None
A_X	Float[Array, 'N M'] | None	Projection matrix K_XZ K_ZZ^{-1}, shape (N, M).	None
S_u	AbstractLinearOperator | None	Optional variational covariance, shape (M, M). When provided, adds diag(A_X S_u A_X^T) to the base variance.	None

Returns:

Type	Description
Float[Array, ' N']	Predictive variances, shape (N,).

Note

Until gaussx 0.7.0 the legacy three-positional-argument form conditional_variance(base_diag, A_X, S_u) (where base_diag was the Schur diagonal already, A_X was the projection, and S_u was the variational covariance) is still accepted: it is detected when the second positional argument is a lineax.AbstractLinearOperator (the old S_u slot type). Such calls emit a DeprecationWarning and compute base_diag + diag(A_X S_u A_X^T) without the Schur subtraction.

Source code in src/gaussx/_linalg/_schur.py

def conditional_variance(
    K_XX_diag: Float[Array, " N"],
    K_XZ: Float[Array, "N M"] | lx.AbstractLinearOperator | None = None,
    A_X: Float[Array, "N M"] | None = None,
    S_u: lx.AbstractLinearOperator | None = None,
) -> Float[Array, " N"]:
    """Predictive variance: Schur complement diagonal plus optional variational
    correction.

    Computes the diagonal of the conditional covariance:

        diag(K_XX - A_X K_XZ^T) + diag(A_X S_u A_X^T)

    where ``A_X = K_XZ K_ZZ^{-1}`` is the projection matrix.
    Negative base variances are clamped to zero.

    Without ``S_u`` this is the exact GP predictive variance (diagonal of the
    Schur complement).  With ``S_u`` it is the sparse GP variational predictive
    variance that adds a variational covariance correction.

    Args:
        K_XX_diag: Prior diagonal variances, shape ``(N,)``.
        K_XZ: Cross-covariance matrix, shape ``(N, M)``.
        A_X: Projection matrix ``K_XZ K_ZZ^{-1}``, shape ``(N, M)``.
        S_u: Optional variational covariance, shape ``(M, M)``.  When
            provided, adds ``diag(A_X S_u A_X^T)`` to the base variance.

    Returns:
        Predictive variances, shape ``(N,)``.

    Note:
        Until gaussx 0.7.0 the legacy three-positional-argument form
        ``conditional_variance(base_diag, A_X, S_u)`` (where
        ``base_diag`` was the *Schur* diagonal already, ``A_X`` was the
        projection, and ``S_u`` was the variational covariance) is
        still accepted: it is detected when the second positional
        argument is a `lineax.AbstractLinearOperator` (the old
        ``S_u`` slot type). Such calls emit a
        `DeprecationWarning` and compute
        ``base_diag + diag(A_X S_u A_X^T)`` without the
        Schur subtraction.
    """
    # Backwards-compat: the legacy 3-positional signature was
    # ``conditional_variance(base_diag, A_X, S_u)``. The pre-#152 docs
    # called it as ``conditional_variance(adjusted_diag, A, S_u)`` with
    # the second arg as the projection matrix and the third as the
    # variational covariance operator. We can detect that pattern by
    # the third positional being an AbstractLinearOperator (it
    # corresponds to the old ``S_u``) while ``S_u`` is left at the
    # default ``None``.
    if (
        S_u is None
        and isinstance(A_X, lx.AbstractLinearOperator)
        and isinstance(K_XZ, jax.Array)
    ):
        warn_deprecated(
            "conditional_variance(base_diag, A_X, S_u) is deprecated and will be "
            "removed in gaussx 0.7.0; "
            "use conditional_variance(K_XX_diag, K_XZ, A_X, S_u=S_u). "
            "The legacy form treats the first argument as the "
            "precomputed Schur diagonal and skips the K_XZ-based "
            "subtraction."
        )
        legacy_A_X = K_XZ  # second positional was the projection matrix
        legacy_S_u: lx.AbstractLinearOperator = A_X  # third was S_u
        S_mat = legacy_S_u.as_matrix()
        AS = legacy_A_X @ S_mat
        diag_ASAt = jnp.sum(AS * legacy_A_X, axis=1)
        return jnp.clip(K_XX_diag, 0.0) + diag_ASAt

    if K_XZ is None or A_X is None:
        raise TypeError(
            "conditional_variance requires K_XZ and A_X (or the legacy "
            f"(base_diag, A_X, S_u) call). Got K_XZ={K_XZ!r}, "
            f"A_X={A_X!r}, S_u={S_u!r}."
        )
    if not isinstance(K_XZ, jax.Array):
        raise TypeError(
            "K_XZ must be a jax array; the legacy 3-positional form "
            "passes an AbstractLinearOperator as the third positional "
            "argument."
        )

    # Diagonal of Schur complement: diag(K_XX - K_XZ K_ZZ^{-1} K_ZX)
    base = jnp.clip(K_XX_diag - jnp.sum(A_X * K_XZ, axis=1), 0.0)
    if S_u is None:
        return base
    # Variational correction: diag(A_X S_u A_X^T)
    S_mat = S_u.as_matrix()
    AS = A_X @ S_mat  # (N, M)
    diag_ASAt = jnp.sum(AS * A_X, axis=1)  # (N,)
    return base + diag_ASAt

diag_conditional_variance(K_XX_diag: Float[Array, ' N'], K_XZ: Float[Array, 'N M'], A_X: Float[Array, 'N M']) -> Float[Array, ' N']

Diagonal of Schur complement: diag(K_XX - A K_ZX).

Thin wrapper around gaussx.conditional_variance (defined in gaussx._linalg._schur.conditional_variance) without a variational covariance. Use gaussx.conditional_variance directly when a variational correction S_u is also needed.

Parameters:

Name	Type	Description	Default
K_XX_diag	Float[Array, ' N']	Prior diagonal variances, shape (N,).	required
K_XZ	Float[Array, 'N M']	Cross-covariance, shape (N, M).	required
A_X	Float[Array, 'N M']	Projection matrix K_XZ K_ZZ^{-1}, shape (N, M).	required

Returns:

Type	Description
Float[Array, ' N']	Conditional variances, shape (N,).

Source code in src/gaussx/_linalg/_linalg.py

def diag_conditional_variance(
    K_XX_diag: Float[Array, " N"],
    K_XZ: Float[Array, "N M"],
    A_X: Float[Array, "N M"],
) -> Float[Array, " N"]:
    """Diagonal of Schur complement: ``diag(K_XX - A K_ZX)``.

    Thin wrapper around `gaussx.conditional_variance` (defined in
    `gaussx._linalg._schur.conditional_variance`) without a
    variational covariance.  Use `gaussx.conditional_variance`
    directly when a variational correction ``S_u`` is also needed.

    Args:
        K_XX_diag: Prior diagonal variances, shape ``(N,)``.
        K_XZ: Cross-covariance, shape ``(N, M)``.
        A_X: Projection matrix ``K_XZ K_ZZ^{-1}``, shape ``(N, M)``.

    Returns:
        Conditional variances, shape ``(N,)``.
    """
    return _conditional_variance(K_XX_diag, K_XZ, A_X)

cov_transform(J: Float[Array, 'M N'] | lx.AbstractLinearOperator, cov_operator: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

cov_transform(J: Float[Array, 'M N'], cov_operator: lx.AbstractLinearOperator) -> lx.MatrixLinearOperator

cov_transform(J: lx.AbstractLinearOperator, cov_operator: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

Covariance propagation through a linear map: J @ Sigma @ J^T.

Used in error propagation, Kalman filter updates, and first-order uncertainty propagation.

Exploits structure where it can:

• Operator-valued J: routes through sandwich, which preserves matched Kronecker / BlockDiag structure and avoids materialising Sigma when either J or cov_operator is diagonal.
• Diagonal cov_operator (dense J): computes (J * d) @ J^T directly, skipping the (N, N) materialization of Sigma.

Otherwise materializes Sigma and forms the dense product. The returned operator is tagged symmetric (and positive-semidefinite when the input is).

Parameters:

Name	Type	Description	Default
J	Float[Array, 'M N'] | AbstractLinearOperator	Jacobian or linear map, shape (M, N) — array or operator.	required
cov_operator	AbstractLinearOperator	Input covariance, shape (N, N).	required

Returns:

Type	Description
AbstractLinearOperator	Transformed covariance operator, shape (M, M). For
AbstractLinearOperator	operator-valued J the structural class of the return type
AbstractLinearOperator	follows sandwich; otherwise it is a
AbstractLinearOperator	lineax.MatrixLinearOperator.

Source code in src/gaussx/_linalg/_linalg.py

def cov_transform(
    J: Float[Array, "M N"] | lx.AbstractLinearOperator,
    cov_operator: lx.AbstractLinearOperator,
) -> lx.AbstractLinearOperator:
    """Covariance propagation through a linear map: ``J @ Sigma @ J^T``.

    Used in error propagation, Kalman filter updates, and
    first-order uncertainty propagation.

    Exploits structure where it can:

    - **Operator-valued** ``J``: routes through `sandwich`, which
      preserves matched ``Kronecker`` / ``BlockDiag`` structure and
      avoids materialising ``Sigma`` when either ``J`` or ``cov_operator``
      is diagonal.
    - **Diagonal** ``cov_operator`` (dense ``J``): computes
      ``(J * d) @ J^T`` directly, skipping the ``(N, N)``
      materialization of ``Sigma``.

    Otherwise materializes ``Sigma`` and forms the dense product. The
    returned operator is tagged symmetric (and positive-semidefinite
    when the input is).

    Args:
        J: Jacobian or linear map, shape ``(M, N)`` — array or operator.
        cov_operator: Input covariance, shape ``(N, N)``.

    Returns:
        Transformed covariance operator, shape ``(M, M)``. For
        operator-valued ``J`` the structural class of the return type
        follows `sandwich`; otherwise it is a
        `lineax.MatrixLinearOperator`.
    """
    if isinstance(J, lx.AbstractLinearOperator):
        return sandwich(J, cov_operator)

    tags = _sandwich_tags(cov_operator)

    if isinstance(cov_operator, lx.DiagonalLinearOperator):
        d = lx.diagonal(cov_operator)
        result = (J * d[None, :]) @ J.T
        return lx.MatrixLinearOperator(result, tags)

    Sigma = cov_operator.as_matrix()
    result = J @ Sigma @ J.T
    return lx.MatrixLinearOperator(result, tags)

sandwich(A: lx.AbstractLinearOperator, P: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

Return A @ P @ A.T exploiting compatible operator structure.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator	Linear map with shape (M, N).	required
P	AbstractLinearOperator	Covariance operator with shape (N, N).	required

Returns:

Type	Description
AbstractLinearOperator	Transformed covariance operator with shape (M, M).

Examples:

A = gaussx.Kronecker(A1, A2)
P = gaussx.Kronecker(P1, P2)
S = gaussx.sandwich(A, P)

Source code in src/gaussx/_linalg/_linalg.py

def sandwich(
    A: lx.AbstractLinearOperator,
    P: lx.AbstractLinearOperator,
) -> lx.AbstractLinearOperator:
    """Return ``A @ P @ A.T`` exploiting compatible operator structure.

    Args:
        A: Linear map with shape ``(M, N)``.
        P: Covariance operator with shape ``(N, N)``.

    Returns:
        Transformed covariance operator with shape ``(M, M)``.

    Examples:
        ```python
        A = gaussx.Kronecker(A1, A2)
        P = gaussx.Kronecker(P1, P2)
        S = gaussx.sandwich(A, P)
        ```
    """
    _check_sandwich_shapes(A, P)
    tags = _sandwich_tags(P)

    if (
        isinstance(A, Kronecker)
        and isinstance(P, Kronecker)
        and len(A.operators) == len(P.operators)
        and all(
            a.in_size() == p.in_size() and p.in_size() == p.out_size()
            for a, p in zip(A.operators, P.operators, strict=True)
        )
    ):
        return Kronecker(
            *(sandwich(a, p) for a, p in zip(A.operators, P.operators, strict=True)),
            tags=tags,
        )

    if (
        isinstance(A, BlockDiag)
        and isinstance(P, BlockDiag)
        and len(A.operators) == len(P.operators)
        and all(
            a.in_size() == p.in_size() and p.in_size() == p.out_size()
            for a, p in zip(A.operators, P.operators, strict=True)
        )
    ):
        return BlockDiag(
            *(sandwich(a, p) for a, p in zip(A.operators, P.operators, strict=True)),
            tags=tags,
        )

    if isinstance(A, lx.DiagonalLinearOperator):
        d = lx.diagonal(A)
        if isinstance(P, lx.DiagonalLinearOperator):
            return lx.TaggedLinearOperator(
                lx.DiagonalLinearOperator(d * lx.diagonal(P) * d),
                tags,
            )
        P_mat = P.as_matrix()
        return lx.MatrixLinearOperator(d[:, None] * P_mat * d[None, :], tags)

    if isinstance(P, lx.DiagonalLinearOperator):
        d = lx.diagonal(P)
        A_mat = A.as_matrix()
        return lx.MatrixLinearOperator((A_mat * d[None, :]) @ A_mat.T, tags)

    A_mat = A.as_matrix()
    P_mat = P.as_matrix()
    return lx.MatrixLinearOperator(A_mat @ P_mat @ A_mat.T, tags)

trace_product(A: lx.AbstractLinearOperator, B: lx.AbstractLinearOperator) -> Float[Array, '']

Trace of a matrix product: tr(A @ B) with structural dispatch.

Uses operator structure where possible to avoid materialization:

• Both diagonal: sum(diag(A) * diag(B)).
• Diagonal × general (or vice versa): contract the diagonal with the diagonal of the other operator (no full materialization).
• Matched gaussx.BlockDiag (same block sizes): sum of per-block trace_product.
• Matched gaussx.Kronecker (same factor structure): prod_i tr(A_i @ B_i).
• Both gaussx.LowRankUpdate (L + U diag(d) V^T): the base-base term recurses, the two cross terms need k matvecs with the other base, and the low-rank-low-rank term costs O(N k_A k_B). For pure low-rank operators (zero-diagonal base, as Nystrom and RFF factorizations produce) this replaces the O(N^2) dense fallback.
• Otherwise falls back to sum(A * B^T) on the materialized matrices — the same O(N²) cost the previous implementation paid.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator	Linear operator, shape (N, N).	required
B	AbstractLinearOperator	Linear operator, shape (N, N).	required

Returns:

Type	Description
Float[Array, '']	Scalar tr(A @ B).

Source code in src/gaussx/_linalg/_linalg.py

def trace_product(
    A: lx.AbstractLinearOperator,
    B: lx.AbstractLinearOperator,
) -> Float[Array, ""]:
    """Trace of a matrix product: ``tr(A @ B)`` with structural dispatch.

    Uses operator structure where possible to avoid materialization:

    - **Both diagonal**: ``sum(diag(A) * diag(B))``.
    - **Diagonal × general** (or vice versa): contract the diagonal with
      the diagonal of the other operator (no full materialization).
    - **Matched** `gaussx.BlockDiag` (same block sizes): sum of
      per-block ``trace_product``.
    - **Matched** `gaussx.Kronecker` (same factor structure):
      ``prod_i tr(A_i @ B_i)``.
    - **Both** `gaussx.LowRankUpdate` (``L + U diag(d) V^T``): the
      base-base term recurses, the two cross terms need ``k`` matvecs with
      the other base, and the low-rank-low-rank term costs
      ``O(N k_A k_B)``. For pure low-rank operators (zero-diagonal base, as
      Nystrom and RFF factorizations produce) this replaces the ``O(N^2)``
      dense fallback.
    - Otherwise falls back to ``sum(A * B^T)`` on the materialized
      matrices — the same O(N²) cost the previous implementation paid.

    Args:
        A: Linear operator, shape ``(N, N)``.
        B: Linear operator, shape ``(N, N)``.

    Returns:
        Scalar ``tr(A @ B)``.
    """

    # Both diagonal: O(N) inner product of diagonals.
    if isinstance(A, lx.DiagonalLinearOperator) and isinstance(
        B, lx.DiagonalLinearOperator
    ):
        return jnp.sum(lx.diagonal(A) * lx.diagonal(B))

    # Diagonal × anything: tr(D @ B) = sum(diag(D) * diag(B)).
    if isinstance(A, lx.DiagonalLinearOperator):
        return jnp.sum(lx.diagonal(A) * diag(B))
    if isinstance(B, lx.DiagonalLinearOperator):
        return jnp.sum(diag(A) * lx.diagonal(B))

    # Matched BlockDiag: tr(blockdiag(A_i) @ blockdiag(B_i)) = sum_i tr(A_i @ B_i).
    if (
        isinstance(A, BlockDiag)
        and isinstance(B, BlockDiag)
        and len(A.operators) == len(B.operators)
        and all(
            a.in_size() == b.in_size()
            for a, b in zip(A.operators, B.operators, strict=True)
        )
    ):
        parts = [
            trace_product(a, b) for a, b in zip(A.operators, B.operators, strict=True)
        ]
        return jnp.sum(jnp.stack(parts))

    # Matched Kronecker: tr((A1⊗A2⊗…) @ (B1⊗B2⊗…)) = prod_i tr(A_i @ B_i).
    if (
        isinstance(A, Kronecker)
        and isinstance(B, Kronecker)
        and len(A.operators) == len(B.operators)
        and all(
            a.in_size() == b.in_size()
            for a, b in zip(A.operators, B.operators, strict=True)
        )
    ):
        parts = [
            trace_product(a, b) for a, b in zip(A.operators, B.operators, strict=True)
        ]
        return jnp.prod(jnp.stack(parts))

    if isinstance(A, LowRankUpdate) and isinstance(B, LowRankUpdate):
        return _trace_product_low_rank(A, B)

    return _trace_product_dense(A, B)

Marginal variances & selected inverses

diag_inv returns \(\operatorname{diag}(Q^{-1})\) — the marginal variances of a Gaussian with precision \(Q\) — and dispatches on structure before falling back to dense Cholesky or Hutchinson:

Structure	Path	Cost
BlockTriDiag (rw1, rw2, ar1, temporal SDE priors)	block Takahashi recursion, selected_inverse	\(O(N d^3)\)
Kronecker \(A \otimes B\)	\(\operatorname{diag}(A^{-1}) \otimes \operatorname{diag}(B^{-1})\)	per factor
KroneckerSum \(A \oplus B\) (grid Laplacians, SPDE on a raster)	\((U_A \circ U_A)\,M\,(U_B \circ U_B)^\top\), \(M_{ij} = 1/(\lambda^A_i + \lambda^B_j)\)	\(O(H^2 W + H W^2)\)
\(A \otimes B + cI\)	same, with \(M_{ij} = 1/(\lambda^A_i \lambda^B_j + c)\)	per factor

pinv=True drops the zero eigenvalues of an intrinsic precision on a grid. A factor that already carries its eigenbasis (DiagonalizedOperator, KroneckerSum) keeps it, so solve, logdet and diag_inv of a space-time \(A \otimes B + cI\) never form the spatial factor \(B\).

# Posterior sd of an AR(1) trend under Gaussian noise: Q + σ⁻² I stays block-tridiagonal
H = gaussx.BlockTriDiag(Q.diagonal + jnp.eye(1) / sigma**2, Q.sub_diagonal)
sd = jnp.sqrt(gaussx.diag_inv(H))  # O(N), not O(N³)

# Prior sd of an intrinsic field on an H × W grid: two small matrix products
sd = jnp.sqrt(gaussx.diag_inv(gaussx.KroneckerSum(L_H, L_W), pinv=True))

For a precision too large to factor (a 10⁶-node mesh, where Cholesky fill is prohibitive), method="xdiag" is the variance-reduced estimator of Epperly, Tropp & Webber (2024), wrapping matfree's leave_one_out_xdiag. It computes the range of \(Q^{-1}\Omega\) exactly and runs Hutchinson only on the remainder, reusing every probe for both. With \(k\) probes it costs \(2k\) solves through solver (Jacobi-preconditioned CG for a precision matrix), and it beats Hutchinson at equal solve counts whenever \(Q^{-1}\) has a few dominant directions, as GMRF covariances do. method="auto" still picks Hutchinson above \(N = 2048\), so XDiag is opt-in.

sd = jnp.sqrt(
    gaussx.diag_inv(
        H,
        method="xdiag",
        num_probes=64,
        key=key,
        solver=gaussx.PreconditionedCGSolver(
            preconditioner=gaussx.JacobiPreconditioner()
        ),
    )
)

Structured linear algebra and Gaussian primitives for JAX.

diag_inv(operator: lx.AbstractLinearOperator, *, method: str = 'auto', num_probes: int = 30, key: jax.Array | None = None, solver: AbstractSolveStrategy | None = None, pinv: bool = False) -> Float[Array, ' N']

Compute the diagonal of the inverse of a linear operator.

Returns diag(A⁻¹) without forming the full inverse matrix — the marginal variances of a Gaussian with precision A.

With method="auto" the operator's structure picks an exact path first:

• DiagonalLinearOperator: 1 / d; BlockDiag: per block, each dispatched in turn; a symmetric LowRankUpdate: the diagonal of its Woodbury inverse, O(N k²); c · A, A / c and -A: the structured path of A, rescaled (gh-365).
• BlockTriDiag (symmetric): the diagonal blocks of gaussx.selected_inverse, O(N d³).
• Kronecker A ⊗ B: diag_inv(A) ⊗ diag_inv(B), each factor dispatched in turn.
• KroneckerSum A ⊕ B (and a DiagonalizedOperator, or a SpectralFunction f(A ⊕ B), with M_ij = 1/f(λ^A_i + λ^B_j)): with A = U_A Λ_A U_Aᵀ, B = U_B Λ_B U_Bᵀ, diag((A ⊕ B)⁻¹) = (U_A ∘ U_A) M (U_B ∘ U_B)ᵀ with M_ij = 1/(λ^A_i + λ^B_j) — two small matrix products, O(H²W + HW²) on an H × W grid. A factor that carries its own eigenbasis keeps it; others are eigendecomposed individually.
• Shifted Kronecker products A ⊗ B + c·I (a SumOfKroneckers or the equivalent lineax sum): the same formula with M_ij = 1/(λ^A_i λ^B_j + c).
• SparseOperator with solver=SparseCholeskySolver(...): Takahashi's selected inverse through the sparse Cholesky factor, exact at O(Σ_j |struct(L_{:,j})|²), at any size.

Anything else uses Cholesky for N ≤ 2048 (the sparse factor for a SparseOperator, dense otherwise) and Hutchinson above.

XDiag (method="xdiag", Epperly, Tropp & Webber, 2024) is the variance-reduced estimator for operators too large to factor (e.g. the Laplace Hessian of a GMRF on a 10⁶-node mesh). For \(B = A^{-1}\), probes \(\Omega = [\omega_1, \dots, \omega_k]\) (random signs; Gaussian when \(N\) is so small that \(k\) sign vectors could be linearly dependent, see below) and \(Q = \operatorname{orth}(B\Omega)\),

\[ \operatorname{diag}(B) = \operatorname{diag}(QQ^\top B) + \operatorname{diag}\big((I - QQ^\top)B\big), \]

the first term exact and the second estimated by Hutchinson, \(\operatorname{diag}(M) \approx \omega \odot M\omega\). XDiag uses every probe for both terms: the \(i\)-th estimate builds \(Q_{-i}\) from the other \(k - 1\) probes and applies Hutchinson with \(\omega_i\), and the result is the mean of the \(k\) leave-one-out estimates. A GMRF covariance has a few large-variance (smooth, large-scale) directions where plain Hutchinson is noisy; the low-rank part captures exactly those. It costs \(2k\) solves (\(B\Omega\) and \(B^\top Q\)), so compare it with Hutchinson at num_probes = 2k. With \(2k \ge N\) it returns the exact diagonal from \(N\) solves instead.

Ω = signs(key, N × k)
Y = A⁻¹ Ω;  Q, R = qr(Y);  Z = A⁻ᵀ Q          # 2k solves
for i in 1..k:                                # in closed form, vectorised
    dᵢ = diag(Q₋ᵢ Q₋ᵢᵀ B) + ωᵢ ⊙ (I − Q₋ᵢ Q₋ᵢᵀ) B ωᵢ
return mean(dᵢ)

It wraps matfree's stochtrace.leave_one_out_xdiag, which reads a rank-deficient \(A^{-1}\Omega\) as a low-rank \(A^{-1}\) and returns the projected diagonal as exact. \(k\) random-sign vectors in \(N\) dimensions are dependent with probability about \(k^2 2^{-N}\) (two equal or opposite columns), so signs, which estimate a diagonal with less variance, are used only when that is below \(2^{-60}\); smaller problems use Gaussian probes, dependent with probability zero. The solves go through solver (any solve strategy; for a precision matrix, Jacobi-preconditioned CG, since the Nyström and RPCholesky preconditioners are for covariance form). Prefer an exact path (the structured ones above, or Takahashi via SparseCholeskySolver) whenever the operator factors; XDiag beats Hutchinson whenever \(A^{-1}\)'s spectrum decays, and matches it when the spectrum is flat.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A linear operator representing A.	required
method	str	Algorithm to use. One of "cholesky" (exact via dense Cholesky), "solve" (exact via repeated solves), "hutchinson" (stochastic estimator), "xdiag" (the variance-reduced estimator above), or "auto" (the structured paths above; otherwise cholesky for N ≤ 2048, hutchinson above). "cholesky" on a SparseOperator is the sparse factor's Takahashi sweep.	'auto'
num_probes	int	Number of Rademacher probe vectors for the hutchinson and xdiag methods (xdiag: \(1 \le k \le N\), and \(2k\) solves).	30
key	Array | None	PRNG key for probe generation in the hutchinson and xdiag methods. When None, defaults to jax.random.PRNGKey(0).	None
solver	AbstractSolveStrategy | None	Optional solve strategy for the "solve", "hutchinson" and "xdiag" methods.	None
pinv	bool	Return the diagonal of the pseudo-inverse instead: eigenvalues below max(n_k) · eps · max|λ| contribute nothing. For intrinsic (singular) precisions on grids, e.g. the exact ICAR / BYM2 scaling constant on a raster. Only the eigenvalue-based paths (Kronecker sums and products of such operators, shifted Kronecker products, DiagonalizedOperator, SpectralFunction) support it.	False

Returns:

Type	Description
Float[Array, ' N']	1D array of shape (N,) with the diagonal entries of A⁻¹.

Raises:

Type	Description
ValueError	For an unknown method, pinv=True on an operator without an eigenvalue-based path, or method="xdiag" with num_probes outside [1, N].

References

Epperly, E. N., Tropp, J. A. & Webber, R. J. (2024). XTrace: making the most of every sample in stochastic trace estimation. SIAM J. Matrix Anal. Appl., 45(1), 1-23.

Examples:

import jax.numpy as jnp
import lineax as lx
import gaussx

# Prior sd of an intrinsic field on a 4 x 5 grid: L_H ⊕ L_W is
# singular (constants), so take the pseudo-inverse.
def path_laplacian(n):
    off = -jnp.ones(n - 1)
    L = jnp.diag(jnp.r_[1.0, 2.0 * jnp.ones(n - 2), 1.0])
    L = L + jnp.diag(off, 1) + jnp.diag(off, -1)
    return lx.MatrixLinearOperator(L, lx.symmetric_tag)

Q = gaussx.KroneckerSum(path_laplacian(4), path_laplacian(5))
variances = gaussx.diag_inv(Q, pinv=True)
dense = jnp.diag(jnp.linalg.pinv(Q.as_matrix()))
assert jnp.allclose(variances, dense, atol=1e-5)

Marginal sd of a precision too large to factor (e.g. a 2M-node global mesh; H a SparseOperator tagged lx.positive_semidefinite_tag, which CG needs), with Jacobi-preconditioned CG for the solves:

sd = jnp.sqrt(
    gaussx.diag_inv(
        H,
        method="xdiag",
        num_probes=64,
        key=key,
        solver=gaussx.PreconditionedCGSolver(
            preconditioner=gaussx.JacobiPreconditioner()
        ),
    )
)

Source code in src/gaussx/_linalg/_diag_inv.py

def diag_inv(
    operator: lx.AbstractLinearOperator,
    *,
    method: str = "auto",
    num_probes: int = 30,
    key: jax.Array | None = None,
    solver: AbstractSolveStrategy | None = None,
    pinv: bool = False,
) -> Float[Array, " N"]:
    r"""Compute the diagonal of the inverse of a linear operator.

    Returns ``diag(A⁻¹)`` without forming the full inverse matrix — the
    marginal variances of a Gaussian with precision ``A``.

    With ``method="auto"`` the operator's structure picks an exact path
    first:

    - ``DiagonalLinearOperator``: ``1 / d``; `BlockDiag`: per block, each
      dispatched in turn; a symmetric `LowRankUpdate`: the diagonal of its
      Woodbury inverse, ``O(N k²)``; ``c · A``, ``A / c`` and ``-A``: the
      structured path of ``A``, rescaled (gh-365).
    - `BlockTriDiag` (symmetric): the diagonal blocks of
      `gaussx.selected_inverse`, ``O(N d³)``.
    - `Kronecker` ``A ⊗ B``: ``diag_inv(A) ⊗ diag_inv(B)``, each factor
      dispatched in turn.
    - `KroneckerSum` ``A ⊕ B`` (and a `DiagonalizedOperator`, or a
      `SpectralFunction` ``f(A ⊕ B)``, with ``M_ij = 1/f(λ^A_i + λ^B_j)``): with
      ``A = U_A Λ_A U_Aᵀ``, ``B = U_B Λ_B U_Bᵀ``,
      ``diag((A ⊕ B)⁻¹) = (U_A ∘ U_A) M (U_B ∘ U_B)ᵀ`` with
      ``M_ij = 1/(λ^A_i + λ^B_j)`` — two small matrix products,
      ``O(H²W + HW²)`` on an ``H × W`` grid. A factor that carries its own
      eigenbasis keeps it; others are eigendecomposed individually.
    - Shifted Kronecker products ``A ⊗ B + c·I`` (a `SumOfKroneckers` or the
      equivalent lineax sum): the same formula with
      ``M_ij = 1/(λ^A_i λ^B_j + c)``.
    - `SparseOperator` with ``solver=SparseCholeskySolver(...)``: Takahashi's
      selected inverse through the sparse Cholesky factor, exact at
      ``O(Σ_j |struct(L_{:,j})|²)``, at any size.

    Anything else uses Cholesky for ``N ≤ 2048`` (the sparse factor for a
    `SparseOperator`, dense otherwise) and Hutchinson above.

    **XDiag** (``method="xdiag"``, Epperly, Tropp & Webber, 2024) is the
    variance-reduced estimator for operators too large to factor (e.g. the
    Laplace Hessian of a GMRF on a 10⁶-node mesh). For $B = A^{-1}$, probes
    $\Omega = [\omega_1, \dots, \omega_k]$ (random signs; Gaussian when
    $N$ is so small that $k$ sign vectors could be linearly dependent, see
    below) and
    $Q = \operatorname{orth}(B\Omega)$,

    $$
    \operatorname{diag}(B) = \operatorname{diag}(QQ^\top B)
    + \operatorname{diag}\big((I - QQ^\top)B\big),
    $$

    the first term exact and the second estimated by Hutchinson,
    $\operatorname{diag}(M) \approx \omega \odot M\omega$. XDiag uses every
    probe for both terms: the $i$-th estimate builds $Q_{-i}$ from the
    other $k - 1$ probes and applies Hutchinson with $\omega_i$, and the
    result is the mean of the $k$ leave-one-out estimates. A GMRF covariance
    has a few large-variance (smooth, large-scale) directions where plain
    Hutchinson is noisy; the low-rank part captures exactly those. It costs
    $2k$ solves ($B\Omega$ and $B^\top Q$), so compare it with Hutchinson
    at ``num_probes = 2k``. With $2k \ge N$ it returns the exact diagonal
    from $N$ solves instead.

    ```text
    Ω = signs(key, N × k)
    Y = A⁻¹ Ω;  Q, R = qr(Y);  Z = A⁻ᵀ Q          # 2k solves
    for i in 1..k:                                # in closed form, vectorised
        dᵢ = diag(Q₋ᵢ Q₋ᵢᵀ B) + ωᵢ ⊙ (I − Q₋ᵢ Q₋ᵢᵀ) B ωᵢ
    return mean(dᵢ)
    ```

    It wraps matfree's ``stochtrace.leave_one_out_xdiag``, which reads a
    rank-deficient $A^{-1}\Omega$ as a low-rank $A^{-1}$ and returns the
    projected diagonal as exact. $k$ random-sign vectors in $N$ dimensions
    are dependent with probability about $k^2 2^{-N}$ (two equal or opposite
    columns), so signs, which estimate a diagonal with less variance, are
    used only when that is below $2^{-60}$; smaller problems use Gaussian
    probes, dependent with probability zero. The solves go
    through ``solver`` (any solve strategy; for a precision matrix,
    Jacobi-preconditioned CG, since the Nyström and RPCholesky
    preconditioners are for covariance form). Prefer an exact path (the
    structured ones above, or Takahashi via ``SparseCholeskySolver``)
    whenever the operator factors; XDiag beats Hutchinson whenever
    $A^{-1}$'s spectrum decays, and matches it when the spectrum is flat.

    Args:
        operator: A linear operator representing A.
        method: Algorithm to use. One of ``"cholesky"`` (exact via
            dense Cholesky), ``"solve"`` (exact via repeated solves),
            ``"hutchinson"`` (stochastic estimator), ``"xdiag"`` (the
            variance-reduced estimator above), or ``"auto"`` (the
            structured paths above; otherwise cholesky for N ≤ 2048,
            hutchinson above). ``"cholesky"`` on a
            `SparseOperator` is the sparse factor's Takahashi sweep.
        num_probes: Number of Rademacher probe vectors for the
            hutchinson and xdiag methods (xdiag: $1 \le k \le N$, and
            $2k$ solves).
        key: PRNG key for probe generation in the hutchinson and xdiag
            methods. When ``None``, defaults to ``jax.random.PRNGKey(0)``.
        solver: Optional solve strategy for the ``"solve"``,
            ``"hutchinson"`` and ``"xdiag"`` methods.
        pinv: Return the diagonal of the pseudo-inverse instead: eigenvalues
            below ``max(n_k) · eps · max|λ|`` contribute nothing. For
            intrinsic (singular) precisions on grids, e.g. the exact ICAR /
            BYM2 scaling constant on a raster. Only the eigenvalue-based
            paths (Kronecker sums and products of such operators, shifted
            Kronecker products, `DiagonalizedOperator`, `SpectralFunction`)
            support it.

    Returns:
        1D array of shape ``(N,)`` with the diagonal entries of A⁻¹.

    Raises:
        ValueError: For an unknown ``method``, ``pinv=True`` on an
            operator without an eigenvalue-based path, or ``method="xdiag"``
            with ``num_probes`` outside ``[1, N]``.

    References:
        Epperly, E. N., Tropp, J. A. & Webber, R. J. (2024). XTrace: making
        the most of every sample in stochastic trace estimation. *SIAM J.
        Matrix Anal. Appl.*, 45(1), 1-23.

    Examples:
        ```python
        import jax.numpy as jnp
        import lineax as lx
        import gaussx

        # Prior sd of an intrinsic field on a 4 x 5 grid: L_H ⊕ L_W is
        # singular (constants), so take the pseudo-inverse.
        def path_laplacian(n):
            off = -jnp.ones(n - 1)
            L = jnp.diag(jnp.r_[1.0, 2.0 * jnp.ones(n - 2), 1.0])
            L = L + jnp.diag(off, 1) + jnp.diag(off, -1)
            return lx.MatrixLinearOperator(L, lx.symmetric_tag)

        Q = gaussx.KroneckerSum(path_laplacian(4), path_laplacian(5))
        variances = gaussx.diag_inv(Q, pinv=True)
        dense = jnp.diag(jnp.linalg.pinv(Q.as_matrix()))
        assert jnp.allclose(variances, dense, atol=1e-5)
        ```

        Marginal sd of a precision too large to factor (e.g. a 2M-node
        global mesh; ``H`` a `SparseOperator` tagged
        ``lx.positive_semidefinite_tag``, which CG needs), with
        Jacobi-preconditioned CG for the solves:

        ```python
        sd = jnp.sqrt(
            gaussx.diag_inv(
                H,
                method="xdiag",
                num_probes=64,
                key=key,
                solver=gaussx.PreconditionedCGSolver(
                    preconditioner=gaussx.JacobiPreconditioner()
                ),
            )
        )
        ```
    """
    n = operator.in_size()

    if method == "auto":
        structured = _diag_inv_structured(
            operator, pinv=pinv, num_probes=num_probes, key=key, solver=solver
        )
        if structured is not None:
            return structured
        method = "cholesky" if n <= 2048 else "hutchinson"

    if pinv:
        msg = (
            "pinv=True needs an eigenvalue-based structure (KroneckerSum, "
            "Kronecker, a shifted Kronecker product or DiagonalizedOperator) "
            f"and method='auto'; got {type(operator).__name__} with "
            f"method={method!r}."
        )
        raise ValueError(msg)

    if method == "cholesky":
        if isinstance(operator, SparseOperator):
            return _sparse_factor(operator, solver).diag_inv()
        return _diag_inv_cholesky(operator)
    if method == "solve":
        return _diag_inv_solve(operator, solver=solver)
    if method == "hutchinson":
        return _diag_inv_hutchinson(
            operator, num_probes=num_probes, key=key, solver=solver
        )
    if method == "xdiag":
        return _diag_inv_xdiag(operator, num_probes=num_probes, key=key, solver=solver)

    msg = (
        f"Unknown method {method!r}; expected 'cholesky', 'solve', "
        "'hutchinson', 'xdiag', or 'auto'."
    )
    raise ValueError(msg)

selected_inverse(operator: BlockTriDiag) -> BlockTriDiag

The band of Σ = Q⁻¹ for a symmetric positive-definite BlockTriDiag.

Takahashi's recursion in block form — the RTS smoother written in precision form. With the block Cholesky factor Q = L Lᵀ (diagonal blocks L_k, sub-diagonal blocks B_k at (k+1, k)), the identity Lᵀ Σ = L⁻¹ gives, sweeping backwards from Σ_NN = L_N⁻ᵀ L_N⁻¹,

\[ \Sigma_{k+1,k} = -\Sigma_{k+1,k+1} B_k L_k^{-1},\qquad \Sigma_{kk} = L_k^{-\top}\big(L_k^{-1} - B_k^\top \Sigma_{k+1,k}\big). \]

Only the diagonal and first off-diagonal blocks are computed, at O(N d³) cost in a lax.scan; with d = 1 this is Rue & Held's scalar recursion. gaussx.diag_inv takes the marginal variances from the diagonal blocks; the off-diagonal blocks are the lag-one covariances.

Parameters:

Name	Type	Description	Default
operator	BlockTriDiag	A symmetric positive-definite block-tridiagonal precision.	required

Returns:

Type	Description
BlockTriDiag	A BlockTriDiag holding Σ_kk and Σ_{k+1,k}. It is the band of
BlockTriDiag	Q⁻¹, not an operator equal to Q⁻¹ (whose blocks are dense).

Raises:

Type	Description
TypeError	If operator is not a BlockTriDiag.
ValueError	If operator is not symmetric.

Examples:

import jax.numpy as jnp
import gaussx

# An AR(1) precision with unit innovation variance and phi = 0.5
phi, n = 0.5, 6
diagonal = jnp.full((n, 1, 1), 1.0 + phi**2).at[-1].set(1.0)
sub_diagonal = jnp.full((n - 1, 1, 1), -phi)
Q = gaussx.BlockTriDiag(diagonal, sub_diagonal)
band = gaussx.selected_inverse(Q)
Sigma = jnp.linalg.inv(Q.as_matrix())
assert jnp.allclose(band.diagonal[:, 0, 0], jnp.diag(Sigma), atol=1e-5)
assert jnp.allclose(band.sub_diagonal[:, 0, 0], jnp.diag(Sigma, -1), atol=1e-5)

Source code in src/gaussx/_linalg/_selected_inverse.py

def selected_inverse(operator: BlockTriDiag) -> BlockTriDiag:
    r"""The band of ``Σ = Q⁻¹`` for a symmetric positive-definite `BlockTriDiag`.

    Takahashi's recursion in block form — the RTS smoother written in
    precision form. With the block Cholesky factor ``Q = L Lᵀ`` (diagonal
    blocks ``L_k``, sub-diagonal blocks ``B_k`` at ``(k+1, k)``), the identity
    ``Lᵀ Σ = L⁻¹`` gives, sweeping backwards from
    ``Σ_NN = L_N⁻ᵀ L_N⁻¹``,

    $$
    \Sigma_{k+1,k} = -\Sigma_{k+1,k+1} B_k L_k^{-1},\qquad
    \Sigma_{kk} = L_k^{-\top}\big(L_k^{-1} - B_k^\top \Sigma_{k+1,k}\big).
    $$

    Only the diagonal and first off-diagonal blocks are computed, at
    ``O(N d³)`` cost in a ``lax.scan``; with ``d = 1`` this is Rue & Held's
    scalar recursion. `gaussx.diag_inv` takes the marginal variances from the
    diagonal blocks; the off-diagonal blocks are the lag-one covariances.

    Args:
        operator: A symmetric positive-definite block-tridiagonal precision.

    Returns:
        A `BlockTriDiag` holding ``Σ_kk`` and ``Σ_{k+1,k}``. It is the band of
        ``Q⁻¹``, not an operator equal to ``Q⁻¹`` (whose blocks are dense).

    Raises:
        TypeError: If ``operator`` is not a `BlockTriDiag`.
        ValueError: If ``operator`` is not symmetric.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx

        # An AR(1) precision with unit innovation variance and phi = 0.5
        phi, n = 0.5, 6
        diagonal = jnp.full((n, 1, 1), 1.0 + phi**2).at[-1].set(1.0)
        sub_diagonal = jnp.full((n - 1, 1, 1), -phi)
        Q = gaussx.BlockTriDiag(diagonal, sub_diagonal)
        band = gaussx.selected_inverse(Q)
        Sigma = jnp.linalg.inv(Q.as_matrix())
        assert jnp.allclose(band.diagonal[:, 0, 0], jnp.diag(Sigma), atol=1e-5)
        assert jnp.allclose(band.sub_diagonal[:, 0, 0], jnp.diag(Sigma, -1), atol=1e-5)
        ```
    """
    if not isinstance(operator, BlockTriDiag):
        msg = (
            "selected_inverse supports BlockTriDiag operators, got "
            f"{type(operator).__name__}."
        )
        raise TypeError(msg)
    if not operator.symmetric:
        msg = (
            "selected_inverse needs a symmetric positive-definite BlockTriDiag "
            "(it runs through the block Cholesky factor)."
        )
        raise ValueError(msg)

    return _block_takahashi(_cholesky_block_tridiag(operator))

Matrix-RHS & batched solves

Solve \(AX = B\) for matrix right-hand sides with one factorization (solve_matrix), per-column or per-row structured dispatch (solve_columns / solve_rows), or \(O(n)\) Thomas-algorithm tridiagonal solves.

Structured linear algebra and Gaussian primitives for JAX.

solve_matrix(operator: lx.AbstractLinearOperator, matrix: Float[Array, 'N K'], *, solver: AbstractSolveStrategy | None = None) -> Float[Array, 'N K']

Solve A X = B with a single factorization on the matrix RHS.

When solver is None and A is a positive semidefinite operator without a structured solve (dense, possibly behind Tagged / scalar wrappers, or a composition or sum), this factors A = L L^T once via gaussx.cholesky and then uses a single cho_solve on the full matrix RHS — avoiding the per-column re-factorization incurred by solve_columns.

Structured operators (LowRankUpdate, Kronecker, BlockDiag, ...), non-PSD operators and a custom solver go to solve_columns, whose per-column structured gaussx.solve (Woodbury, Roth's lemma, per block) is cheaper than densifying a Cholesky factor (gh-356).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator A, shape (N, N).	required
matrix	Float[Array, 'N K']	Right-hand side B, shape (N, K).	required
solver	AbstractSolveStrategy | None	Optional solver strategy. When provided, dispatch is delegated column-by-column via solve_columns.	None

Returns:

Type	Description
Float[Array, 'N K']	Solution X = A⁻¹ B, shape (N, K).

Source code in src/gaussx/_linalg/_linalg.py

def solve_matrix(
    operator: lx.AbstractLinearOperator,
    matrix: Float[Array, "N K"],
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, "N K"]:
    """Solve ``A X = B`` with a single factorization on the matrix RHS.

    When ``solver`` is ``None`` and ``A`` is a positive semidefinite
    operator without a structured solve (dense, possibly behind ``Tagged`` /
    scalar wrappers, or a composition or sum), this factors
    ``A = L L^T`` once via `gaussx.cholesky` and then uses a single
    ``cho_solve`` on the full matrix RHS — avoiding the per-column
    re-factorization incurred by `solve_columns`.

    Structured operators (``LowRankUpdate``, ``Kronecker``, ``BlockDiag``,
    ...), non-PSD operators and a custom ``solver`` go to `solve_columns`,
    whose per-column structured `gaussx.solve` (Woodbury, Roth's lemma,
    per block) is cheaper than densifying a Cholesky factor (gh-356).

    Args:
        operator: Linear operator A, shape ``(N, N)``.
        matrix: Right-hand side B, shape ``(N, K)``.
        solver: Optional solver strategy. When provided, dispatch is
            delegated column-by-column via `solve_columns`.

    Returns:
        Solution X = A⁻¹ B, shape ``(N, K)``.
    """
    if (
        solver is None
        and lx.is_positive_semidefinite(operator)
        and not _has_structured_solve(operator)
    ):
        L = cholesky(operator).as_matrix()
        return jax.scipy.linalg.cho_solve((L, True), matrix)
    return solve_columns(operator, matrix, solver=solver)

solve_columns(operator: lx.AbstractLinearOperator, matrix: Float[Array, 'N K'], *, solver: AbstractSolveStrategy | None = None) -> Float[Array, 'N K']

Solve A X = B column-by-column via vmap.

Equivalent to A⁻¹ @ B but uses structured dispatch per column. Use this when A is a lineax operator with efficient per-vector solve (e.g., Cholesky, diagonal, block-diagonal).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator A, shape (N, N).	required
matrix	Float[Array, 'N K']	Right-hand side B, shape (N, K).	required
solver	AbstractSolveStrategy | None	Optional solver strategy.	None

Returns:

Type	Description
Float[Array, 'N K']	Solution X = A⁻¹ B, shape (N, K).

Source code in src/gaussx/_linalg/_linalg.py

def solve_columns(
    operator: lx.AbstractLinearOperator,
    matrix: Float[Array, "N K"],
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, "N K"]:
    """Solve A X = B column-by-column via vmap.

    Equivalent to ``A⁻¹ @ B`` but uses structured dispatch per column.
    Use this when ``A`` is a lineax operator with efficient per-vector
    solve (e.g., Cholesky, diagonal, block-diagonal).

    Args:
        operator: Linear operator A, shape ``(N, N)``.
        matrix: Right-hand side B, shape ``(N, K)``.
        solver: Optional solver strategy.

    Returns:
        Solution X = A⁻¹ B, shape ``(N, K)``.
    """
    if solver is not None:
        return jax.vmap(
            lambda col: dispatch_solve(operator, col, solver),
            in_axes=1,
            out_axes=1,
        )(matrix)
    return jax.vmap(
        lambda col: solve(operator, col),
        in_axes=1,
        out_axes=1,
    )(matrix)

solve_rows(operator: lx.AbstractLinearOperator, matrix: Float[Array, 'K N'], *, solver: AbstractSolveStrategy | None = None) -> Float[Array, 'K N']

Solve A x = bᵢ for each row bᵢ of a matrix via vmap.

Equivalent to B @ A⁻¹ row-by-row. Used in Kalman gain, Schur complement, and GP prediction computations.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator A, shape (N, N).	required
matrix	Float[Array, 'K N']	Rows of right-hand sides, shape (K, N).	required
solver	AbstractSolveStrategy | None	Optional solver strategy.	None

Returns:

Type	Description
Float[Array, 'K N']	Solutions, shape (K, N).

Source code in src/gaussx/_linalg/_linalg.py

def solve_rows(
    operator: lx.AbstractLinearOperator,
    matrix: Float[Array, "K N"],
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, "K N"]:
    """Solve A x = bᵢ for each row bᵢ of a matrix via vmap.

    Equivalent to ``B @ A⁻¹`` row-by-row. Used in Kalman gain, Schur
    complement, and GP prediction computations.

    Args:
        operator: Linear operator A, shape ``(N, N)``.
        matrix: Rows of right-hand sides, shape ``(K, N)``.
        solver: Optional solver strategy.

    Returns:
        Solutions, shape ``(K, N)``.
    """
    if solver is not None:
        return jax.vmap(
            lambda row: dispatch_solve(operator, row, solver),
        )(matrix)
    return jax.vmap(lambda row: solve(operator, row))(matrix)

tridiagonal_solve(lower: Float[Array, ' n_minus_1'], diag: Float[Array, ' n'], upper: Float[Array, ' n_minus_1'], rhs: Float[Array, ' n']) -> Float[Array, ' n']

Solve a tridiagonal system A x = d.

Thin wrapper over lineax.TridiagonalLinearOperator, which delegates to jax.lax.linalg.tridiagonal_solve (LAPACK / cuSPARSE).

Parameters:

Name	Type	Description	Default
lower	Float[Array, ' n_minus_1']	Sub-diagonal of A, length n - 1.	required
diag	Float[Array, ' n']	Main diagonal of A, length n.	required
upper	Float[Array, ' n_minus_1']	Super-diagonal of A, length n - 1.	required
rhs	Float[Array, ' n']	Right-hand side, length n.	required

Returns:

Type	Description
Float[Array, ' n']	The solution x, length n.

Source code in src/gaussx/_linalg/_tridiagonal.py

def tridiagonal_solve(
    lower: Float[Array, " n_minus_1"],
    diag: Float[Array, " n"],
    upper: Float[Array, " n_minus_1"],
    rhs: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve a tridiagonal system ``A x = d``.

    Thin wrapper over `lineax.TridiagonalLinearOperator`, which delegates
    to ``jax.lax.linalg.tridiagonal_solve`` (LAPACK / cuSPARSE).

    Args:
        lower: Sub-diagonal of ``A``, length ``n - 1``.
        diag: Main diagonal of ``A``, length ``n``.
        upper: Super-diagonal of ``A``, length ``n - 1``.
        rhs: Right-hand side, length ``n``.

    Returns:
        The solution ``x``, length ``n``.
    """
    operator = lx.TridiagonalLinearOperator(diag, lower, upper)
    return lx.linear_solve(operator, rhs, solver=lx.Tridiagonal()).value

tridiagonal_solve_batched(lower: Float[Array, '*batch n_minus_1'], diag: Float[Array, '*batch n'], upper: Float[Array, '*batch n_minus_1'], rhs: Float[Array, '*batch n']) -> Float[Array, '*batch n']

Solve independent tridiagonal systems over leading batch dimensions.

Applies tridiagonal_solve vmapped over all leading dimensions.

Parameters:

Name	Type	Description	Default
lower	Float[Array, '*batch n_minus_1']	Sub-diagonals, shape (*batch, n - 1).	required
diag	Float[Array, '*batch n']	Main diagonals, shape (*batch, n).	required
upper	Float[Array, '*batch n_minus_1']	Super-diagonals, shape (*batch, n - 1).	required
rhs	Float[Array, '*batch n']	Right-hand sides, shape (*batch, n).	required

Returns:

Type	Description
Float[Array, '*batch n']	Solutions, shape (*batch, n).

Source code in src/gaussx/_linalg/_tridiagonal.py

def tridiagonal_solve_batched(
    lower: Float[Array, "*batch n_minus_1"],
    diag: Float[Array, "*batch n"],
    upper: Float[Array, "*batch n_minus_1"],
    rhs: Float[Array, "*batch n"],
) -> Float[Array, "*batch n"]:
    """Solve independent tridiagonal systems over leading batch dimensions.

    Applies `tridiagonal_solve` vmapped over all leading dimensions.

    Args:
        lower: Sub-diagonals, shape ``(*batch, n - 1)``.
        diag: Main diagonals, shape ``(*batch, n)``.
        upper: Super-diagonals, shape ``(*batch, n - 1)``.
        rhs: Right-hand sides, shape ``(*batch, n)``.

    Returns:
        Solutions, shape ``(*batch, n)``.
    """
    n_batch = diag.ndim - 1
    fn = tridiagonal_solve
    for _ in range(n_batch):
        fn = eqx.filter_vmap(fn)
    return fn(lower, diag, upper, rhs)

Matrix diagonalization & shifted Kronecker-sum solves

Factor each 1D operator once with EigenDecomposition (non-symmetric diagonalizable factors with a real spectrum are supported), then solve \((A_0 \oplus A_1 \oplus \dots - \sigma I)\,x = b\) in tensor form for any shift \(\sigma\) with kronecker_sum_solve — the matrix-diagonalization method for separable operators such as tensor-product spectral Laplacians.

Structured linear algebra and Gaussian primitives for JAX.

EigenDecomposition

Bases: Module

Precomputed eigendecomposition A = V diag(λ) V⁻¹ of a square matrix.

Built once (on the host, outside jit) and then reused for any number of shifted solves (A − σI)⁻¹ b, each costing two dense mat-vecs — O(n²) — for any shift σ, which may be a traced value (so solves are jit/grad-compatible in σ).

Only real spectra are supported: the factorization is kept in real arithmetic, and from_matrix rejects matrices whose eigenvalues have non-negligible imaginary parts.

Attributes:

Name	Type	Description
eigenvalues	Float[Array, ' n']	Eigenvalues λ, shape (n,).
eigenvectors	Float[Array, 'n n']	Right eigenvectors V as columns, shape (n, n).
eigenvectors_inv	Float[Array, 'n n']	V⁻¹, shape (n, n) (Vᵀ for symmetric A).

Example

fac = gaussx.EigenDecomposition.from_matrix(A)
x = fac.solve_shifted(b, shift=2.0)  # (A - 2 I) x = b

Source code in src/gaussx/_linalg/_eigen_factorization.py

class EigenDecomposition(eqx.Module):
    r"""Precomputed eigendecomposition ``A = V diag(λ) V⁻¹`` of a square matrix.

    Built once (on the host, outside ``jit``) and then reused for any number
    of shifted solves ``(A − σI)⁻¹ b``, each costing two dense mat-vecs —
    ``O(n²)`` — for *any* shift ``σ``, which may be a traced value (so solves
    are ``jit``/``grad``-compatible in ``σ``).

    Only real spectra are supported: the factorization is kept in real
    arithmetic, and `from_matrix` rejects matrices whose eigenvalues have
    non-negligible imaginary parts.

    Attributes:
        eigenvalues: Eigenvalues ``λ``, shape ``(n,)``.
        eigenvectors: Right eigenvectors ``V`` as columns, shape ``(n, n)``.
        eigenvectors_inv: ``V⁻¹``, shape ``(n, n)`` (``Vᵀ`` for symmetric ``A``).

    Example:
        ```python
        fac = gaussx.EigenDecomposition.from_matrix(A)
        x = fac.solve_shifted(b, shift=2.0)  # (A - 2 I) x = b
        ```
    """

    eigenvalues: Float[Array, " n"]
    eigenvectors: Float[Array, "n n"]
    eigenvectors_inv: Float[Array, "n n"]

    @classmethod
    def from_matrix(
        cls,
        matrix: ArrayLike | lx.AbstractLinearOperator,
        *,
        symmetric: bool | None = None,
    ) -> EigenDecomposition:
        """Eigendecompose a concrete square matrix on the host.

        Uses NumPy (``eigh`` for symmetric input, ``eig`` otherwise), so it
        works for non-symmetric matrices on every JAX backend. Because of
        that, ``matrix`` must be concrete: call this outside ``jit`` (e.g.
        when constructing a solver), not on traced values.

        Args:
            matrix: Square matrix, or a lineax operator (materialized via
                ``as_matrix()``).
            symmetric: Use the symmetric path. Defaults to
                ``lx.is_symmetric(matrix)`` for operators and ``False`` for
                arrays.

        Returns:
            The factorization.

        Raises:
            ValueError: If ``matrix`` is not square, is a tracer, or has a
                spectrum with non-negligible imaginary parts.
        """
        if isinstance(matrix, lx.AbstractLinearOperator):
            if symmetric is None:
                symmetric = lx.is_symmetric(matrix)
            matrix = matrix.as_matrix()
        try:
            mat = np.asarray(matrix)
        except jax.errors.TracerArrayConversionError as err:
            raise ValueError(
                "EigenDecomposition.from_matrix needs a concrete matrix "
                "(it factorizes on the host with NumPy); call it outside jit."
            ) from err
        if mat.ndim != 2 or mat.shape[0] != mat.shape[1]:
            raise ValueError(f"matrix must be square, got shape {mat.shape}.")

        if symmetric:
            lam, V = np.linalg.eigh(mat)
            V_inv = V.T
        else:
            lam, V = np.linalg.eig(mat)
            scale = max(float(np.max(np.abs(lam))), 1.0)
            tol = _IMAG_TOLERANCE_FACTOR * np.finfo(mat.dtype).eps * scale
            if np.max(np.abs(lam.imag)) > tol:
                raise ValueError(
                    "EigenDecomposition supports real spectra only; the largest "
                    f"imaginary part is {np.max(np.abs(lam.imag)):.3e}."
                )
            lam, V = lam.real, V.real
            V_inv = np.linalg.inv(V)
        return cls(jnp.asarray(lam), jnp.asarray(V), jnp.asarray(V_inv))

    def as_matrix(self) -> Float[Array, "n n"]:
        """Reassemble ``V diag(λ) V⁻¹``."""
        return (self.eigenvectors * self.eigenvalues) @ self.eigenvectors_inv

    def solve_shifted(
        self,
        rhs: Float[Array, " n *batch"],
        shift: float | Float[Array, ""] = 0.0,
        *,
        drop: Bool[Array, " n"] | None = None,
    ) -> Float[Array, " n *batch"]:
        """Solve ``(A − shift·I) x = rhs`` along the leading axis of ``rhs``.

        ``x = V diag(1 / (λ − shift)) V⁻¹ rhs``.

        Args:
            rhs: Right-hand side(s), shape ``(n,)`` or ``(n, *batch)``.
            shift: Scalar shift ``σ`` (may be traced).
            drop: Optional boolean mask over eigenmodes. Masked modes get a
                zero coefficient instead of ``1 / (λ − σ)`` — a restricted
                pseudo-inverse, e.g. to project out a known null space such
                as the constant mode of a pure-Neumann Laplacian.

        Returns:
            The solution, same shape as ``rhs``.
        """
        return kronecker_sum_solve((self,), rhs, shift, drop=drop)

from_matrix(matrix: ArrayLike | lx.AbstractLinearOperator, *, symmetric: bool | None = None) -> EigenDecomposition classmethod

Eigendecompose a concrete square matrix on the host.

Uses NumPy (eigh for symmetric input, eig otherwise), so it works for non-symmetric matrices on every JAX backend. Because of that, matrix must be concrete: call this outside jit (e.g. when constructing a solver), not on traced values.

Parameters:

Name	Type	Description	Default
matrix	ArrayLike | AbstractLinearOperator	Square matrix, or a lineax operator (materialized via as_matrix()).	required
symmetric	bool | None	Use the symmetric path. Defaults to lx.is_symmetric(matrix) for operators and False for arrays.	None

Returns:

Type	Description
EigenDecomposition	The factorization.

Raises:

Type	Description
ValueError	If matrix is not square, is a tracer, or has a spectrum with non-negligible imaginary parts.

Source code in src/gaussx/_linalg/_eigen_factorization.py

@classmethod
def from_matrix(
    cls,
    matrix: ArrayLike | lx.AbstractLinearOperator,
    *,
    symmetric: bool | None = None,
) -> EigenDecomposition:
    """Eigendecompose a concrete square matrix on the host.

    Uses NumPy (``eigh`` for symmetric input, ``eig`` otherwise), so it
    works for non-symmetric matrices on every JAX backend. Because of
    that, ``matrix`` must be concrete: call this outside ``jit`` (e.g.
    when constructing a solver), not on traced values.

    Args:
        matrix: Square matrix, or a lineax operator (materialized via
            ``as_matrix()``).
        symmetric: Use the symmetric path. Defaults to
            ``lx.is_symmetric(matrix)`` for operators and ``False`` for
            arrays.

    Returns:
        The factorization.

    Raises:
        ValueError: If ``matrix`` is not square, is a tracer, or has a
            spectrum with non-negligible imaginary parts.
    """
    if isinstance(matrix, lx.AbstractLinearOperator):
        if symmetric is None:
            symmetric = lx.is_symmetric(matrix)
        matrix = matrix.as_matrix()
    try:
        mat = np.asarray(matrix)
    except jax.errors.TracerArrayConversionError as err:
        raise ValueError(
            "EigenDecomposition.from_matrix needs a concrete matrix "
            "(it factorizes on the host with NumPy); call it outside jit."
        ) from err
    if mat.ndim != 2 or mat.shape[0] != mat.shape[1]:
        raise ValueError(f"matrix must be square, got shape {mat.shape}.")

    if symmetric:
        lam, V = np.linalg.eigh(mat)
        V_inv = V.T
    else:
        lam, V = np.linalg.eig(mat)
        scale = max(float(np.max(np.abs(lam))), 1.0)
        tol = _IMAG_TOLERANCE_FACTOR * np.finfo(mat.dtype).eps * scale
        if np.max(np.abs(lam.imag)) > tol:
            raise ValueError(
                "EigenDecomposition supports real spectra only; the largest "
                f"imaginary part is {np.max(np.abs(lam.imag)):.3e}."
            )
        lam, V = lam.real, V.real
        V_inv = np.linalg.inv(V)
    return cls(jnp.asarray(lam), jnp.asarray(V), jnp.asarray(V_inv))

as_matrix() -> Float[Array, 'n n']

Reassemble V diag(λ) V⁻¹.

Source code in src/gaussx/_linalg/_eigen_factorization.py

def as_matrix(self) -> Float[Array, "n n"]:
    """Reassemble ``V diag(λ) V⁻¹``."""
    return (self.eigenvectors * self.eigenvalues) @ self.eigenvectors_inv

solve_shifted(rhs: Float[Array, ' n *batch'], shift: float | Float[Array, ''] = 0.0, *, drop: Bool[Array, ' n'] | None = None) -> Float[Array, ' n *batch']

Solve (A − shift·I) x = rhs along the leading axis of rhs.

x = V diag(1 / (λ − shift)) V⁻¹ rhs.

Parameters:

Name	Type	Description	Default
rhs	Float[Array, ' n *batch']	Right-hand side(s), shape (n,) or (n, *batch).	required
shift	float | Float[Array, '']	Scalar shift σ (may be traced).	0.0
drop	Bool[Array, ' n'] | None	Optional boolean mask over eigenmodes. Masked modes get a zero coefficient instead of 1 / (λ − σ) — a restricted pseudo-inverse, e.g. to project out a known null space such as the constant mode of a pure-Neumann Laplacian.	None

Returns:

Type	Description
Float[Array, ' n *batch']	The solution, same shape as rhs.

Source code in src/gaussx/_linalg/_eigen_factorization.py

def solve_shifted(
    self,
    rhs: Float[Array, " n *batch"],
    shift: float | Float[Array, ""] = 0.0,
    *,
    drop: Bool[Array, " n"] | None = None,
) -> Float[Array, " n *batch"]:
    """Solve ``(A − shift·I) x = rhs`` along the leading axis of ``rhs``.

    ``x = V diag(1 / (λ − shift)) V⁻¹ rhs``.

    Args:
        rhs: Right-hand side(s), shape ``(n,)`` or ``(n, *batch)``.
        shift: Scalar shift ``σ`` (may be traced).
        drop: Optional boolean mask over eigenmodes. Masked modes get a
            zero coefficient instead of ``1 / (λ − σ)`` — a restricted
            pseudo-inverse, e.g. to project out a known null space such
            as the constant mode of a pure-Neumann Laplacian.

    Returns:
        The solution, same shape as ``rhs``.
    """
    return kronecker_sum_solve((self,), rhs, shift, drop=drop)

kronecker_sum_solve(factors: Sequence[EigenDecomposition], rhs: Float[Array, '*dims'], shift: float | Float[Array, ''] = 0.0, *, drop: Bool[Array, '*dims'] | None = None) -> Float[Array, '*dims']

Solve a shifted Kronecker-sum system in tensor (unvectorized) form.

With A_k acting on axis k of the unknown tensor X (shape (n_0, …, n_{d−1}), possibly followed by trailing batch axes), solves

\[ \sum_k X \times_k A_k \;-\; \sigma X \;=\; R, \]

i.e. (A_0 ⊕ A_1 ⊕ … ⊕ A_{d−1} − σ I) vec(X) = vec(R). In 2D this is the Sylvester equation A_0 X + X A_1ᵀ − σX = R. With A_k = V_k Λ_k V_k⁻¹:

1. X̂ = R ×_0 V_0⁻¹ ×_1 V_1⁻¹ … (rotate into the eigenbasis),
2. X̂[i_0, …] /= λ⁰_{i_0} + λ¹_{i_1} + … − σ,
3. X = X̂ ×_0 V_0 ×_1 V_1 … (rotate back).

Cost is O(N · Σ_k n_k) for N = Π_k n_k unknowns, and the N × N Kronecker matrix is never formed; factors need not be symmetric.

Parameters:

Name	Type	Description	Default
factors	Sequence[EigenDecomposition]	One EigenDecomposition per leading axis of rhs.	required
rhs	Float[Array, '*dims']	Right-hand side, shape (n_0, …, n_{d−1}, *batch).	required
shift	float | Float[Array, '']	Scalar shift σ (may be traced).	0.0
drop	Bool[Array, '*dims'] | None	Optional boolean mask of shape (n_0, …, n_{d−1}); masked eigen-combinations get a zero coefficient (restricted pseudo-inverse), as in EigenDecomposition.solve_shifted.	None

Returns:

Type	Description
Float[Array, '*dims']	X, same shape as rhs.

Raises:

Type	Description
ValueError	If the factor sizes do not match the leading axes of rhs.

Source code in src/gaussx/_linalg/_eigen_factorization.py

def kronecker_sum_solve(
    factors: Sequence[EigenDecomposition],
    rhs: Float[Array, "*dims"],
    shift: float | Float[Array, ""] = 0.0,
    *,
    drop: Bool[Array, "*dims"] | None = None,
) -> Float[Array, "*dims"]:
    r"""Solve a shifted Kronecker-sum system in tensor (unvectorized) form.

    With ``A_k`` acting on axis ``k`` of the unknown tensor ``X`` (shape
    ``(n_0, …, n_{d−1})``, possibly followed by trailing batch axes), solves

    $$
    \sum_k X \times_k A_k \;-\; \sigma X \;=\; R,
    $$

    i.e. ``(A_0 ⊕ A_1 ⊕ … ⊕ A_{d−1} − σ I) vec(X) = vec(R)``. In 2D this is
    the Sylvester equation ``A_0 X + X A_1ᵀ − σX = R``. With
    ``A_k = V_k Λ_k V_k⁻¹``:

    1. ``X̂ = R ×_0 V_0⁻¹ ×_1 V_1⁻¹ …`` (rotate into the eigenbasis),
    2. ``X̂[i_0, …] /= λ⁰_{i_0} + λ¹_{i_1} + … − σ``,
    3. ``X = X̂ ×_0 V_0 ×_1 V_1 …`` (rotate back).

    Cost is ``O(N · Σ_k n_k)`` for ``N = Π_k n_k`` unknowns, and the
    ``N × N`` Kronecker matrix is never formed; factors need not be
    symmetric.

    Args:
        factors: One `EigenDecomposition` per leading axis of ``rhs``.
        rhs: Right-hand side, shape ``(n_0, …, n_{d−1}, *batch)``.
        shift: Scalar shift ``σ`` (may be traced).
        drop: Optional boolean mask of shape ``(n_0, …, n_{d−1})``; masked
            eigen-combinations get a zero coefficient (restricted
            pseudo-inverse), as in `EigenDecomposition.solve_shifted`.

    Returns:
        ``X``, same shape as ``rhs``.

    Raises:
        ValueError: If the factor sizes do not match the leading axes of
            ``rhs``.
    """
    d = len(factors)
    sizes = tuple(f.eigenvalues.shape[0] for f in factors)
    if rhs.shape[:d] != sizes:
        raise ValueError(
            f"rhs leading shape {rhs.shape[:d]} does not match factor sizes {sizes}."
        )

    x = rhs
    for k, fac in enumerate(factors):
        x = _apply_along_axis(fac.eigenvectors_inv, x, k)

    # Σ_k λᵏ broadcast over the d eigen axes: shape (n_0, …, n_{d−1}).
    lam_sum = jnp.zeros((), dtype=factors[0].eigenvalues.dtype)
    for k, fac in enumerate(factors):
        other_axes = tuple(j for j in range(d) if j != k)
        lam_sum = lam_sum + jnp.expand_dims(fac.eigenvalues, other_axes)
    denom = lam_sum - shift
    if drop is None:
        inv = 1.0 / denom
    else:
        inv = jnp.where(drop, 0.0, 1.0 / jnp.where(drop, 1.0, denom))
    inv = jnp.expand_dims(inv, tuple(range(d, x.ndim)))  # broadcast over batch axes
    x = x * inv

    for k, fac in enumerate(factors):
        x = _apply_along_axis(fac.eigenvectors, x, k)
    return x

Stable distances & Lyapunov

Structured linear algebra and Gaussian primitives for JAX.

stable_squared_distances(X: Float[Array, 'N D'], Z: Float[Array, 'M D'], *, compute_dtype: DTypeLike = jnp.float32, accumulate_dtype: DTypeLike | None = None) -> Float[Array, 'N M']

Squared Euclidean distances with mixed-precision stability.

The expansion ||x - z||^2 = ||x||^2 + ||z||^2 - 2 x^T z suffers catastrophic cancellation when the points are far from the origin relative to their separation, producing negative distances and non-PSD kernels. Two measures keep it accurate:

1. Both point sets are centred on their common mean first. Distances are translation-invariant, and centring removes the large shared offset that the expansion would otherwise cancel. This is what protects near-duplicate points, in any dtype.
2. The norms, the cross term and the subtraction are all computed in accumulate_dtype. Widening only the subtraction would not help: the digits are already lost in the rounded norms and cross term.

The result is cast back to compute_dtype.

Parameters:

Name	Type	Description	Default
X	Float[Array, 'N D']	First set of points, shape (N, D).	required
Z	Float[Array, 'M D']	Second set of points, shape (M, D).	required
compute_dtype	DTypeLike	Dtype of the inputs and the result (default float32).	float32
accumulate_dtype	DTypeLike | None	Dtype for the centred norms, cross term and subtraction. None (the default) means the widest float JAX has enabled: float64 with x64, float32 without.	None

Returns:

Type	Description
Float[Array, 'N M']	Squared distances, shape (N, M), guaranteed non-negative.

Raises:

Type	Description
ValueError	If accumulate_dtype is explicitly requested but not available, e.g. float64 while jax_enable_x64 is off. JAX would otherwise truncate it to float32 with only a warning.

Source code in src/gaussx/_linalg/_mixed_precision.py

def stable_squared_distances(
    X: Float[Array, "N D"],
    Z: Float[Array, "M D"],
    *,
    compute_dtype: DTypeLike = jnp.float32,
    accumulate_dtype: DTypeLike | None = None,
) -> Float[Array, "N M"]:
    r"""Squared Euclidean distances with mixed-precision stability.

    The expansion ``||x - z||^2 = ||x||^2 + ||z||^2 - 2 x^T z`` suffers
    catastrophic cancellation when the points are far from the origin
    relative to their separation, producing negative distances and
    non-PSD kernels. Two measures keep it accurate:

    1. Both point sets are centred on their common mean first. Distances
       are translation-invariant, and centring removes the large shared
       offset that the expansion would otherwise cancel. This is what
       protects near-duplicate points, in any dtype.
    2. The norms, the cross term and the subtraction are all computed in
       ``accumulate_dtype``. Widening only the subtraction would not help:
       the digits are already lost in the rounded norms and cross term.

    The result is cast back to ``compute_dtype``.

    Args:
        X: First set of points, shape ``(N, D)``.
        Z: Second set of points, shape ``(M, D)``.
        compute_dtype: Dtype of the inputs and the result (default
            float32).
        accumulate_dtype: Dtype for the centred norms, cross term and
            subtraction. ``None`` (the default) means the widest float JAX
            has enabled: float64 with x64, float32 without.

    Returns:
        Squared distances, shape ``(N, M)``, guaranteed non-negative.

    Raises:
        ValueError: If ``accumulate_dtype`` is explicitly requested but not
            available, e.g. float64 while ``jax_enable_x64`` is off. JAX
            would otherwise truncate it to float32 with only a warning.
    """
    if accumulate_dtype is None:
        accumulate_dtype = jax.dtypes.canonicalize_dtype(jnp.float64)
    elif jax.dtypes.canonicalize_dtype(accumulate_dtype) != jnp.dtype(accumulate_dtype):
        raise ValueError(
            f"accumulate_dtype={jnp.dtype(accumulate_dtype)} is not available "
            "(is jax_enable_x64 off?). Pass accumulate_dtype=None to use the "
            "widest available float."
        )

    X_c = X.astype(compute_dtype)
    Z_c = Z.astype(compute_dtype)

    # Centre on the common mean (translation-invariant), then widen.
    centre = reduce(jnp.concatenate([X_c, Z_c]), "N D -> D", "mean")
    X_a = einx.subtract("N D, D -> N D", X_c, centre).astype(accumulate_dtype)
    Z_a = einx.subtract("M D, D -> M D", Z_c, centre).astype(accumulate_dtype)

    X_sq = reduce(X_a**2, "N D -> N", "sum")
    Z_sq = reduce(Z_a**2, "M D -> M", "sum")
    cross = einsum(X_a, Z_a, "N D, M D -> N M")
    dist_sq = einx.add("N, M -> N M", X_sq, Z_sq) - 2.0 * cross

    # Clamp and cast back
    dist_sq = jnp.maximum(dist_sq, 0.0)
    return dist_sq.astype(compute_dtype)

discrete_lyapunov_solve(G: Float[Array, 'N N'], Q: Float[Array, 'N N']) -> Float[Array, 'N N']

Solve the discrete Lyapunov equation P - G P G^T = Q.

Uses the eigendecomposition G = V Λ V^{-1} so that

\[ \tilde{P} - \Lambda \tilde{P} \Lambda^T = V^{-1} Q V^{-T}, \quad \tilde{P}_{ij} = \frac{(V^{-1} Q V^{-T})_{ij}}{1 - \lambda_i \lambda_j}, \]

and then P = V \tilde{P} V^T. This avoids materializing the (N², N²) Kronecker matrix the vectorized formulation requires.

Cost is O(N³) (one general eigendecomposition + a couple of matrix multiplies) versus O(N⁶) for the vectorized solve, and the memory footprint drops from O(N⁴) to O(N²).

The discrete Lyapunov equation has a unique solution iff λ_i λ_j ≠ 1 for all eigenvalue pairs (λ_i, λ_j) of G. The standard sufficient condition — G is stable (spectral radius < 1) — guarantees this.

Warning

This implementation assumes G is diagonalizable. For defective G (e.g., a Jordan block with eigenvalue magnitude < 1) the eigenvector matrix V is singular and this solve will return NaN / inf. The discrete Lyapunov equation still has a unique solution in that case, but recovering it requires a Schur decomposition (Bartels-Stewart), which JAX does not currently expose. Fall back to the vectorized (I − G ⊗ G) solve for those operators.

Parameters:

Name	Type	Description	Default
G	Float[Array, 'N N']	Square matrix, shape (N, N). Should be stable for a unique steady-state solution.	required
Q	Float[Array, 'N N']	Right-hand side, shape (N, N). Typically symmetric PSD in Kalman-smoother applications.	required

Returns:

Type	Description
Float[Array, 'N N']	Symmetric matrix P of shape (N, N) solving
Float[Array, 'N N']	P - G P G^T = Q.

Source code in src/gaussx/_linalg/_lyapunov.py

def discrete_lyapunov_solve(
    G: Float[Array, "N N"],
    Q: Float[Array, "N N"],
) -> Float[Array, "N N"]:
    r"""Solve the discrete Lyapunov equation ``P - G P G^T = Q``.

    Uses the eigendecomposition ``G = V Λ V^{-1}`` so that

    $$
    \tilde{P} - \Lambda \tilde{P} \Lambda^T
        = V^{-1} Q V^{-T}, \quad
    \tilde{P}_{ij} = \frac{(V^{-1} Q V^{-T})_{ij}}{1 - \lambda_i \lambda_j},
    $$

    and then ``P = V \tilde{P} V^T``. This avoids materializing the
    ``(N², N²)`` Kronecker matrix the vectorized formulation requires.

    Cost is ``O(N³)`` (one general eigendecomposition + a couple of
    matrix multiplies) versus ``O(N⁶)`` for the vectorized solve, and
    the memory footprint drops from ``O(N⁴)`` to ``O(N²)``.

    The discrete Lyapunov equation has a unique solution iff
    ``λ_i λ_j ≠ 1`` for all eigenvalue pairs ``(λ_i, λ_j)`` of ``G``.
    The standard sufficient condition — ``G`` is stable
    (spectral radius < 1) — guarantees this.

    !!! warning

        This implementation assumes ``G`` is **diagonalizable**. For
        defective ``G`` (e.g., a Jordan block with eigenvalue magnitude
        ``< 1``) the eigenvector matrix ``V`` is singular and this
        solve will return ``NaN`` / ``inf``. The discrete Lyapunov
        equation still has a unique solution in that case, but
        recovering it requires a Schur decomposition (Bartels-Stewart),
        which JAX does not currently expose. Fall back to the
        vectorized ``(I − G ⊗ G)`` solve for those operators.

    Args:
        G: Square matrix, shape ``(N, N)``. Should be stable for a
            unique steady-state solution.
        Q: Right-hand side, shape ``(N, N)``. Typically symmetric PSD
            in Kalman-smoother applications.

    Returns:
        Symmetric matrix ``P`` of shape ``(N, N)`` solving
        ``P - G P G^T = Q``.
    """
    eigs, V = jnp.linalg.eig(G)
    # Use solves rather than an explicit ``inv(V)``: more stable, and
    # cheaper when V is well-conditioned. We need
    #     Q_tilde = V^{-1} Q V^{-T}
    # so first solve V Y = Q for Y = V^{-1} Q, then solve V Z = Y^T for
    # Z = V^{-1} Y^T = V^{-1} Q (V^{-1})^T (so Q_tilde = Z^T = V^{-1} Q V^{-T}).
    Q_complex = Q.astype(V.dtype)
    Y = jnp.linalg.solve(V, Q_complex)
    Q_tilde = jnp.linalg.solve(V, Y.T).T
    denom = 1.0 - eigs[:, None] * eigs[None, :]
    P_tilde = Q_tilde / denom
    P = V @ P_tilde @ V.T
    P_real = jnp.real(P)
    # Symmetrize to eliminate residual floating-point asymmetry —
    # consistent with conditional / infinite_horizon_smoother covariance
    # post-processing.
    return symmetrize(P_real)