Gaussian Processes¶

Layer 3 recipes for GP inference: conditioning, whitening, prediction caches, pathwise (Matheron) sampling, variational bounds, and cross-validation — all expressed over covariance operators so structured kernels keep their fast paths end to end. The modelling shell (kernels with hyperparameter priors, NumPyro sites) lives downstream; gaussx owns the math.

Conditioning & prediction¶

The standard posterior

\[ \mu_* = K_{*f}\,K^{-1}(y - \mu), \qquad \Sigma_{**} = K_{**} - K_{*f}\,K^{-1}K_{f*} \]

plus a precomputed-cache variant for repeated test-time queries and a Kronecker-structured path for separable kernels on grids.

The sparse-GP helpers share one convention (gh-353):

Object	Name	Shape / type
Inducing covariance (first argument)	`K_zz`	`(M, M)` array or PSD operator
Cross-covariance (data × inducing)	`K_xz`	`(N, M)`
Prior at the evaluation points	`K_xx` / `K_xx_diag`	`(N, N)` / `(N,)`
Variational mean	`q_mu`	`(M, R)` or `(M,)`
Variational root	`q_sqrt`	`(R, M, M)` or `(M, R)`

sparse_conditional follows it; base_conditional (cross-covariance (M, N)) is deprecated in its favour. collapsed_elbo's K_diag= keyword is now K_xx_diag=, and the unused solver= of collapsed_elbo and gauss_kl is deprecated; the old forms warn until gaussx 0.7.0.

Structured linear algebra and Gaussian primitives for JAX.

`PredictionCache` ¶

Bases: Module

Cached training solve (and factor) for amortized predictions.

A cache is tied to the operator it was built from: rebuild it whenever the kernel hyperparameters (and so K_y) change.

Attributes:

Name	Type	Description
`alpha`	`Float[Array, ' N']`	Solved weights `K_y^{-1} y`, shape `(N,)`.
`operator`	`AbstractLinearOperator \| None`	The training covariance `K_y` the cache was built from, or `None` for a cache constructed from `alpha` alone.
`factor`	`AbstractLinearOperator \| None`	Lower Cholesky factor `L` of `K_y` (structured where `gaussx.cholesky` preserves structure), or `None` when `K_y` is not factorised (iterative solvers, structured operators with their own solve).

Source code in src/gaussx/_gp/_prediction_cache.py

class PredictionCache(eqx.Module):
    """Cached training solve (and factor) for amortized predictions.

    A cache is tied to the operator it was built from: rebuild it whenever
    the kernel hyperparameters (and so ``K_y``) change.

    Attributes:
        alpha: Solved weights ``K_y^{-1} y``, shape ``(N,)``.
        operator: The training covariance ``K_y`` the cache was built
            from, or ``None`` for a cache constructed from ``alpha`` alone.
        factor: Lower Cholesky factor ``L`` of ``K_y`` (structured where
            `gaussx.cholesky` preserves structure), or ``None`` when
            ``K_y`` is not factorised (iterative solvers, structured
            operators with their own solve).
    """

    alpha: Float[Array, " N"]
    operator: lx.AbstractLinearOperator | None = None
    factor: lx.AbstractLinearOperator | None = None

`sparse_conditional(K_zz: Float[Array, 'M M'] | lx.AbstractLinearOperator, K_xz: Float[Array, 'N M'], K_xx: Float[Array, 'N N'] | Float[Array, ' N'], q_mu: Float[Array, 'M R'] | Float[Array, ' M'], *, q_sqrt: Float[Array, 'R M M'] | Float[Array, 'M R'] | Float[Array, 'M M'] | Float[Array, ' M'] | None = None, white: bool = False) -> tuple[Float[Array, 'N R'], Float[Array, ...]]` ¶

Sparse-GP conditional q(f_x) = \int p(f_x | u) q(u) du.

Follows the gaussx._gp sparse-GP convention: inducing covariance K_zz first (array or operator), cross-covariance K_xz with shape (N, M) (data × inducing), prior K_xx at the evaluation points, variational q(u) = N(q_mu, q_sqrt q_sqrtᵀ).

With L = chol(K_zz) and A = L^{-1} K_xz^\top (shape (M, N)),

\[ \begin{aligned} \mu &= A^\top L^{-1} q_\mu \;(= K_{xz} K_{zz}^{-1} q_\mu), \qquad \mu = A^\top q_\mu \text{ if white}, \\ \Sigma_r &= K_{xx} - A^\top A + W_r^\top W_r, \qquad W_r = S_r^\top P, \quad P = \begin{cases} A & \text{white} \\ L^{-\top} A & \text{otherwise} \end{cases}, \end{aligned} \]

where S_r is the r-th q_sqrt factor (diag(q_sqrt[:, r]) for the diagonal layout). Only the diagonal of Σ_r is formed when K_xx is a vector. K_zz is factorised with gaussx.cholesky and every solve with its factor is a structured triangular solve, so a Kronecker or BlockDiag K_zz is never densified.

Parameters:

Name	Type	Description	Default
`K_zz`	`Float[Array, 'M M'] \| AbstractLinearOperator`	Prior covariance at the inducing points, `(M, M)` array or PSD operator.	required
`K_xz`	`Float[Array, 'N M']`	Cross-covariance `k(X, Z)`, shape `(N, M)`.	required
`K_xx`	`Float[Array, 'N N'] \| Float[Array, ' N']`	Prior covariance at `X`: full `(N, N)` or diagonal `(N,)`.	required
`q_mu`	`Float[Array, 'M R'] \| Float[Array, ' M']`	Variational mean `(M, R)`, or `(M,)` for a single output (the outputs then drop their trailing `R` axis).	required
`q_sqrt`	`Float[Array, 'R M M'] \| Float[Array, 'M R'] \| Float[Array, 'M M'] \| Float[Array, ' M'] \| None`	Optional variational root. Full `(R, M, M)` (lower triangular) or diagonal `(M, R)`; with a 1-D `q_mu` also `(M, M)` or `(M,)`.	`None`
`white`	`bool`	If `True`, `q_mu` and `q_sqrt` are in whitened space (prior `N(0, I)`).	`False`

Returns:

Type	Description
`Float[Array, 'N R']`	`(mean, var)`: `mean` `(N, R)`; `var` `(N, N, R)` (full
`Float[Array, ...]`	`K_xx`) or `(N, R)` (diagonal); no trailing `R` for a 1-D
`tuple[Float[Array, 'N R'], Float[Array, ...]]`	`q_mu`. Diagonal variances are clipped at 0 against round-off; a
`tuple[Float[Array, 'N R'], Float[Array, ...]]`	full covariance is returned unclipped.

Raises:

Type	Description
`ValueError`	If `q_mu` / `q_sqrt` do not match `M` and `R`.

Source code in src/gaussx/_gp/_base_conditional.py

def sparse_conditional(
    K_zz: Float[Array, "M M"] | lx.AbstractLinearOperator,
    K_xz: Float[Array, "N M"],
    K_xx: Float[Array, "N N"] | Float[Array, " N"],
    q_mu: Float[Array, "M R"] | Float[Array, " M"],
    *,
    q_sqrt: Float[Array, "R M M"]
    | Float[Array, "M R"]
    | Float[Array, "M M"]
    | Float[Array, " M"]
    | None = None,
    white: bool = False,
) -> tuple[Float[Array, "N R"], Float[Array, ...]]:
    r"""Sparse-GP conditional ``q(f_x) = \int p(f_x | u) q(u) du``.

    Follows the `gaussx._gp` sparse-GP convention: inducing covariance
    ``K_zz`` first (array or operator), cross-covariance ``K_xz`` with
    shape ``(N, M)`` (data × inducing), prior ``K_xx`` at the evaluation
    points, variational ``q(u) = N(q_mu, q_sqrt q_sqrtᵀ)``.

    With ``L = chol(K_zz)`` and ``A = L^{-1} K_xz^\top`` (shape ``(M, N)``),

    $$
    \begin{aligned}
    \mu &= A^\top L^{-1} q_\mu
        \;(= K_{xz} K_{zz}^{-1} q_\mu), \qquad
        \mu = A^\top q_\mu \text{ if white}, \\
    \Sigma_r &= K_{xx} - A^\top A + W_r^\top W_r, \qquad
    W_r = S_r^\top P, \quad
    P = \begin{cases} A & \text{white} \\ L^{-\top} A & \text{otherwise}
        \end{cases},
    \end{aligned}
    $$

    where ``S_r`` is the ``r``-th ``q_sqrt`` factor (``diag(q_sqrt[:, r])``
    for the diagonal layout). Only the diagonal of ``Σ_r`` is formed when
    ``K_xx`` is a vector. ``K_zz`` is factorised with `gaussx.cholesky`
    and every solve with its factor is a structured triangular solve, so a
    `Kronecker` or `BlockDiag` ``K_zz`` is never densified.

    Args:
        K_zz: Prior covariance at the inducing points, ``(M, M)`` array or
            PSD operator.
        K_xz: Cross-covariance ``k(X, Z)``, shape ``(N, M)``.
        K_xx: Prior covariance at ``X``: full ``(N, N)`` or diagonal ``(N,)``.
        q_mu: Variational mean ``(M, R)``, or ``(M,)`` for a single output
            (the outputs then drop their trailing ``R`` axis).
        q_sqrt: Optional variational root. Full ``(R, M, M)`` (lower
            triangular) or diagonal ``(M, R)``; with a 1-D ``q_mu`` also
            ``(M, M)`` or ``(M,)``.
        white: If ``True``, ``q_mu`` and ``q_sqrt`` are in whitened space
            (prior ``N(0, I)``).

    Returns:
        ``(mean, var)``: ``mean`` ``(N, R)``; ``var`` ``(N, N, R)`` (full
        ``K_xx``) or ``(N, R)`` (diagonal); no trailing ``R`` for a 1-D
        ``q_mu``. Diagonal variances are clipped at 0 against round-off; a
        full covariance is returned unclipped.

    Raises:
        ValueError: If ``q_mu`` / ``q_sqrt`` do not match ``M`` and ``R``.
    """
    K_zz_op = _as_psd_operator(K_zz)
    M = K_zz_op.in_size()
    single_output = q_mu.ndim == 1
    f, q_sqrt = _check_shapes(M, q_mu, q_sqrt)
    R = f.shape[1]

    L = cholesky(K_zz_op)  # lower factor, structured where possible
    A = solve_columns(L, rearrange(K_xz, "n m -> m n"))  # (M, N)

    if white:
        mean = einsum(A, f, "m n, m r -> n r")
        P = A
    else:
        mean = einsum(A, solve_columns(L, f), "m n, m r -> n r")
        P = solve_columns(L.transpose(), A)  # L^{-T} A = K_zz^{-1} K_zx

    full_cov = K_xx.ndim == 2
    if full_cov:
        var_base = K_xx - einsum(A, A, "m i, m j -> i j")  # (N, N)
    else:
        var_base = K_xx - reduce(A**2, "m n -> n", "sum")  # (N,)

    if q_sqrt is None:
        if full_cov:
            var = repeat(var_base, "i j -> i j r", r=R)
        else:
            var = repeat(var_base, "n -> n r", r=R)
    else:
        if q_sqrt.ndim == 2:  # (M, R) diagonal roots
            W = einx.multiply("m r, m n -> r m n", q_sqrt, P)
        else:  # (R, M, M) lower-triangular roots: W_r = S_rᵀ P
            W = einsum(q_sqrt, P, "r m k, m n -> r k n")
        if full_cov:
            var = einx.add(
                "i j, i j r -> i j r", var_base, einsum(W, W, "r k i, r k j -> i j r")
            )
        else:
            var = einx.add(
                "n, n r -> n r", var_base, reduce(W**2, "r k n -> n r", "sum")
            )

    if not full_cov:
        # K_xx - diag(Q_xx) can round below zero (gh-363).
        var = jnp.maximum(var, 0.0)
    if single_output:
        return mean[:, 0], var[..., 0]
    return mean, var

`build_prediction_cache(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None) -> PredictionCache` ¶

Solve K_y alpha = y and cache the result.

For a PSD dense, diagonal, Kronecker or BlockDiag operator with solver None or DenseSolver, computes L = cholesky(K_y) once, solves alpha = L^{-T} L^{-1} y through it and stores L so that predict_variance needs only triangular solves. Otherwise (iterative strategies, other structured operators) alpha comes from solver and factor is None.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Training covariance operator `K_y`, shape `(N, N)`.	required
`y`	`Float[Array, ' N']`	Training targets, shape `(N,)`.	required
`solver`	`AbstractSolveStrategy \| None`	Optional solve strategy. When `None`, falls back to structural-dispatch `gaussx.solve`.	`None`

Returns:

Type	Description
`PredictionCache`	A `PredictionCache` holding `alpha`, `K_y` and, where
`PredictionCache`	available, its Cholesky factor.

Source code in src/gaussx/_gp/_prediction_cache.py

def build_prediction_cache(
    operator: lx.AbstractLinearOperator,
    y: Float[Array, " N"],
    *,
    solver: AbstractSolveStrategy | None = None,
) -> PredictionCache:
    """Solve ``K_y alpha = y`` and cache the result.

    For a PSD dense, diagonal, `Kronecker` or `BlockDiag` operator with
    ``solver`` ``None`` or `DenseSolver`, computes ``L = cholesky(K_y)``
    once, solves ``alpha = L^{-T} L^{-1} y`` through it and stores ``L``
    so that `predict_variance` needs only triangular solves. Otherwise
    (iterative strategies, other structured operators) ``alpha`` comes
    from ``solver`` and ``factor`` is ``None``.

    Args:
        operator: Training covariance operator ``K_y``, shape ``(N, N)``.
        y: Training targets, shape ``(N,)``.
        solver: Optional solve strategy. When ``None``, falls back
            to structural-dispatch `gaussx.solve`.

    Returns:
        A `PredictionCache` holding ``alpha``, ``K_y`` and, where
        available, its Cholesky factor.
    """
    if (
        isinstance(operator, _FACTORABLE)
        and lx.is_positive_semidefinite(operator)
        and (solver is None or isinstance(solver, DenseSolver))
    ):
        L = cholesky(operator)
        alpha = solve(L.transpose(), solve(L, y))
        return PredictionCache(alpha=alpha, operator=operator, factor=L)
    alpha = dispatch_solve(operator, y, solver)
    return PredictionCache(alpha=alpha, operator=operator)

`predict_mean(cache: PredictionCache, K_cross: Float[Array, 'Nt N']) -> Float[Array, ' Nt']` ¶

Predictive mean: mu* = K_*f @ alpha.

Parameters:

Name	Type	Description	Default
`cache`	`PredictionCache`	Prediction cache from `build_prediction_cache`.	required
`K_cross`	`Float[Array, 'Nt N']`	Cross-covariance matrix, shape `(Nt, N)`.	required

Returns:

Type	Description
`Float[Array, ' Nt']`	Predictive mean, shape `(Nt,)`.

Source code in src/gaussx/_gp/_prediction_cache.py

def predict_mean(
    cache: PredictionCache,
    K_cross: Float[Array, "Nt N"],
) -> Float[Array, " Nt"]:
    """Predictive mean: ``mu* = K_*f @ alpha``.

    Args:
        cache: Prediction cache from `build_prediction_cache`.
        K_cross: Cross-covariance matrix, shape ``(Nt, N)``.

    Returns:
        Predictive mean, shape ``(Nt,)``.
    """
    return K_cross @ cache.alpha

`predict_variance(cache: PredictionCache | Float[Array, 'Nt N'] | None = None, K_cross: Float[Array, 'Nt N'] | Float[Array, ' Nt'] | None = None, K_test_diag: Float[Array, ' Nt'] | lx.AbstractLinearOperator | None = None, *, operator: lx.AbstractLinearOperator | None = None, solver: AbstractSolveStrategy | None = None) -> Float[Array, ' Nt']` ¶

Predictive variance sigma^2_* = k_** - diag(K_*f K_y^{-1} K_f*).

Call as predict_variance(cache, K_cross, K_test_diag), with the cache first as in predict_mean. With a cached factor L:

\[ V = L^{-1} K_{*f}^\top, \qquad \sigma^2_{*,i} = k_{**,i} - \sum_n V_{ni}^2, \]

triangular solves only (no refactorisation of K_y). Without a factor, solves K_y v_i = K_cross[i, :] for each test row with solver and returns k_{**,i} - K_cross[i, :] v_i.

The old order predict_variance(K_cross, K_test_diag, operator) (or operator= as a keyword) still works for one release, emits a DeprecationWarning, and refactorises K_y on every call.

Parameters:

Name	Type	Description	Default
`cache`	`PredictionCache \| Float[Array, 'Nt N'] \| None`	Prediction cache from `build_prediction_cache`.	`None`
`K_cross`	`Float[Array, 'Nt N'] \| Float[Array, ' Nt'] \| None`	Cross-covariance matrix, shape `(Nt, N)`.	`None`
`K_test_diag`	`Float[Array, ' Nt'] \| AbstractLinearOperator \| None`	Prior variance at test points, shape `(Nt,)`.	`None`
`operator`	`AbstractLinearOperator \| None`	Deprecated; only for the old call form.	`None`
`solver`	`AbstractSolveStrategy \| None`	Optional solve strategy for the no-factor path. When `None`, falls back to structural-dispatch `gaussx.solve`.	`None`

Returns:

Type	Description
`Float[Array, ' Nt']`	Predictive variance, shape `(Nt,)`.

Raises:

Type	Description
`ValueError`	If the cache holds neither a factor nor an operator.

Source code in src/gaussx/_gp/_prediction_cache.py

def predict_variance(
    cache: PredictionCache | Float[Array, "Nt N"] | None = None,
    K_cross: Float[Array, "Nt N"] | Float[Array, " Nt"] | None = None,
    K_test_diag: Float[Array, " Nt"] | lx.AbstractLinearOperator | None = None,
    *,
    operator: lx.AbstractLinearOperator | None = None,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, " Nt"]:
    r"""Predictive variance ``sigma^2_* = k_** - diag(K_*f K_y^{-1} K_f*)``.

    Call as ``predict_variance(cache, K_cross, K_test_diag)``, with the
    cache first as in `predict_mean`. With a cached factor ``L``:

    $$
    V = L^{-1} K_{*f}^\top, \qquad
    \sigma^2_{*,i} = k_{**,i} - \sum_n V_{ni}^2,
    $$

    triangular solves only (no refactorisation of ``K_y``). Without a
    factor, solves ``K_y v_i = K_cross[i, :]`` for each test row with
    ``solver`` and returns ``k_{**,i} - K_cross[i, :] v_i``.

    The old order ``predict_variance(K_cross, K_test_diag, operator)``
    (or ``operator=`` as a keyword) still works for one release, emits a
    `DeprecationWarning`, and refactorises ``K_y`` on every call.

    Args:
        cache: Prediction cache from `build_prediction_cache`.
        K_cross: Cross-covariance matrix, shape ``(Nt, N)``.
        K_test_diag: Prior variance at test points, shape ``(Nt,)``.
        operator: Deprecated; only for the old call form.
        solver: Optional solve strategy for the no-factor path. When
            ``None``, falls back to structural-dispatch `gaussx.solve`.

    Returns:
        Predictive variance, shape ``(Nt,)``.

    Raises:
        ValueError: If the cache holds neither a factor nor an operator.
    """

    if not isinstance(cache, PredictionCache):
        warn_deprecated(
            "predict_variance(K_cross, K_test_diag, operator) is deprecated "
            "and will be removed in gaussx 0.7.0; build a cache "
            "with build_prediction_cache(operator, y) and call "
            "predict_variance(cache, K_cross, K_test_diag), which reuses the "
            "cached factorisation."
        )
        if operator is None:  # all positional: (K_cross, K_test_diag, operator)
            old = (cache, K_cross, K_test_diag)
        elif cache is not None:  # (K_cross, [K_test_diag], operator=...)
            old = (cache, K_cross if K_test_diag is None else K_test_diag, operator)
        else:  # all keywords
            old = (K_cross, K_test_diag, operator)
        if not isinstance(old[2], lx.AbstractLinearOperator):
            msg = "predict_variance needs a PredictionCache or an operator."
            raise TypeError(msg)
        Kc, kd = cast(Array, old[0]), cast(Array, old[1])
        V = solve_rows(old[2], Kc, solver=solver)
        return kd - reduce(Kc * V, "t n -> t", "sum")

    if operator is not None:
        msg = "operator= is only for the deprecated call form; the cache holds K_y."
        raise TypeError(msg)
    Kc, kd = cast(Array, K_cross), cast(Array, K_test_diag)
    if cache.factor is not None:
        V = solve_columns(cache.factor, rearrange(Kc, "t n -> n t"))
        return kd - reduce(V**2, "n t -> t", "sum")
    if cache.operator is None:
        msg = (
            "This PredictionCache holds only alpha; build it with "
            "build_prediction_cache(operator, y) to predict variances."
        )
        raise ValueError(msg)
    V = solve_rows(cache.operator, Kc, solver=solver)
    return kd - reduce(Kc * V, "t n -> t", "sum")

`conditional_interpolate(A_fwd: Float[Array, 'd d'], Q_fwd: Float[Array, 'd d'], A_bwd: Float[Array, 'd d'], Q_bwd: Float[Array, 'd d'], mu_prev: Float[Array, ' d'], P_prev: Float[Array, 'd d'], mu_next: Float[Array, ' d'], P_next: Float[Array, 'd d'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' d'], Float[Array, 'd d']]` ¶

Two-filter fusion at t from a filtered state and a backward message.

For an SDE-discretised state-space model with no observation at t,

\[ x_t \mid x_{t^-} \sim \mathcal{N}(A_{fwd} x_{t^-}, Q_{fwd}), \qquad x_{t^+} \mid x_t \sim \mathcal{N}(A_{bwd} x_t, Q_{bwd}), \]

combines two conditionally independent sources of information about x_t by adding precisions (the two-filter smoother):

(mu_prev, P_prev): the forward-filtered state at t^-, p(x_{t^-} \mid y_{\le t^-}) = N(mu_prev, P_prev);
(mu_next, P_next): the backward-information message at t^+, i.e. the likelihood of the data at and after t^+ written as p(y_{\ge t^+} \mid x_{t^+}) \propto N(mu_next; x_{t^+}, P_next) (no prior on x_{t^+}).

With those inputs it returns p(x_t \mid y_{\le t^-}, y_{\ge t^+}) exactly:

\[ \begin{aligned} m_{fwd} &= A_{fwd}\,\mu_{prev}, \quad P_{fwd} = A_{fwd} P_{prev} A_{fwd}^\top + Q_{fwd}, \\ \Lambda_{bwd} &= A_{bwd}^\top (P_{next} + Q_{bwd})^{-1} A_{bwd}, \quad \eta_{bwd} = A_{bwd}^\top (P_{next} + Q_{bwd})^{-1} \mu_{next}, \\ P &= (P_{fwd}^{-1} + \Lambda_{bwd})^{-1}, \quad \mu = P\,(P_{fwd}^{-1} m_{fwd} + \eta_{bwd}). \end{aligned} \]

Warning

Do not pass smoothed marginals (e.g. from rts_smoother) for both ends: each already contains all the data, so the data is counted twice and the result is over-confident (about 20% too small a covariance on a 3-state chain, gh-288). To interpolate from a filtered state at t^- and a smoothed state at t^+, use rts_interpolate.

Parameters:

Name	Type	Description	Default
`A_fwd`	`Float[Array, 'd d']`	Forward transition from `t^-` to `t`, shape `(d, d)`.	required
`Q_fwd`	`Float[Array, 'd d']`	Forward process noise, shape `(d, d)`.	required
`A_bwd`	`Float[Array, 'd d']`	Transition from `t` to `t^+`, shape `(d, d)`.	required
`Q_bwd`	`Float[Array, 'd d']`	Process noise from `t` to `t^+`, shape `(d, d)`.	required
`mu_prev`	`Float[Array, ' d']`	Forward-filtered mean at `t^-`, shape `(d,)`.	required
`P_prev`	`Float[Array, 'd d']`	Forward-filtered covariance at `t^-`, shape `(d, d)`.	required
`mu_next`	`Float[Array, ' d']`	Mean of the backward-information message at `t^+`, shape `(d,)`.	required
`P_next`	`Float[Array, 'd d']`	Covariance of the backward-information message at `t^+`, shape `(d, d)`.	required
`solver`	`AbstractSolverStrategy \| None`	Optional solver strategy for structured linear algebra. When `None`, falls back to structural dispatch.	`None`

Returns:

Type	Description
`tuple[Float[Array, ' d'], Float[Array, 'd d']]`	Tuple `(mean, cov)` of `x_t` given the data on both sides.

Source code in src/gaussx/_gp/_interpolation.py

def conditional_interpolate(
    A_fwd: Float[Array, "d d"],
    Q_fwd: Float[Array, "d d"],
    A_bwd: Float[Array, "d d"],
    Q_bwd: Float[Array, "d d"],
    mu_prev: Float[Array, " d"],
    P_prev: Float[Array, "d d"],
    mu_next: Float[Array, " d"],
    P_next: Float[Array, "d d"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " d"], Float[Array, "d d"]]:
    r"""Two-filter fusion at ``t`` from a filtered state and a backward message.

    For an SDE-discretised state-space model with no observation at ``t``,

    $$
    x_t \mid x_{t^-} \sim \mathcal{N}(A_{fwd} x_{t^-}, Q_{fwd}), \qquad
    x_{t^+} \mid x_t \sim \mathcal{N}(A_{bwd} x_t, Q_{bwd}),
    $$

    combines two **conditionally independent** sources of information about
    ``x_t`` by adding precisions (the two-filter smoother):

    - ``(mu_prev, P_prev)``: the **forward-filtered** state at ``t^-``,
      ``p(x_{t^-} \mid y_{\le t^-}) = N(mu_prev, P_prev)``;
    - ``(mu_next, P_next)``: the **backward-information message** at
      ``t^+``, i.e. the likelihood of the data at and after ``t^+`` written
      as ``p(y_{\ge t^+} \mid x_{t^+}) \propto N(mu_next; x_{t^+}, P_next)``
      (no prior on ``x_{t^+}``).

    With those inputs it returns ``p(x_t \mid y_{\le t^-}, y_{\ge t^+})``
    exactly:

    $$
    \begin{aligned}
    m_{fwd} &= A_{fwd}\,\mu_{prev}, \quad
    P_{fwd} = A_{fwd} P_{prev} A_{fwd}^\top + Q_{fwd}, \\
    \Lambda_{bwd} &= A_{bwd}^\top (P_{next} + Q_{bwd})^{-1} A_{bwd}, \quad
    \eta_{bwd} = A_{bwd}^\top (P_{next} + Q_{bwd})^{-1} \mu_{next}, \\
    P &= (P_{fwd}^{-1} + \Lambda_{bwd})^{-1}, \quad
    \mu = P\,(P_{fwd}^{-1} m_{fwd} + \eta_{bwd}).
    \end{aligned}
    $$

    Warning:
        Do **not** pass smoothed marginals (e.g. from `rts_smoother`) for
        both ends: each already contains all the data, so the data is
        counted twice and the result is over-confident (about 20% too small
        a covariance on a 3-state chain, gh-288). To interpolate from a
        filtered state at ``t^-`` and a smoothed state at ``t^+``, use
        `rts_interpolate`.

    Args:
        A_fwd: Forward transition from ``t^-`` to ``t``, shape ``(d, d)``.
        Q_fwd: Forward process noise, shape ``(d, d)``.
        A_bwd: Transition from ``t`` to ``t^+``, shape ``(d, d)``.
        Q_bwd: Process noise from ``t`` to ``t^+``, shape ``(d, d)``.
        mu_prev: Forward-filtered mean at ``t^-``, shape ``(d,)``.
        P_prev: Forward-filtered covariance at ``t^-``, shape ``(d, d)``.
        mu_next: Mean of the backward-information message at ``t^+``,
            shape ``(d,)``.
        P_next: Covariance of the backward-information message at ``t^+``,
            shape ``(d, d)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(mean, cov)`` of ``x_t`` given the data on both sides.
    """
    # Forward prediction to t
    m_fwd = A_fwd @ mu_prev
    P_fwd = A_fwd @ P_prev @ A_fwd.T + Q_fwd

    # Forward information
    P_fwd_op = lx.MatrixLinearOperator(P_fwd, lx.positive_semidefinite_tag)
    Lambda_fwd = inv(P_fwd_op).as_matrix()
    eta1_fwd = dispatch_solve(P_fwd_op, m_fwd, solver)

    # Backward information from t+
    S_bwd = P_next + Q_bwd
    S_bwd_op = lx.MatrixLinearOperator(S_bwd, lx.positive_semidefinite_tag)
    Lambda_bwd = A_bwd.T @ solve_columns(S_bwd_op, A_bwd, solver=solver)
    eta1_bwd = A_bwd.T @ dispatch_solve(S_bwd_op, mu_next, solver)

    # Fuse forward and backward
    Lambda = Lambda_fwd + Lambda_bwd
    Lambda_op = lx.MatrixLinearOperator(Lambda, lx.positive_semidefinite_tag)
    P = inv(Lambda_op).as_matrix()
    m = P @ (eta1_fwd + eta1_bwd)

    return m, P

`rts_interpolate(A_fwd: Float[Array, 'd d'], Q_fwd: Float[Array, 'd d'], A_bwd: Float[Array, 'd d'], Q_bwd: Float[Array, 'd d'], filtered_mean_prev: Float[Array, ' d'], filtered_cov_prev: Float[Array, 'd d'], smoothed_mean_next: Float[Array, ' d'], smoothed_cov_next: Float[Array, 'd d'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' d'], Float[Array, 'd d']]` ¶

Smoothed marginal at an unobserved t from filtered and smoothed states.

For the model of conditional_interpolate (no observation at t, t^- < t < t^+), takes the forward-filtered state at t^- and the smoothed state at t^+ — what kalman_filter and rts_smoother return — and performs one Rauch-Tung-Striebel backward step through t:

\[ \begin{aligned} m_t^- &= A_{fwd} m_f, \quad P_t^- = A_{fwd} P_f A_{fwd}^\top + Q_{fwd}, \\ m_+ &= A_{bwd} m_t^-, \quad P_+ = A_{bwd} P_t^- A_{bwd}^\top + Q_{bwd}, \\ G &= P_t^- A_{bwd}^\top P_+^{-1}, \\ m_t &= m_t^- + G (m_s - m_+), \quad P_t = P_t^- + G (P_s - P_+) G^\top . \end{aligned} \]

This is exact: given x_{t^+}, x_t is independent of the data at and after t^+, so p(x_t | y) = \int p(x_t | x_{t^+}, y_{\le t^-}) p(x_{t^+} | y) dx_{t^+}.

Parameters:

Name	Type	Description	Default
`A_fwd`	`Float[Array, 'd d']`	Forward transition from `t^-` to `t`, shape `(d, d)`.	required
`Q_fwd`	`Float[Array, 'd d']`	Forward process noise, shape `(d, d)`.	required
`A_bwd`	`Float[Array, 'd d']`	Transition from `t` to `t^+`, shape `(d, d)`.	required
`Q_bwd`	`Float[Array, 'd d']`	Process noise from `t` to `t^+`, shape `(d, d)`.	required
`filtered_mean_prev`	`Float[Array, ' d']`	Filtered mean `m_f` at `t^-`, shape `(d,)`.	required
`filtered_cov_prev`	`Float[Array, 'd d']`	Filtered covariance `P_f` at `t^-`, shape `(d, d)`.	required
`smoothed_mean_next`	`Float[Array, ' d']`	Smoothed mean `m_s` at `t^+`, shape `(d,)`.	required
`smoothed_cov_next`	`Float[Array, 'd d']`	Smoothed covariance `P_s` at `t^+`, shape `(d, d)`.	required
`solver`	`AbstractSolverStrategy \| None`	Optional solver strategy for the `P_+^{-1}` solve.	`None`

Returns:

Type	Description
`tuple[Float[Array, ' d'], Float[Array, 'd d']]`	Tuple `(mean, cov)` — the smoothed marginal of `x_t`.

Source code in src/gaussx/_gp/_interpolation.py

def rts_interpolate(
    A_fwd: Float[Array, "d d"],
    Q_fwd: Float[Array, "d d"],
    A_bwd: Float[Array, "d d"],
    Q_bwd: Float[Array, "d d"],
    filtered_mean_prev: Float[Array, " d"],
    filtered_cov_prev: Float[Array, "d d"],
    smoothed_mean_next: Float[Array, " d"],
    smoothed_cov_next: Float[Array, "d d"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " d"], Float[Array, "d d"]]:
    r"""Smoothed marginal at an unobserved ``t`` from filtered and smoothed states.

    For the model of `conditional_interpolate` (no observation at ``t``,
    ``t^- < t < t^+``), takes the **forward-filtered** state at ``t^-`` and
    the **smoothed** state at ``t^+`` — what `kalman_filter` and
    `rts_smoother` return — and performs one Rauch-Tung-Striebel backward
    step through ``t``:

    $$
    \begin{aligned}
    m_t^- &= A_{fwd} m_f, \quad
    P_t^- = A_{fwd} P_f A_{fwd}^\top + Q_{fwd}, \\
    m_+ &= A_{bwd} m_t^-, \quad
    P_+ = A_{bwd} P_t^- A_{bwd}^\top + Q_{bwd}, \\
    G &= P_t^- A_{bwd}^\top P_+^{-1}, \\
    m_t &= m_t^- + G (m_s - m_+), \quad
    P_t = P_t^- + G (P_s - P_+) G^\top .
    \end{aligned}
    $$

    This is exact: given ``x_{t^+}``, ``x_t`` is independent of the data at
    and after ``t^+``, so
    ``p(x_t | y) = \int p(x_t | x_{t^+}, y_{\le t^-}) p(x_{t^+} | y) dx_{t^+}``.

    Args:
        A_fwd: Forward transition from ``t^-`` to ``t``, shape ``(d, d)``.
        Q_fwd: Forward process noise, shape ``(d, d)``.
        A_bwd: Transition from ``t`` to ``t^+``, shape ``(d, d)``.
        Q_bwd: Process noise from ``t`` to ``t^+``, shape ``(d, d)``.
        filtered_mean_prev: Filtered mean ``m_f`` at ``t^-``, shape ``(d,)``.
        filtered_cov_prev: Filtered covariance ``P_f`` at ``t^-``,
            shape ``(d, d)``.
        smoothed_mean_next: Smoothed mean ``m_s`` at ``t^+``, shape ``(d,)``.
        smoothed_cov_next: Smoothed covariance ``P_s`` at ``t^+``,
            shape ``(d, d)``.
        solver: Optional solver strategy for the ``P_+^{-1}`` solve.

    Returns:
        Tuple ``(mean, cov)`` — the smoothed marginal of ``x_t``.
    """
    m_pred = A_fwd @ filtered_mean_prev
    P_pred = symmetrize(
        einsum(A_fwd, filtered_cov_prev, A_fwd, "i j, j k, l k -> i l") + Q_fwd
    )
    m_plus = A_bwd @ m_pred
    P_plus = symmetrize(einsum(A_bwd, P_pred, A_bwd, "i j, j k, l k -> i l") + Q_bwd)
    P_plus_op = lx.MatrixLinearOperator(P_plus, lx.positive_semidefinite_tag)
    # G = P_pred A_bwd^T P_plus^{-1}, row by row (P_plus is symmetric).
    G = solve_rows(P_plus_op, einsum(P_pred, A_bwd, "i j, k j -> i k"), solver=solver)
    m = m_pred + G @ (smoothed_mean_next - m_plus)
    P = symmetrize(
        P_pred + einsum(G, smoothed_cov_next - P_plus, G, "i j, j k, l k -> i l")
    )
    return m, P

`kronecker_posterior_predictive(K_factors: list[lx.AbstractLinearOperator], y: Float[Array, ' N'], noise_var: float, grid_shape: tuple[int, ...], K_cross_factors: list[Float[Array, 'Ni_test Ni_train']], *, K_test_diag_factors: list[Float[Array, ' Ni_test']]) -> tuple[Float[Array, ' N_test'], Float[Array, ' N_test']]` ¶

Posterior mean and variance for a Kronecker GP at test points.

Uses the eigendecomposition trick: projects cross-covariances onto the eigenbasis and weights by inverse eigenvalues:

mu_* = K_{*f} (K_{ff} + sigma^2 I)^{-1} y
var_* = k_{**} - K_{*f} (K_{ff} + sigma^2 I)^{-1} K_{f*}

Both computed via per-factor eigendecomposition in O(sum n_i^3 + N + N_test) time.

Parameters:

Name	Type	Description	Default
`K_factors`	`list[AbstractLinearOperator]`	List of per-dimension training kernel operators.	required
`y`	`Float[Array, ' N']`	Observations, shape `(N,)` where `N = prod(grid_shape)`.	required
`noise_var`	`float`	Observation noise variance `sigma^2`.	required
`grid_shape`	`tuple[int, ...]`	Grid shape, e.g. `(n1, n2)`.	required
`K_cross_factors`	`list[Float[Array, 'Ni_test Ni_train']]`	Per-dimension cross-covariance matrices, each shape `(n_i_test, n_i_train)`.	required
`K_test_diag_factors`	`list[Float[Array, ' Ni_test']]`	Per-dimension prior diagonals at the test points, each shape `(n_i_test,)`.	required

Returns:

Type	Description
`tuple[Float[Array, ' N_test'], Float[Array, ' N_test']]`	Tuple `(mean, variance)` at test points.

Raises:

Type	Description
`ValueError`	If `grid_shape` does not match the factors or `y` (as in `kronecker_mll`), or a per-factor list has the wrong length.

Source code in src/gaussx/_gp/_kronecker_gp.py

def kronecker_posterior_predictive(
    K_factors: list[lx.AbstractLinearOperator],
    y: Float[Array, " N"],
    noise_var: float,
    grid_shape: tuple[int, ...],
    K_cross_factors: list[Float[Array, "Ni_test Ni_train"]],
    *,
    K_test_diag_factors: list[Float[Array, " Ni_test"]],
) -> tuple[Float[Array, " N_test"], Float[Array, " N_test"]]:
    r"""Posterior mean and variance for a Kronecker GP at test points.

    Uses the eigendecomposition trick: projects cross-covariances onto
    the eigenbasis and weights by inverse eigenvalues:

        mu_* = K_{*f} (K_{ff} + sigma^2 I)^{-1} y
        var_* = k_{**} - K_{*f} (K_{ff} + sigma^2 I)^{-1} K_{f*}

    Both computed via per-factor eigendecomposition in
    ``O(sum n_i^3 + N + N_test)`` time.

    Args:
        K_factors: List of per-dimension training kernel operators.
        y: Observations, shape ``(N,)`` where ``N = prod(grid_shape)``.
        noise_var: Observation noise variance ``sigma^2``.
        grid_shape: Grid shape, e.g. ``(n1, n2)``.
        K_cross_factors: Per-dimension cross-covariance matrices,
            each shape ``(n_i_test, n_i_train)``.
        K_test_diag_factors: Per-dimension prior diagonals at the test points,
            each shape ``(n_i_test,)``.

    Returns:
        Tuple ``(mean, variance)`` at test points.

    Raises:
        ValueError: If ``grid_shape`` does not match the factors or ``y``
            (as in `kronecker_mll`), or a per-factor list has the wrong length.
    """
    _validate_grid(K_factors, grid_shape, y.shape[0])
    if len(K_cross_factors) != len(K_factors):
        msg = "K_cross_factors must have one matrix per Kronecker factor"
        raise ValueError(msg)
    if len(K_test_diag_factors) != len(K_factors):
        msg = "K_test_diag_factors must have one vector per Kronecker factor"
        raise ValueError(msg)

    # Per-factor eigendecomposition of training kernels
    factor_eigs = [eig(K_i) for K_i in K_factors]
    all_vals = [vals for vals, _ in factor_eigs]
    all_vecs = [vecs for _, vecs in factor_eigs]

    # Combined training eigenvalues
    combined_vals = ft.reduce(jnp.kron, all_vals)
    inv_noisy_vals = 1.0 / (combined_vals + noise_var)

    # Rotate observations: alpha = Q^T y
    alpha = _kron_rotate(all_vecs, y, grid_shape)

    # Weights in eigenbasis: w = diag(1/(lambda + sigma^2)) Q^T y
    w = inv_noisy_vals * alpha

    # Per-factor cross-covariance projected onto eigenbasis: A_i = K_cross_i @ Q_i
    A_factors = [K_cross_factors[i] @ all_vecs[i] for i in range(len(K_factors))]

    # Posterior mean: mu_* = (A_1 kron A_2 kron ...) w
    # Mean: project back from eigenbasis
    mean = _kron_matvec(A_factors, w, grid_shape)

    # Variance: k_** - sum_j (A_j^2 / (lambda_j + sigma^2))
    # = k_** - (A_1 kron A_2 kron ...)^2 @ inv_noisy_vals element-wise
    A_sq_factors = [A**2 for A in A_factors]
    var_reduction = _kron_matvec(A_sq_factors, inv_noisy_vals, grid_shape)

    # Prior diagonal at test points: diag(K_**) = kron_i diag(K_i(test, test))
    K_test_prior = ft.reduce(jnp.kron, K_test_diag_factors)

    variance = jnp.clip(K_test_prior - var_reduction, 0.0)

    return mean, variance

`kronecker_mll(K_factors: list[lx.AbstractLinearOperator], y: Float[Array, ' N'], noise_var: float, grid_shape: tuple[int, ...]) -> Float[Array, '']` ¶

Exact marginal log-likelihood for a Kronecker-structured GP.

For a GP with covariance K = K_1 \otimes K_2 \otimes \ldots + sigma^2 I, computes the log marginal likelihood via per-factor eigendecomposition:

log p(y) = -0.5 * (y^T (K + sigma^2 I)^{-1} y
                   + log|K + sigma^2 I|
                   + N log(2 pi))

The Kronecker eigendecomposition avoids forming the full N x N matrix: if K_i = Q_i Lambda_i Q_i^T, the combined eigenvalues are the outer products of per-factor eigenvalues and the combined eigenvectors are the Kronecker product of per-factor eigenvectors.

Complexity: O(sum n_i^3 + N) instead of O(N^3) where N = prod n_i.

Parameters:

Name	Type	Description	Default
`K_factors`	`list[AbstractLinearOperator]`	List of per-dimension kernel operators. Each must be square and symmetric.	required
`y`	`Float[Array, ' N']`	Observations, shape `(N,)` where `N = prod(grid_shape)`.	required
`noise_var`	`float`	Observation noise variance `sigma^2`.	required
`grid_shape`	`tuple[int, ...]`	Shape of the grid, e.g. `(n1, n2)` for 2D.	required

Returns:

Type	Description
`Float[Array, '']`	Scalar log marginal likelihood.

Raises:

Type	Description
`ValueError`	If `grid_shape` does not have one entry per factor matching that factor's size, or `prod(grid_shape) != N`.

Source code in src/gaussx/_gp/_kronecker_gp.py

def kronecker_mll(
    K_factors: list[lx.AbstractLinearOperator],
    y: Float[Array, " N"],
    noise_var: float,
    grid_shape: tuple[int, ...],
) -> Float[Array, ""]:
    r"""Exact marginal log-likelihood for a Kronecker-structured GP.

    For a GP with covariance ``K = K_1 \otimes K_2 \otimes \ldots + sigma^2 I``,
    computes the log marginal likelihood via per-factor eigendecomposition:

        log p(y) = -0.5 * (y^T (K + sigma^2 I)^{-1} y
                           + log|K + sigma^2 I|
                           + N log(2 pi))

    The Kronecker eigendecomposition avoids forming the full ``N x N`` matrix:
    if ``K_i = Q_i Lambda_i Q_i^T``, the combined eigenvalues are the outer
    products of per-factor eigenvalues and the combined eigenvectors are the
    Kronecker product of per-factor eigenvectors.

    Complexity: ``O(sum n_i^3 + N)`` instead of ``O(N^3)`` where ``N = prod n_i``.

    Args:
        K_factors: List of per-dimension kernel operators. Each must be
            square and symmetric.
        y: Observations, shape ``(N,)`` where ``N = prod(grid_shape)``.
        noise_var: Observation noise variance ``sigma^2``.
        grid_shape: Shape of the grid, e.g. ``(n1, n2)`` for 2D.

    Returns:
        Scalar log marginal likelihood.

    Raises:
        ValueError: If ``grid_shape`` does not have one entry per factor
            matching that factor's size, or ``prod(grid_shape) != N``.
    """
    N = y.shape[0]
    _validate_grid(K_factors, grid_shape, N)

    # Per-factor eigendecomposition
    factor_eigs = [eig(K_i) for K_i in K_factors]
    all_vals = [vals for vals, _ in factor_eigs]
    all_vecs = [vecs for _, vecs in factor_eigs]

    # Combined eigenvalues: outer product of per-factor eigenvalues
    combined_vals = ft.reduce(jnp.kron, all_vals)  # (N,)
    noisy_vals = combined_vals + noise_var  # (N,)

    # Rotate data into eigenbasis: alpha = Q^T y
    # Q = Q_1 kron Q_2 kron ..., so Q^T y can be computed factor-by-factor
    alpha = _kron_rotate(all_vecs, y, grid_shape)

    # MLL in eigenbasis: all operations are O(N)
    data_fit = jnp.sum(alpha**2 / noisy_vals)
    log_det = jnp.sum(jnp.log(noisy_vals))
    const = N * jnp.log(2.0 * jnp.pi)

    return -0.5 * (data_fit + log_det + const)

Pathwise sampling¶

Matheron's rule turns joint prior draws \((a, b)\) into posterior draws: \(a + \mathrm{Cov}(a,b)\,\mathrm{Cov}(b,b)^{-1}(\beta - b)\).

Structured linear algebra and Gaussian primitives for JAX.

`matheron_update(prior_sample_target: Float[Array, 'S N_star'], prior_sample_conditioning: Float[Array, 'S M'], observed_value: Float[Array, ' M'], cross_covariance: lx.AbstractLinearOperator, conditioning_covariance: lx.AbstractLinearOperator, *, solver: AbstractSolveStrategy | None = None) -> Float[Array, 'S N_star']` ¶

Posterior samples via Matheron's-rule correction.

Given joint prior draws (a, b) and an observed conditioning value β, Matheron's rule samples from a | b = β by applying

\[ a + \operatorname{Cov}(a, b)\operatorname{Cov}(b, b)^{-1}(β - b). \]

This helper keeps both covariance arguments as lineax operators, so the conditioning solve uses the existing GaussX structural dispatch and the target correction is a rectangular matvec.

Parameters:

Name	Type	Description	Default
`prior_sample_target`	`Float[Array, 'S N_star']`	Prior samples at target points, shape `(S, N_star)`.	required
`prior_sample_conditioning`	`Float[Array, 'S M']`	Joint prior samples at conditioning points, shape `(S, M)`.	required
`observed_value`	`Float[Array, ' M']`	Observed conditioning value, shape `(M,)`.	required
`cross_covariance`	`AbstractLinearOperator`	Cross-covariance operator `Cov(a, b)`, shape `(N_star, M)`.	required
`conditioning_covariance`	`AbstractLinearOperator`	Conditioning covariance operator `Cov(b, b)`, shape `(M, M)`.	required
`solver`	`AbstractSolveStrategy \| None`	Optional solver strategy for the conditioning solve (e.g. `gaussx.CGSolver`, `gaussx.BBMMSolver`). When `None`, routes through structural dispatch on `conditioning_covariance`.	`None`

Returns:

Type	Description
`Float[Array, 'S N_star']`	Corrected posterior samples, shape `(S, N_star)`.

Source code in src/gaussx/_gp/_matheron.py

def matheron_update(
    prior_sample_target: Float[Array, "S N_star"],
    prior_sample_conditioning: Float[Array, "S M"],
    observed_value: Float[Array, " M"],
    cross_covariance: lx.AbstractLinearOperator,
    conditioning_covariance: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, "S N_star"]:
    r"""Posterior samples via Matheron's-rule correction.

    Given joint prior draws ``(a, b)`` and an observed conditioning value
    ``β``, Matheron's rule samples from ``a | b = β`` by applying

    $$
    a + \operatorname{Cov}(a, b)\operatorname{Cov}(b, b)^{-1}(β - b).
    $$

    This helper keeps both covariance arguments as lineax operators, so the
    conditioning solve uses the existing GaussX structural dispatch and the
    target correction is a rectangular matvec.

    Args:
        prior_sample_target: Prior samples at target points, shape
            ``(S, N_star)``.
        prior_sample_conditioning: Joint prior samples at conditioning
            points, shape ``(S, M)``.
        observed_value: Observed conditioning value, shape ``(M,)``.
        cross_covariance: Cross-covariance operator ``Cov(a, b)``, shape
            ``(N_star, M)``.
        conditioning_covariance: Conditioning covariance operator
            ``Cov(b, b)``, shape ``(M, M)``.
        solver: Optional solver strategy for the conditioning solve
            (e.g. `gaussx.CGSolver`, `gaussx.BBMMSolver`).
            When ``None``, routes through structural dispatch on
            ``conditioning_covariance``.

    Returns:
        Corrected posterior samples, shape ``(S, N_star)``.
    """
    prior_sample_target = jnp.asarray(prior_sample_target)
    prior_sample_conditioning = jnp.asarray(prior_sample_conditioning)
    observed_value = jnp.asarray(observed_value)
    dtype = jnp.result_type(
        prior_sample_target,
        prior_sample_conditioning,
        observed_value,
    )
    # Promote sample arrays and observations together so the residual,
    # structured solve, and matvec all use the same precision.
    prior_sample_target = prior_sample_target.astype(dtype)
    prior_sample_conditioning = prior_sample_conditioning.astype(dtype)
    observed_value = observed_value.astype(dtype)

    if prior_sample_target.ndim != 2:
        raise ValueError("prior_sample_target must have shape (S, N_star).")
    if prior_sample_conditioning.ndim != 2:
        raise ValueError("prior_sample_conditioning must have shape (S, M).")
    if observed_value.ndim != 1:
        raise ValueError("observed_value must have shape (M,).")
    if prior_sample_target.shape[0] != prior_sample_conditioning.shape[0]:
        raise ValueError("prior samples must have the same sample dimension S.")

    num_conditioning = prior_sample_conditioning.shape[1]
    num_target = prior_sample_target.shape[1]
    if observed_value.shape[0] != num_conditioning:
        raise ValueError("observed_value must match the conditioning dimension M.")
    if conditioning_covariance.in_size() != num_conditioning:
        raise ValueError(
            "conditioning_covariance input size must match the conditioning "
            "dimension M."
        )
    if conditioning_covariance.out_size() != num_conditioning:
        raise ValueError(
            "conditioning_covariance output size must match the conditioning "
            "dimension M."
        )
    if cross_covariance.in_size() != num_conditioning:
        raise ValueError(
            "cross_covariance input size must match the conditioning dimension M."
        )
    if cross_covariance.out_size() != num_target:
        raise ValueError(
            "cross_covariance output size must match the target dimension N_star."
        )

    residuals = observed_value[None, :] - prior_sample_conditioning
    if solver is not None:
        solve_one = lambda r: dispatch_solve(conditioning_covariance, r, solver)
    else:
        solve_one = lambda r: solve(conditioning_covariance, r)
    solves = jax.vmap(solve_one)(residuals)
    corrections = jax.vmap(cross_covariance.mv)(solves)
    return prior_sample_target + corrections

Whitening¶

The whitened parameterization \(u = Lv\), \(v \sim \mathcal{N}(0, I)\) that keeps sparse-variational optimization well-conditioned, and the whitened SVGP predictive that consumes it.

Structured linear algebra and Gaussian primitives for JAX.

`whiten_covariance = unwhiten_covariance` `module-attribute` ¶

`unwhiten(m_tilde: Float[Array, ' M'], L: lx.AbstractLinearOperator) -> Float[Array, ' M']` ¶

Unwhiten variational mean: m = L @ m_tilde.

Parameters:

Name	Type	Description	Default
`m_tilde`	`Float[Array, ' M']`	Whitened mean vector, shape `(M,)`.	required
`L`	`AbstractLinearOperator`	Cholesky factor, shape `(M, M)`.	required

Returns:

Type	Description
`Float[Array, ' M']`	Unwhitened mean m, shape `(M,)`.

Source code in src/gaussx/_gp/_unwhiten.py

def unwhiten(
    m_tilde: Float[Array, " M"],
    L: lx.AbstractLinearOperator,
) -> Float[Array, " M"]:
    """Unwhiten variational mean: ``m = L @ m_tilde``.

    Args:
        m_tilde: Whitened mean vector, shape ``(M,)``.
        L: Cholesky factor, shape ``(M, M)``.

    Returns:
        Unwhitened mean m, shape ``(M,)``.
    """
    return L.mv(m_tilde)

`unwhiten_covariance(L: lx.AbstractLinearOperator, S_tilde: lx.AbstractLinearOperator) -> lx.MatrixLinearOperator` ¶

Unwhiten variational covariance: S = L S̃ Lᵀ.

Delegates to cov_transform.

Parameters:

Name	Type	Description	Default
`L`	`AbstractLinearOperator`	Cholesky factor, shape `(M, M)`.	required
`S_tilde`	`AbstractLinearOperator`	Whitened variational covariance, shape `(M, M)`.	required

Returns:

Type	Description
`MatrixLinearOperator`	Unwhitened covariance operator S.

Source code in src/gaussx/_gp/_unwhiten.py

def unwhiten_covariance(
    L: lx.AbstractLinearOperator,
    S_tilde: lx.AbstractLinearOperator,
) -> lx.MatrixLinearOperator:
    """Unwhiten variational covariance: S = L S̃ Lᵀ.

    Delegates to `cov_transform`.

    Args:
        L: Cholesky factor, shape ``(M, M)``.
        S_tilde: Whitened variational covariance, shape ``(M, M)``.

    Returns:
        Unwhitened covariance operator S.
    """
    return cov_transform(L.as_matrix(), S_tilde)

`whitened_svgp_predict(K_zz_op: lx.AbstractLinearOperator, K_xz: Float[Array, 'N M'], u_mean: Float[Array, ' M'], u_chol: Float[Array, 'M M'], K_xx_diag: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']]` ¶

Whitened SVGP prediction: mean and variance at test points.

Computes the predictive mean and variance for a sparse variational GP using the whitened parameterization:

L_{zz} = cholesky(K_{zz})
A = L_{zz}^{-1} K_{zx}           (triangular solve)
f_{loc} = A^T u_{mean}
Q_{xx} = sum(A^2, axis=0)         (prior variance reduction)
W = u_{chol}^T A
S_{contrib} = sum(W^2, axis=0)    (posterior variance contribution)
f_{var} = K_{xx,diag} - Q_{xx} + S_{contrib}

Parameters:

Name	Type	Description	Default
`K_zz_op`	`AbstractLinearOperator`	Inducing-point covariance operator, shape `(M, M)`.	required
`K_xz`	`Float[Array, 'N M']`	Cross-covariance matrix, shape `(N, M)`.	required
`u_mean`	`Float[Array, ' M']`	Whitened variational mean, shape `(M,)`.	required
`u_chol`	`Float[Array, 'M M']`	Whitened variational Cholesky factor, shape `(M, M)`. Lower-triangular matrix such that the variational covariance in whitened space is `u_chol @ u_chol^T`.	required
`K_xx_diag`	`Float[Array, ' N']`	Prior diagonal variances at test points, shape `(N,)`.	required

Returns:

Type	Description
`Float[Array, ' N']`	Tuple `(f_loc, f_var)` — predictive mean shape `(N,)` and
`Float[Array, ' N']`	predictive variance shape `(N,)`.

Source code in src/gaussx/_gp/_svgp.py

def whitened_svgp_predict(
    K_zz_op: lx.AbstractLinearOperator,
    K_xz: Float[Array, "N M"],
    u_mean: Float[Array, " M"],
    u_chol: Float[Array, "M M"],
    K_xx_diag: Float[Array, " N"],
) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
    r"""Whitened SVGP prediction: mean and variance at test points.

    Computes the predictive mean and variance for a sparse variational
    GP using the whitened parameterization:

        L_{zz} = cholesky(K_{zz})
        A = L_{zz}^{-1} K_{zx}           (triangular solve)
        f_{loc} = A^T u_{mean}
        Q_{xx} = sum(A^2, axis=0)         (prior variance reduction)
        W = u_{chol}^T A
        S_{contrib} = sum(W^2, axis=0)    (posterior variance contribution)
        f_{var} = K_{xx,diag} - Q_{xx} + S_{contrib}

    Args:
        K_zz_op: Inducing-point covariance operator, shape ``(M, M)``.
        K_xz: Cross-covariance matrix, shape ``(N, M)``.
        u_mean: Whitened variational mean, shape ``(M,)``.
        u_chol: Whitened variational Cholesky factor, shape ``(M, M)``.
            Lower-triangular matrix such that the variational covariance
            in whitened space is ``u_chol @ u_chol^T``.
        K_xx_diag: Prior diagonal variances at test points, shape ``(N,)``.

    Returns:
        Tuple ``(f_loc, f_var)`` — predictive mean shape ``(N,)`` and
        predictive variance shape ``(N,)``.
    """
    L_zz = cholesky(K_zz_op)

    # A = L_zz^{-1} K_xz^T  -> shape (M, N)
    # Solve L_zz @ A_col = K_xz^T_col for each column of K_xzᵀ

    K_zx = K_xz.T  # (M, N)
    A = solve_columns(L_zz, K_zx)

    # Predictive mean: f_loc = A^T @ u_mean = K_xz @ L_zz^{-T} @ u_mean
    f_loc = A.T @ u_mean

    # Prior variance reduction: Q_xx = Σₘ Aₘₙ²
    Q_xx = reduce(A**2, "M N -> N", "sum")

    # Posterior variance contribution: W = u_cholᵀ @ A, S = Σₖ Wₖₙ²
    W = u_chol.T @ A
    S_contrib = reduce(W**2, "K N -> N", "sum")

    # Predictive variance
    f_var = jnp.clip(K_xx_diag - Q_xx + S_contrib, 0.0)

    return f_loc, f_var

`svgp_variance_adjustment(K_zz_op: lx.AbstractLinearOperator, S_u: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator` ¶

Compute the SVGP variance adjustment operator.

Builds the operator Q = K_{zz}^{-1} S_u K_{zz}^{-1} - K_{zz}^{-1} which appears in every sparse GP predictive variance computation:

Var[f_*] = k_{**} - k_{*z} (K_{zz}^{-1} - Q) k_{z*}

The returned value is exposed as a linear operator, but the current implementation materializes dense (M, M) intermediates while building it.

Parameters:

Name	Type	Description	Default
`K_zz_op`	`AbstractLinearOperator`	Inducing-point covariance operator, shape `(M, M)`.	required
`S_u`	`AbstractLinearOperator`	Variational covariance operator, shape `(M, M)`.	required

Returns:

Type	Description
`AbstractLinearOperator`	Operator `Q` of shape `(M, M)` such that
`AbstractLinearOperator`	`Q @ v = K_{zz}^{-1} S_u K_{zz}^{-1} v - K_{zz}^{-1} v`.

Source code in src/gaussx/_gp/_svgp_variance.py

def svgp_variance_adjustment(
    K_zz_op: lx.AbstractLinearOperator,
    S_u: lx.AbstractLinearOperator,
) -> lx.AbstractLinearOperator:
    r"""Compute the SVGP variance adjustment operator.

    Builds the operator ``Q = K_{zz}^{-1} S_u K_{zz}^{-1} - K_{zz}^{-1}``
    which appears in every sparse GP predictive variance computation:

        Var[f_*] = k_{**} - k_{*z} (K_{zz}^{-1} - Q) k_{z*}

    The returned value is exposed as a linear operator, but the current
    implementation materializes dense ``(M, M)`` intermediates while building
    it.

    Args:
        K_zz_op: Inducing-point covariance operator, shape ``(M, M)``.
        S_u: Variational covariance operator, shape ``(M, M)``.

    Returns:
        Operator ``Q`` of shape ``(M, M)`` such that
        ``Q @ v = K_{zz}^{-1} S_u K_{zz}^{-1} v - K_{zz}^{-1} v``.
    """
    K_inv = inv(K_zz_op)
    K_inv_S = lx.MatrixLinearOperator(
        _compose_dense(K_inv, S_u),
    )
    # Q = K_inv @ S_u @ K_inv - K_inv
    # Build as (K_inv @ S_u - I) @ K_inv
    import jax.numpy as jnp

    M = K_zz_op.out_structure().shape[0]
    K_inv_S_minus_I = lx.AddLinearOperator(
        K_inv_S,
        lx.DiagonalLinearOperator(-jnp.ones(M)),
    )
    Q_dense = _compose_dense(K_inv_S_minus_I, K_inv)
    return lx.MatrixLinearOperator(Q_dense)

Variational bounds & KL¶

ELBOs for Gaussian and Monte-Carlo variational families, the collapsed (Titsias) sparse bound, and the Gaussian-to-Gaussian KL term.

Structured linear algebra and Gaussian primitives for JAX.

`variational_elbo_gaussian(y: Float[Array, ' N'], f_loc: Float[Array, ' N'], f_var: Float[Array, ' N'], noise_var: float, kl: Float[Array, '']) -> Float[Array, '']` ¶

Uncollapsed (SVGP) ELBO for a Gaussian likelihood.

The stochastic-variational bound of Hensman et al. (2013): the variational distribution q(u) is kept explicit, so the caller supplies its marginals f_loc, f_var and the KL(q(u) || p(u)) term. Computes:

ELBO = E_q[log p(y|f)] - KL(q||p)

where the expected log-likelihood under a Gaussian variational distribution with diagonal variance has the closed form:

E_q[log N(y|f, sigma^2 I)]
    = -0.5 * N * log(2 pi sigma^2)
      -0.5 / sigma^2 * (||y - f_loc||^2 + sum(f_var))

Parameters:

Name	Type	Description	Default
`y`	`Float[Array, ' N']`	Observations, shape `(N,)`.	required
`f_loc`	`Float[Array, ' N']`	Variational mean, shape `(N,)`.	required
`f_var`	`Float[Array, ' N']`	Variational marginal variances, shape `(N,)`.	required
`noise_var`	`float`	Observation noise variance (scalar).	required
`kl`	`Float[Array, '']`	KL divergence term `KL(q \|\| p)` (scalar).	required

Returns:

Type	Description
`Float[Array, '']`	Scalar ELBO value.

`variational_elbo_mc(log_likelihood_fn: Callable[[Float[Array, ' N']], Float[Array, '']], f_samples: Float[Array, 'S N'], kl: Float[Array, '']) -> Float[Array, '']` ¶

Monte Carlo ELBO for non-conjugate likelihoods.

Computes:

ELBO = (1/S) sum_s log p(y|f_s) - KL(q||p)

where f_s ~ q(f) are samples from the variational distribution. Supports any likelihood (Poisson, Bernoulli, Pareto, etc.).

Parameters:

Name	Type	Description	Default
`log_likelihood_fn`	`Callable[[Float[Array, ' N']], Float[Array, '']]`	Function mapping latent samples to scalar log-likelihood. Signature `(N,) -> scalar`.	required
`f_samples`	`Float[Array, 'S N']`	Samples from the variational posterior, shape `(S, N)` where S is the number of samples.	required
`kl`	`Float[Array, '']`	KL divergence term `KL(q \|\| p)` (scalar).	required

Returns:

Type	Description
`Float[Array, '']`	Scalar ELBO value.

Source code in src/gaussx/_gp/_elbo.py

def variational_elbo_mc(
    log_likelihood_fn: Callable[[Float[Array, " N"]], Float[Array, ""]],
    f_samples: Float[Array, "S N"],
    kl: Float[Array, ""],
) -> Float[Array, ""]:
    """Monte Carlo ELBO for non-conjugate likelihoods.

    Computes:

        ELBO = (1/S) sum_s log p(y|f_s) - KL(q||p)

    where ``f_s ~ q(f)`` are samples from the variational distribution.
    Supports any likelihood (Poisson, Bernoulli, Pareto, etc.).

    Args:
        log_likelihood_fn: Function mapping latent samples to scalar
            log-likelihood. Signature ``(N,) -> scalar``.
        f_samples: Samples from the variational posterior, shape
            ``(S, N)`` where S is the number of samples.
        kl: KL divergence term ``KL(q || p)`` (scalar).

    Returns:
        Scalar ELBO value.
    """
    import jax

    ell = jnp.mean(jax.vmap(log_likelihood_fn)(f_samples))
    return ell - kl

`collapsed_elbo(y: Float[Array, ' N'], K_xx_diag: Float[Array, ' N'], K_xz: Float[Array, 'N M'], K_zz: Float[Array, 'M M'] | lx.AbstractLinearOperator, noise_var: float, *, jitter: float = 1e-06, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']` ¶

Collapsed ELBO (Titsias bound) for sparse GP regression.

Computes the variational lower bound on the log marginal likelihood using the matrix determinant lemma for O(NM² + M³) cost:

ELBO = log 𝒩(y | 0, Q_ff + σ²I) − ½σ⁻² tr(K_ff − Q_ff)

where Q_ff = K_xz K_zz⁻¹ K_xzᵀ is the Nyström approximation. With L L^T = K_zz and V = L⁻¹ K_xzᵀ (structured triangular solves), B = I + σ⁻² V Vᵀ is the only dense (M, M) factorisation.

Parameters:

Name	Type	Description	Default
`y`	`Float[Array, ' N']`	Observations, shape `(N,)`.	required
`K_xx_diag`	`Float[Array, ' N']`	Diagonal of the full kernel matrix K_ff, shape `(N,)`. (Formerly `K_diag`; the old keyword is deprecated.)	required
`K_xz`	`Float[Array, 'N M']`	Cross-covariance between data and inducing points, shape `(N, M)`.	required
`K_zz`	`Float[Array, 'M M'] \| AbstractLinearOperator`	Inducing point kernel matrix, `(M, M)` array or PSD operator. A `Kronecker` / `BlockDiag` operator is factorised per factor.	required
`noise_var`	`float`	Observation noise variance σ² (scalar).	required
`jitter`	`float`	Diagonal jitter added to `K_zz` before the Cholesky factorisation, for a dense array, `lineax.MatrixLinearOperator` or `lineax.DiagonalLinearOperator`. It is not added to other structured operators (it would destroy their structure): include any jitter in the operator itself, e.g. per factor.	`1e-06`
`solver`	`AbstractSolverStrategy \| None`	Deprecated and ignored (the Cholesky factorisations take no solver); passing one warns. Removed in gaussx 0.7.0.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar ELBO value.

`gauss_kl(q_mu: Float[Array, 'M R'], q_sqrt: Float[Array, 'R M M'] | Float[Array, 'M R'], K: Float[Array, 'M M'] | lx.AbstractLinearOperator | None = None, *, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']` ¶

KL(first ‖ second): KL[q(u) || p(u)] with p(u) = N(0, K).

Cholesky-parameterised variant of gaussian_kl designed for GP/SVGP models. The Cholesky representation avoids explicit covariance matrix construction and supports both full and diagonal q_sqrt. For lineax-operator covariances, use gaussian_kl instead.

Computes the KL divergence where:

q = N(q_mu, q_sqrt @ q_sqrt^T)
p = N(0, K) or p = N(0, I) if K is None (white prior)

Handles both full and diagonal q_sqrt:

Full q_sqrt: shape (R, M, M) — lower-triangular Cholesky factors of the variational covariance per output dimension.
Diagonal q_sqrt: shape (M, R) — diagonal standard deviations.

Parameters:

Name	Type	Description	Default
`q_mu`	`Float[Array, 'M R']`	Variational mean, shape `(M, R)`.	required
`q_sqrt`	`Float[Array, 'R M M'] \| Float[Array, 'M R']`	Variational Cholesky factor or diagonal std devs.	required
`K`	`Float[Array, 'M M'] \| AbstractLinearOperator \| None`	Prior covariance, `(M, M)` array or PSD operator (a `Kronecker` / `BlockDiag` operator is factorised per factor). If `None`, uses white prior (identity).	`None`
`solver`	`AbstractSolverStrategy \| None`	Deprecated and ignored (the Cholesky factorisation takes no solver); passing one warns. Removed in gaussx 0.7.0.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar KL divergence summed over all `R` output dimensions.

Cross-validation & diagnostics¶

LOVE-style cached predictive variances and closed-form leave-one-out cross-validation from a single factorization.

Structured linear algebra and Gaussian primitives for JAX.

`LOVECache` ¶

Bases: Module

Cached rank-k Lanczos approximation of K^{-1}.

k steps of symmetric Lanczos from a start vector v_0 build an orthonormal basis of the Krylov space \mathcal{K}_k(K, v_0) = \mathrm{span}\{v_0, K v_0, \dots, K^{k-1} v_0\} and the tridiagonal T_k = Q_k^\top K Q_k. With the eigendecomposition T_k = S \Theta S^\top, the cache stores the Ritz vectors Q = Q_k S and the inverse Ritz values 1 / \theta_i, so that

\[ K^{-1} \approx Q_k T_k^{-1} Q_k^\top = Q \Theta^{-1} Q^\top . \]

Ritz pairs approximate eigenpairs of K only once they have converged (first at the extremes of the spectrum); for k < N they are not the eigenpairs of K. K^{-1} - Q \Theta^{-1} Q^\top is positive semi-definite, so the approximation never over-estimates a quadratic form k_*^\top K^{-1} k_* (see love_variance).

Attributes:

Name	Type	Description
`Q`	`Float[Array, 'N k']`	Ritz vectors (an orthonormal basis of the Krylov space), shape `(N, k)`.
`inv_eigvals`	`Float[Array, ' k']`	Inverse Ritz values `1 / \theta_i`, shape `(k,)`.

Source code in src/gaussx/_gp/_love.py

class LOVECache(eqx.Module):
    r"""Cached rank-``k`` Lanczos approximation of ``K^{-1}``.

    ``k`` steps of symmetric Lanczos from a start vector ``v_0`` build an
    orthonormal basis of the Krylov space
    ``\mathcal{K}_k(K, v_0) = \mathrm{span}\{v_0, K v_0, \dots, K^{k-1} v_0\}``
    and the tridiagonal ``T_k = Q_k^\top K Q_k``. With the eigendecomposition
    ``T_k = S \Theta S^\top``, the cache stores the Ritz vectors
    ``Q = Q_k S`` and the inverse Ritz values ``1 / \theta_i``, so that

    $$
    K^{-1} \approx Q_k T_k^{-1} Q_k^\top = Q \Theta^{-1} Q^\top .
    $$

    Ritz pairs approximate eigenpairs of ``K`` only once they have converged
    (first at the extremes of the spectrum); for ``k < N`` they are not the
    eigenpairs of ``K``. ``K^{-1} - Q \Theta^{-1} Q^\top`` is positive
    semi-definite, so the approximation never over-estimates a quadratic form
    ``k_*^\top K^{-1} k_*`` (see `love_variance`).

    Attributes:
        Q: Ritz vectors (an orthonormal basis of the Krylov space),
            shape ``(N, k)``.
        inv_eigvals: Inverse Ritz values ``1 / \theta_i``, shape ``(k,)``.
    """

    Q: Float[Array, "N k"]
    inv_eigvals: Float[Array, " k"]

`LOOResult` ¶

Bases: Module

Result of leave-one-out cross-validation.

Attributes:

Name	Type	Description
`loo_log_likelihood`	`Float[Array, '']`	Scalar LOO-CV log-likelihood.
`loo_means`	`Float[Array, ' N']`	Per-point LOO predictive means, shape `(N,)`.
`loo_variances`	`Float[Array, ' N']`	Per-point LOO predictive variances, shape `(N,)`.

Source code in src/gaussx/_gp/_loo.py

class LOOResult(eqx.Module):
    """Result of leave-one-out cross-validation.

    Attributes:
        loo_log_likelihood: Scalar LOO-CV log-likelihood.
        loo_means: Per-point LOO predictive means, shape ``(N,)``.
        loo_variances: Per-point LOO predictive variances, shape ``(N,)``.
    """

    loo_log_likelihood: Float[Array, ""]
    loo_means: Float[Array, " N"]
    loo_variances: Float[Array, " N"]

`love_cache(K_op: lx.AbstractLinearOperator, lanczos_order: int = 50, key: jax.Array | None = None, *, initial_vector: Float[Array, ' N'] | None = None) -> LOVECache` ¶

Precompute a Lanczos approximation of K^{-1} for fast variance.

Runs k = min(lanczos_order, N) steps of symmetric Lanczos (with full reorthogonalisation) on K from the start vector v_0 and returns the Galerkin approximation

\[ K^{-1} \approx Q_k T_k^{-1} Q_k^\top = K^{-1/2} P_k K^{-1/2}, \]

where P_k is the orthogonal projector onto K^{1/2} \mathcal{K}_k(K, v_0). Because P_k \preceq I the approximation is a lower bound on K^{-1} in the Loewner order: the quadratic form returned by love_variance is biased low and the GP predictive variance k_{**} - k_*^\top K^{-1} k_* built from it is biased high, until the Krylov space captures every direction in which k_* has weight. At k = N the result is exact.

Choosing lanczos_order: the required k is set by the number of eigenvalues of K well above the noise level (the kernel's effective rank at the noise variance), not by N. A smooth kernel with small noise can need noticeably more steps than its numerical rank suggests. Check convergence with love_residual / love_variance_error_bound on representative test points.

Pseudocode:

v0 = initial_vector or normal(key, (N,))
Q_k, T_k = lanczos(K, v0, k)            # k matvecs with K
theta, S = eigh(T_k)
Q = Q_k @ S;  inv_eigvals = 1 / theta

Once cached, each test point costs O(Nk) instead of an O(N^2) solve (Pleiss et al., 2018).

Parameters:

Name	Type	Description	Default
`K_op`	`AbstractLinearOperator`	Training kernel operator `K`, shape `(N, N)`. Must be symmetric positive definite.	required
`lanczos_order`	`int`	Number of Lanczos steps `k` (rank of the approximation), clamped to `N`. Default `50`.	`50`
`key`	`Array \| None`	PRNG key for a random start vector. If `None` (and no `initial_vector` is given), uses `jax.random.PRNGKey(0)`.	`None`
`initial_vector`	`Float[Array, ' N'] \| None`	Deterministic start vector `v_0`, shape `(N,)`, e.g. the training targets `y` (the right-hand side of the mean solve, as in LOVE). Takes precedence over `key`. A zero vector is replaced by the all-ones vector.	`None`

A lineax.DiagonalLinearOperator, BlockDiag, Kronecker or KroneckerSum K is not run through Lanczos: its exact structured eigendecomposition is truncated to the top k eigenpairs (key and initial_vector are unused). That is also a Galerkin approximation, with the same one-signed bias.

Returns:

Type	Description
`LOVECache`	A `LOVECache` object.

Source code in src/gaussx/_gp/_love.py

def love_cache(
    K_op: lx.AbstractLinearOperator,
    lanczos_order: int = 50,
    key: jax.Array | None = None,
    *,
    initial_vector: Float[Array, " N"] | None = None,
) -> LOVECache:
    r"""Precompute a Lanczos approximation of ``K^{-1}`` for fast variance.

    Runs ``k = min(lanczos_order, N)`` steps of symmetric Lanczos (with full
    reorthogonalisation) on ``K`` from the start vector ``v_0`` and returns
    the Galerkin approximation

    $$
    K^{-1} \approx Q_k T_k^{-1} Q_k^\top
    = K^{-1/2} P_k K^{-1/2},
    $$

    where ``P_k`` is the orthogonal projector onto
    ``K^{1/2} \mathcal{K}_k(K, v_0)``. Because ``P_k \preceq I`` the
    approximation is a lower bound on ``K^{-1}`` in the Loewner order: the
    quadratic form returned by `love_variance` is biased low and the GP
    predictive variance ``k_{**} - k_*^\top K^{-1} k_*`` built from it is
    biased **high**, until the Krylov space captures every direction in
    which ``k_*`` has weight. At ``k = N`` the result is exact.

    Choosing ``lanczos_order``: the required ``k`` is set by the number of
    eigenvalues of ``K`` well above the noise level (the kernel's effective
    rank at the noise variance), not by ``N``. A smooth kernel with small
    noise can need noticeably more steps than its numerical rank suggests.
    Check convergence with `love_residual` / `love_variance_error_bound` on
    representative test points.

    Pseudocode:

    ```text
    v0 = initial_vector or normal(key, (N,))
    Q_k, T_k = lanczos(K, v0, k)            # k matvecs with K
    theta, S = eigh(T_k)
    Q = Q_k @ S;  inv_eigvals = 1 / theta
    ```

    Once cached, each test point costs ``O(Nk)`` instead of an ``O(N^2)``
    solve (Pleiss et al., 2018).

    Args:
        K_op: Training kernel operator ``K``, shape ``(N, N)``. Must be
            symmetric positive definite.
        lanczos_order: Number of Lanczos steps ``k`` (rank of the
            approximation), clamped to ``N``. Default ``50``.
        key: PRNG key for a random start vector. If ``None`` (and no
            ``initial_vector`` is given), uses ``jax.random.PRNGKey(0)``.
        initial_vector: Deterministic start vector ``v_0``, shape ``(N,)``,
            e.g. the training targets ``y`` (the right-hand side of the mean
            solve, as in LOVE). Takes precedence over ``key``. A zero vector
            is replaced by the all-ones vector.

    A `lineax.DiagonalLinearOperator`, `BlockDiag`, `Kronecker` or
    `KroneckerSum` ``K`` is not run through Lanczos: its exact structured
    eigendecomposition is truncated to the top ``k`` eigenpairs (``key``
    and ``initial_vector`` are unused). That is also a Galerkin
    approximation, with the same one-signed bias.

    Returns:
        A `LOVECache` object.
    """
    n = K_op.in_size()
    rank = min(lanczos_order, n)
    if isinstance(K_op, _EXACT_EIG):
        # Exact (structured) eigendecomposition, top-``rank`` pairs: no
        # Krylov breakdown on repeated eigenvalues (e.g. a diagonal K).
        root = root_inv_decomposition(K_op, rank=rank, method="lanczos", key=key).root
        inv_eigvals = reduce(root**2, "n k -> k", "sum")  # ||R[:, i]||² = 1/λᵢ
        scale = jnp.sqrt(_safe_psd(inv_eigvals))
        return LOVECache(
            Q=einx.divide("n k, k -> n k", root, scale), inv_eigvals=inv_eigvals
        )

    dtype = K_op.in_structure().dtype
    if initial_vector is None:
        if key is None:
            key = jr.PRNGKey(0)
        v0 = jr.normal(key, (n,), dtype=dtype)
    else:
        v0 = jnp.asarray(initial_vector, dtype=dtype)
        # A zero start vector would make the first Lanczos normalisation 0/0.
        v0 = jnp.where(jnp.linalg.norm(v0) > 0, v0, jnp.ones_like(v0))

    tridiag = matfree.decomp.tridiag_sym(rank, reortho="full")
    # matfree returns Ritz values (k,) and Ritz vectors (k, N).
    vals, vecs = matfree.eig.eigh_partial(tridiag)(K_op.mv, v0)
    vals, Q = _top_real_eigenpairs(vals, rearrange(vecs, "k n -> n k"), rank)
    return LOVECache(Q=Q, inv_eigvals=1.0 / _safe_psd(vals))

`love_variance(cache: LOVECache, K_star_row: Float[Array, ' N']) -> Float[Array, '']` ¶

Approximate k_*^\top K^{-1} k_* from a LOVE cache.

Computes, in O(Nk),

\[ k_*^\top Q \Theta^{-1} Q^\top k_* = \sum_i (q_i^\top k_*)^2 / \theta_i \;\le\; k_*^\top K^{-1} k_* . \]

The inequality is the one-signed bias of the Lanczos approximation (see love_cache): the GP predictive variance

\[ \sigma_*^2 = k_{**} - \mathrm{love\_variance}(\mathrm{cache}, k_*) \]

is an upper bound on the exact predictive variance, tight once the cache has converged for this k_* (check with love_residual).

Parameters:

Name	Type	Description	Default
`cache`	`LOVECache`	A `LOVECache` from `love_cache`.	required
`K_star_row`	`Float[Array, ' N']`	Cross-covariance vector `k(X_{train}, x_*)`, shape `(N,)`.	required

Returns:

Type	Description
`Float[Array, '']`	Scalar approximation of `k_^T K^{-1} k_`.

Source code in src/gaussx/_gp/_love.py

def love_variance(
    cache: LOVECache,
    K_star_row: Float[Array, " N"],
) -> Float[Array, ""]:
    r"""Approximate ``k_*^\top K^{-1} k_*`` from a LOVE cache.

    Computes, in ``O(Nk)``,

    $$
    k_*^\top Q \Theta^{-1} Q^\top k_*
    = \sum_i (q_i^\top k_*)^2 / \theta_i
    \;\le\; k_*^\top K^{-1} k_* .
    $$

    The inequality is the one-signed bias of the Lanczos approximation
    (see `love_cache`): the GP predictive variance

    $$
    \sigma_*^2 = k_{**} - \mathrm{love\_variance}(\mathrm{cache}, k_*)
    $$

    is an **upper bound** on the exact predictive variance, tight once the
    cache has converged for this ``k_*`` (check with `love_residual`).

    Args:
        cache: A `LOVECache` from `love_cache`.
        K_star_row: Cross-covariance vector ``k(X_{train}, x_*)``,
            shape ``(N,)``.

    Returns:
        Scalar approximation of ``k_*^T K^{-1} k_*``.
    """
    z = einsum(cache.Q, K_star_row, "n k, n -> k")
    return jnp.sum(z**2 * cache.inv_eigvals)

`love_residual(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N']) -> Float[Array, '']` ¶

Relative solve residual of a LOVE cache for one cross-covariance.

Convergence diagnostic for love_variance: with a = Q \Theta^{-1} Q^\top k_* (the cache's approximation of K^{-1} k_*), returns

\[ r(k_*) = \frac{\lVert K a - k_* \rVert_2}{\lVert k_* \rVert_2}. \]

r = 0 exactly when the cache solves K a = k_*, in which case love_variance is exact for this k_* (k_* = 0 returns 0). A value well above the working precision means lanczos_order is too small for this test point. A small value is not sufficient on its own for an ill-conditioned K: the variance error is r^\top K^{-1} r, which can be large for a small \lVert r \rVert when K has tiny eigenvalues; love_variance_error_bound turns the residual into a bound in variance units. Costs one matvec with K.

Parameters:

Name	Type	Description	Default
`cache`	`LOVECache`	A `LOVECache` from `love_cache`.	required
`K_op`	`AbstractLinearOperator`	The operator the cache was built from, shape `(N, N)`.	required
`K_star_row`	`Float[Array, ' N']`	Cross-covariance vector `k(X_{train}, x_*)`, shape `(N,)`.	required

Returns:

Type	Description
`Float[Array, '']`	Scalar relative residual `r(k_*)`.

Source code in src/gaussx/_gp/_love.py

def love_residual(
    cache: LOVECache,
    K_op: lx.AbstractLinearOperator,
    K_star_row: Float[Array, " N"],
) -> Float[Array, ""]:
    r"""Relative solve residual of a LOVE cache for one cross-covariance.

    Convergence diagnostic for `love_variance`: with
    ``a = Q \Theta^{-1} Q^\top k_*`` (the cache's approximation of
    ``K^{-1} k_*``), returns

    $$
    r(k_*) = \frac{\lVert K a - k_* \rVert_2}{\lVert k_* \rVert_2}.
    $$

    ``r = 0`` exactly when the cache solves ``K a = k_*``, in which case
    `love_variance` is exact for this ``k_*`` (``k_* = 0`` returns ``0``).
    A value well above the working precision means ``lanczos_order`` is
    too small for this test point. A *small* value is not sufficient on its
    own for an ill-conditioned ``K``: the variance error is
    ``r^\top K^{-1} r``, which can be large for a small ``\lVert r \rVert``
    when ``K`` has tiny eigenvalues; `love_variance_error_bound` turns the
    residual into a bound in variance units. Costs one matvec with ``K``.

    Args:
        cache: A `LOVECache` from `love_cache`.
        K_op: The operator the cache was built from, shape ``(N, N)``.
        K_star_row: Cross-covariance vector ``k(X_{train}, x_*)``,
            shape ``(N,)``.

    Returns:
        Scalar relative residual ``r(k_*)``.
    """
    r = _residual(cache, K_op, K_star_row)
    num = jnp.linalg.norm(r)
    den = jnp.linalg.norm(K_star_row)
    # k_* = 0 is solved exactly (a = 0, r = 0): report 0, not 0/0.
    return jnp.where(den > 0, num / jnp.where(den > 0, den, 1.0), 0.0)

`love_variance_error_bound(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N'], lambda_min: float | Float[Array, '']) -> Float[Array, '']` ¶

Upper bound on the variance over-estimate of a LOVE cache.

The cache is a Galerkin approximation, so with a = Q \Theta^{-1} Q^\top k_* and r = K a - k_* (Q^\top r = 0) the error of the quadratic form is exactly r^\top K^{-1} r:

\[ 0 \le k_*^\top K^{-1} k_* - \mathrm{love\_variance}(\mathrm{cache}, k_*) = r^\top K^{-1} r \le \frac{\lVert r \rVert_2^2}{\lambda_{\min}(K)} . \]

Unlike the relative residual of love_residual, this is in the units of the variance and accounts for the conditioning of K. For a GP with K = K_f + \sigma^2 I the noise variance \sigma^2 is a valid lambda_min.

Parameters:

Name	Type	Description	Default
`cache`	`LOVECache`	A `LOVECache` from `love_cache`.	required
`K_op`	`AbstractLinearOperator`	The operator the cache was built from, shape `(N, N)`.	required
`K_star_row`	`Float[Array, ' N']`	Cross-covariance vector, shape `(N,)`.	required
`lambda_min`	`float \| Float[Array, '']`	A lower bound on the smallest eigenvalue of `K`.	required

Returns:

Type	Description
`Float[Array, '']`	Scalar upper bound on the absolute predictive-variance error.

Source code in src/gaussx/_gp/_love.py

def love_variance_error_bound(
    cache: LOVECache,
    K_op: lx.AbstractLinearOperator,
    K_star_row: Float[Array, " N"],
    lambda_min: float | Float[Array, ""],
) -> Float[Array, ""]:
    r"""Upper bound on the variance over-estimate of a LOVE cache.

    The cache is a Galerkin approximation, so with
    ``a = Q \Theta^{-1} Q^\top k_*`` and ``r = K a - k_*`` (``Q^\top r = 0``)
    the error of the quadratic form is exactly ``r^\top K^{-1} r``:

    $$
    0 \le k_*^\top K^{-1} k_* - \mathrm{love\_variance}(\mathrm{cache}, k_*)
      = r^\top K^{-1} r \le \frac{\lVert r \rVert_2^2}{\lambda_{\min}(K)} .
    $$

    Unlike the relative residual of `love_residual`, this is in the units
    of the variance and accounts for the conditioning of ``K``. For a GP
    with ``K = K_f + \sigma^2 I`` the noise variance ``\sigma^2`` is a
    valid ``lambda_min``.

    Args:
        cache: A `LOVECache` from `love_cache`.
        K_op: The operator the cache was built from, shape ``(N, N)``.
        K_star_row: Cross-covariance vector, shape ``(N,)``.
        lambda_min: A lower bound on the smallest eigenvalue of ``K``.

    Returns:
        Scalar upper bound on the absolute predictive-variance error.
    """
    r = _residual(cache, K_op, K_star_row)
    return jnp.sum(r**2) / lambda_min

`leave_one_out_cv(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None, diag_inv_method: str = 'solve', diag_inv_num_probes: int = 30, diag_inv_key: jax.Array | None = None) -> LOOResult` ¶

LOO-CV via the bordered-system identity.

Computes leave-one-out predictive means, variances, and log-likelihood without refitting the model N times.

Math:

alpha = K_y^{-1} y
mu_LOO_i   = y_i - alpha_i / [K_y^{-1}]_{ii}
sigma^2_LOO_i = 1 / [K_y^{-1}]_{ii}
LOO-CV = -(1/2) sum_i [ log sigma^2_LOO_i
                        + (y_i - mu_LOO_i)^2 / sigma^2_LOO_i
                        + log 2 pi ]

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A linear operator representing the (noise-inclusive) kernel matrix K_y.	required
`y`	`Float[Array, ' N']`	Observation vector of shape `(N,)`.	required
`solver`	`AbstractSolveStrategy \| None`	Optional solve strategy for computing K_y^{-1} y and for the `diag_inv` computation. When `None`, falls back to structural dispatch.	`None`
`diag_inv_method`	`str`	Method passed to `diag_inv`. Defaults to `"solve"` (exact, deterministic, and honours `solver`) rather than `diag_inv`'s own `"auto"`, which switches to the stochastic Hutchinson estimator for unstructured `N > 2048`. `"cholesky"` is equally exact and deterministic and is usually faster for a dense `K_y` (one factorisation instead of `N` solves); `"auto"` is preferable for structured operators with an exact fast path (`Kronecker`, `BlockTriDiag`, ...).	`'solve'`
`diag_inv_num_probes`	`int`	Number of Hutchinson probes when `diag_inv_method="hutchinson"`.	`30`
`diag_inv_key`	`Array \| None`	PRNG key for probe generation when `diag_inv_method="hutchinson"`.	`None`

Returns:

Type	Description
`LOOResult`	A `LOOResult` containing the LOO log-likelihood,
`LOOResult`	predictive means, and predictive variances.

Source code in src/gaussx/_gp/_loo.py

def leave_one_out_cv(
    operator: lx.AbstractLinearOperator,
    y: Float[Array, " N"],
    *,
    solver: AbstractSolveStrategy | None = None,
    diag_inv_method: str = "solve",
    diag_inv_num_probes: int = 30,
    diag_inv_key: jax.Array | None = None,
) -> LOOResult:
    """LOO-CV via the bordered-system identity.

    Computes leave-one-out predictive means, variances, and
    log-likelihood without refitting the model N times.

    Math:

        alpha = K_y^{-1} y
        mu_LOO_i   = y_i - alpha_i / [K_y^{-1}]_{ii}
        sigma^2_LOO_i = 1 / [K_y^{-1}]_{ii}
        LOO-CV = -(1/2) sum_i [ log sigma^2_LOO_i
                                + (y_i - mu_LOO_i)^2 / sigma^2_LOO_i
                                + log 2 pi ]

    Args:
        operator: A linear operator representing the (noise-inclusive)
            kernel matrix K_y.
        y: Observation vector of shape ``(N,)``.
        solver: Optional solve strategy for computing K_y^{-1} y
            and for the ``diag_inv`` computation. When ``None``,
            falls back to structural dispatch.
        diag_inv_method: Method passed to `diag_inv`. Defaults to
            ``"solve"`` (exact, deterministic, and honours ``solver``)
            rather than `diag_inv`'s own ``"auto"``, which switches to the
            stochastic Hutchinson estimator for unstructured ``N > 2048``.
            ``"cholesky"`` is equally exact and deterministic and is
            usually faster for a dense ``K_y`` (one factorisation instead
            of ``N`` solves); ``"auto"`` is preferable for structured
            operators with an exact fast path (`Kronecker`,
            `BlockTriDiag`, ...).
        diag_inv_num_probes: Number of Hutchinson probes when
            ``diag_inv_method="hutchinson"``.
        diag_inv_key: PRNG key for probe generation when
            ``diag_inv_method="hutchinson"``.

    Returns:
        A `LOOResult` containing the LOO log-likelihood,
        predictive means, and predictive variances.
    """
    alpha = dispatch_solve(operator, y, solver)
    diag_Kinv = diag_inv(
        operator,
        method=diag_inv_method,
        num_probes=diag_inv_num_probes,
        key=diag_inv_key,
        solver=solver,
    )

    loo_means = y - alpha / diag_Kinv
    loo_variances = 1.0 / diag_Kinv

    # (y_i - mu_i)^2 / sigma^2_i simplifies to alpha_i^2 / diag_Kinv_i
    loo_ll = -0.5 * jnp.sum(
        jnp.log(loo_variances) + alpha**2 / diag_Kinv + jnp.log(2.0 * jnp.pi)
    )

    return LOOResult(
        loo_log_likelihood=loo_ll,
        loo_means=loo_means,
        loo_variances=loo_variances,
    )

Multi-output projections¶

The orthogonal instantaneous linear mixing model (OILMM): project multi-output observations into independent latent processes and back.

Structured linear algebra and Gaussian primitives for JAX.

`oilmm_project(Y: Float[Array, 'N P'], W: Float[Array, 'P L'], noise_var: Float[Array, ' P'] | float, *, check_orthonormal: bool = False, orthonormal_tol: float | None = None) -> tuple[Float[Array, 'N L'], Float[Array, ' L']]` ¶

Project multi-output data to independent latent GPs via OILMM.

Given an orthonormal mixing matrix W ∈ ℝᴾˣᴸ with WᵀW = I_L, projects P-output observations to L independent latent channels:

Y_latent    = Y W                      (N, L)
σ²_latent   = diag(Wᵀ D W) = (W ⊙ W)ᵀ σ²   (L,)

with D = diag(σ²). For isotropic noise Wᵀ D W = σ² I is exactly diagonal. For heteroscedastic noise Wᵀ D W has off-diagonal terms; only its diagonal is returned, which is the OILMM independence approximation (the latent channels are treated as independent).

WᵀW = I is a precondition and is not checked by default: with a non-orthonormal W the projection Y W is not the least-squares one (that would be Y W (WᵀW)⁻¹) and the latent noise is wrong.

Parameters:

Name	Type	Description	Default
`Y`	`Float[Array, 'N P']`	Observations, shape `(N, P)`.	required
`W`	`Float[Array, 'P L']`	Orthonormal mixing matrix, shape `(P, L)` with WᵀW = I_L.	required
`noise_var`	`Float[Array, ' P'] \| float`	Observation noise variance. Scalar for isotropic noise, or shape `(P,)` for heteroscedastic noise.	required
`check_orthonormal`	`bool`	If `True`, raise (jit-safe, via `equinox.error_if`) when `max \|WᵀW − I\| > orthonormal_tol`.	`False`
`orthonormal_tol`	`float \| None`	Tolerance for `check_orthonormal`. Default `sqrt(eps)` of `W`'s dtype.	`None`

Returns:

Type	Description
`Float[Array, 'N L']`	Tuple `(Y_latent, noise_latent)` with shapes `(N, L)`
`Float[Array, ' L']`	and `(L,)`.

Source code in src/gaussx/_gp/_oilmm.py

def oilmm_project(
    Y: Float[Array, "N P"],
    W: Float[Array, "P L"],
    noise_var: Float[Array, " P"] | float,
    *,
    check_orthonormal: bool = False,
    orthonormal_tol: float | None = None,
) -> tuple[Float[Array, "N L"], Float[Array, " L"]]:
    """Project multi-output data to independent latent GPs via OILMM.

    Given an orthonormal mixing matrix W ∈ ℝᴾˣᴸ with WᵀW = I_L, projects
    P-output observations to L independent latent channels:

        Y_latent    = Y W                      (N, L)
        σ²_latent   = diag(Wᵀ D W) = (W ⊙ W)ᵀ σ²   (L,)

    with ``D = diag(σ²)``. For isotropic noise ``Wᵀ D W = σ² I`` is exactly
    diagonal. For heteroscedastic noise ``Wᵀ D W`` has off-diagonal terms;
    only its diagonal is returned, which is the OILMM independence
    approximation (the latent channels are treated as independent).

    ``WᵀW = I`` is a precondition and is **not** checked by default: with a
    non-orthonormal ``W`` the projection ``Y W`` is not the least-squares
    one (that would be ``Y W (WᵀW)⁻¹``) and the latent noise is wrong.

    Args:
        Y: Observations, shape ``(N, P)``.
        W: Orthonormal mixing matrix, shape ``(P, L)`` with WᵀW = I_L.
        noise_var: Observation noise variance. Scalar for isotropic noise,
            or shape ``(P,)`` for heteroscedastic noise.
        check_orthonormal: If ``True``, raise (jit-safe, via
            `equinox.error_if`) when ``max |WᵀW − I| > orthonormal_tol``.
        orthonormal_tol: Tolerance for ``check_orthonormal``. Default
            ``sqrt(eps)`` of ``W``'s dtype.

    Returns:
        Tuple ``(Y_latent, noise_latent)`` with shapes ``(N, L)``
        and ``(L,)``.
    """
    if check_orthonormal:
        gram = einsum(W, W, "p i, p j -> i j")
        tol = (
            math.sqrt(float(jnp.finfo(W.dtype).eps))
            if orthonormal_tol is None
            else orthonormal_tol
        )
        err = jnp.max(jnp.abs(gram - jnp.eye(W.shape[1], dtype=W.dtype)))
        W = eqx.error_if(
            W, err > tol, "oilmm_project: W is not orthonormal (WᵀW != I)."
        )
    Y_latent = Y @ W  # (N, L)
    noise_var = jnp.broadcast_to(jnp.asarray(noise_var), (W.shape[0],))  # (P,)
    noise_latent = einsum(W**2, noise_var, "p l, p -> l")  # (L,)
    return Y_latent, noise_latent

`oilmm_back_project(f_means: Float[Array, 'N L'], f_vars: Float[Array, 'N L'], W: Float[Array, 'P L']) -> tuple[Float[Array, 'N P'], Float[Array, 'N P']]` ¶

Back-project latent GP predictions to the observation space.

Reconstructs observation-space predictions via:

y_means = f_means Wᵀ              (N, P)
y_vars  = f_vars (W ⊙ W)ᵀ        (N, P)

Parameters:

Name	Type	Description	Default
`f_means`	`Float[Array, 'N L']`	Latent predictive means, shape `(N, L)`.	required
`f_vars`	`Float[Array, 'N L']`	Latent predictive variances, shape `(N, L)`.	required
`W`	`Float[Array, 'P L']`	Orthogonal mixing matrix, shape `(P, L)` with WᵀW = I_L.	required

Returns:

Type	Description
`tuple[Float[Array, 'N P'], Float[Array, 'N P']]`	Tuple `(y_means, y_vars)` with shapes `(N, P)` and `(N, P)`.

Source code in src/gaussx/_gp/_oilmm.py

def oilmm_back_project(
    f_means: Float[Array, "N L"],
    f_vars: Float[Array, "N L"],
    W: Float[Array, "P L"],
) -> tuple[Float[Array, "N P"], Float[Array, "N P"]]:
    """Back-project latent GP predictions to the observation space.

    Reconstructs observation-space predictions via:

        y_means = f_means Wᵀ              (N, P)
        y_vars  = f_vars (W ⊙ W)ᵀ        (N, P)

    Args:
        f_means: Latent predictive means, shape ``(N, L)``.
        f_vars: Latent predictive variances, shape ``(N, L)``.
        W: Orthogonal mixing matrix, shape ``(P, L)`` with WᵀW = I_L.

    Returns:
        Tuple ``(y_means, y_vars)`` with shapes ``(N, P)`` and ``(N, P)``.
    """
    y_means = f_means @ W.T  # (N, P)
    y_vars = f_vars @ (W**2).T  # (N, P)
    return y_means, y_vars

Deprecated aliases¶

These names still work but warn on use and will be removed in gaussx 0.7.0 (see the deprecation policy).

Structured linear algebra and Gaussian primitives for JAX.

`base_conditional(K_mm: Float[Array, 'M M'], K_mn: Float[Array, 'M N'], K_nn: Float[Array, 'N N'] | Float[Array, ' N'], f: Float[Array, 'M R'] | Float[Array, ' M'], *, q_sqrt: Float[Array, 'R M M'] | Float[Array, 'M R'] | Float[Array, 'M M'] | Float[Array, ' M'] | None = None, white: bool = False, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, 'N R'], Float[Array, ...]]` ¶

Deprecated: use sparse_conditional (note K_xz = K_mn.T).

base_conditional(K_mm, K_mn, K_nn, f, q_sqrt=..., white=...) equals sparse_conditional(K_mm, K_mn.T, K_nn, f, q_sqrt=..., white=...). It takes the cross-covariance as (M, N) (inducing × data), the transpose of every other sparse-GP helper; for M == N a transposed argument is not detectable, so the convention is changed under a new name rather than in place. solver was never used. Removed in gaussx 0.7.0.

Parameters:

Name	Type	Description	Default
`K_mm`	`Float[Array, 'M M']`	Prior covariance at inducing points, shape `(M, M)`.	required
`K_mn`	`Float[Array, 'M N']`	Cross-covariance, shape `(M, N)`.	required
`K_nn`	`Float[Array, 'N N'] \| Float[Array, ' N']`	Test-point covariance, full `(N, N)` or diagonal `(N,)`.	required
`f`	`Float[Array, 'M R'] \| Float[Array, ' M']`	Inducing function values, shape `(M, R)` or `(M,)`.	required
`q_sqrt`	`Float[Array, 'R M M'] \| Float[Array, 'M R'] \| Float[Array, 'M M'] \| Float[Array, ' M'] \| None`	Optional variational root, as in `sparse_conditional`.	`None`
`white`	`bool`	Whitened parameterisation, as in `sparse_conditional`.	`False`
`solver`	`AbstractSolverStrategy \| None`	Ignored.	`None`

Returns:

Type	Description
`tuple[Float[Array, 'N R'], Float[Array, ...]]`	`(mean, var)` as in `sparse_conditional`.

Source code in src/gaussx/_gp/_base_conditional.py

def base_conditional(
    K_mm: Float[Array, "M M"],
    K_mn: Float[Array, "M N"],
    K_nn: Float[Array, "N N"] | Float[Array, " N"],
    f: Float[Array, "M R"] | Float[Array, " M"],
    *,
    q_sqrt: Float[Array, "R M M"]
    | Float[Array, "M R"]
    | Float[Array, "M M"]
    | Float[Array, " M"]
    | None = None,
    white: bool = False,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "N R"], Float[Array, ...]]:
    r"""Deprecated: use `sparse_conditional` (note ``K_xz = K_mn.T``).

    ``base_conditional(K_mm, K_mn, K_nn, f, q_sqrt=..., white=...)`` equals
    ``sparse_conditional(K_mm, K_mn.T, K_nn, f, q_sqrt=..., white=...)``.
    It takes the cross-covariance as ``(M, N)`` (inducing × data), the
    transpose of every other sparse-GP helper; for ``M == N`` a transposed
    argument is not detectable, so the convention is changed under a new
    name rather than in place. ``solver`` was never used. Removed in
    gaussx 0.7.0.

    Args:
        K_mm: Prior covariance at inducing points, shape ``(M, M)``.
        K_mn: Cross-covariance, shape ``(M, N)``.
        K_nn: Test-point covariance, full ``(N, N)`` or diagonal ``(N,)``.
        f: Inducing function values, shape ``(M, R)`` or ``(M,)``.
        q_sqrt: Optional variational root, as in `sparse_conditional`.
        white: Whitened parameterisation, as in `sparse_conditional`.
        solver: Ignored.

    Returns:
        ``(mean, var)`` as in `sparse_conditional`.
    """
    del solver
    warn_deprecated(
        "base_conditional(K_mm, K_mn, ...) is deprecated and will be removed in "
        "gaussx 0.7.0; use sparse_conditional(K_zz, K_xz, K_xx, q_mu, "
        "...) with K_xz = K_mn.T (shape (N, M))."
    )
    return sparse_conditional(
        K_mm, rearrange(K_mn, "m n -> n m"), K_nn, f, q_sqrt=q_sqrt, white=white
    )

Gaussian Processes¶

Conditioning & prediction¶

PredictionCache ¶

build_prediction_cache(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None) -> PredictionCache ¶

predict_mean(cache: PredictionCache, K_cross: Float[Array, 'Nt N']) -> Float[Array, ' Nt'] ¶

kronecker_mll(K_factors: list[lx.AbstractLinearOperator], y: Float[Array, ' N'], noise_var: float, grid_shape: tuple[int, ...]) -> Float[Array, ''] ¶

Pathwise sampling¶

Whitening¶

whiten_covariance = unwhiten_covariance module-attribute ¶

unwhiten(m_tilde: Float[Array, ' M'], L: lx.AbstractLinearOperator) -> Float[Array, ' M'] ¶

unwhiten_covariance(L: lx.AbstractLinearOperator, S_tilde: lx.AbstractLinearOperator) -> lx.MatrixLinearOperator ¶

whitened_svgp_predict(K_zz_op: lx.AbstractLinearOperator, K_xz: Float[Array, 'N M'], u_mean: Float[Array, ' M'], u_chol: Float[Array, 'M M'], K_xx_diag: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']] ¶

svgp_variance_adjustment(K_zz_op: lx.AbstractLinearOperator, S_u: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator ¶

Variational bounds & KL¶

variational_elbo_gaussian(y: Float[Array, ' N'], f_loc: Float[Array, ' N'], f_var: Float[Array, ' N'], noise_var: float, kl: Float[Array, '']) -> Float[Array, ''] ¶

variational_elbo_mc(log_likelihood_fn: Callable[[Float[Array, ' N']], Float[Array, '']], f_samples: Float[Array, 'S N'], kl: Float[Array, '']) -> Float[Array, ''] ¶

collapsed_elbo(y: Float[Array, ' N'], K_xx_diag: Float[Array, ' N'], K_xz: Float[Array, 'N M'], K_zz: Float[Array, 'M M'] | lx.AbstractLinearOperator, noise_var: float, *, jitter: float = 1e-06, solver: AbstractSolverStrategy | None = None) -> Float[Array, ''] ¶

gauss_kl(q_mu: Float[Array, 'M R'], q_sqrt: Float[Array, 'R M M'] | Float[Array, 'M R'], K: Float[Array, 'M M'] | lx.AbstractLinearOperator | None = None, *, solver: AbstractSolverStrategy | None = None) -> Float[Array, ''] ¶

Cross-validation & diagnostics¶

LOVECache ¶

LOOResult ¶

love_cache(K_op: lx.AbstractLinearOperator, lanczos_order: int = 50, key: jax.Array | None = None, *, initial_vector: Float[Array, ' N'] | None = None) -> LOVECache ¶

love_variance(cache: LOVECache, K_star_row: Float[Array, ' N']) -> Float[Array, ''] ¶

love_residual(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N']) -> Float[Array, ''] ¶

love_variance_error_bound(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N'], lambda_min: float | Float[Array, '']) -> Float[Array, ''] ¶

leave_one_out_cv(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None, diag_inv_method: str = 'solve', diag_inv_num_probes: int = 30, diag_inv_key: jax.Array | None = None) -> LOOResult ¶

Multi-output projections¶

oilmm_project(Y: Float[Array, 'N P'], W: Float[Array, 'P L'], noise_var: Float[Array, ' P'] | float, *, check_orthonormal: bool = False, orthonormal_tol: float | None = None) -> tuple[Float[Array, 'N L'], Float[Array, ' L']] ¶

oilmm_back_project(f_means: Float[Array, 'N L'], f_vars: Float[Array, 'N L'], W: Float[Array, 'P L']) -> tuple[Float[Array, 'N P'], Float[Array, 'N P']] ¶

Deprecated aliases¶

`PredictionCache` ¶

`build_prediction_cache(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None) -> PredictionCache` ¶

`predict_mean(cache: PredictionCache, K_cross: Float[Array, 'Nt N']) -> Float[Array, ' Nt']` ¶

`kronecker_mll(K_factors: list[lx.AbstractLinearOperator], y: Float[Array, ' N'], noise_var: float, grid_shape: tuple[int, ...]) -> Float[Array, '']` ¶

`whiten_covariance = unwhiten_covariance` `module-attribute` ¶

`unwhiten(m_tilde: Float[Array, ' M'], L: lx.AbstractLinearOperator) -> Float[Array, ' M']` ¶

`unwhiten_covariance(L: lx.AbstractLinearOperator, S_tilde: lx.AbstractLinearOperator) -> lx.MatrixLinearOperator` ¶

`whitened_svgp_predict(K_zz_op: lx.AbstractLinearOperator, K_xz: Float[Array, 'N M'], u_mean: Float[Array, ' M'], u_chol: Float[Array, 'M M'], K_xx_diag: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']]` ¶

`svgp_variance_adjustment(K_zz_op: lx.AbstractLinearOperator, S_u: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator` ¶

`variational_elbo_gaussian(y: Float[Array, ' N'], f_loc: Float[Array, ' N'], f_var: Float[Array, ' N'], noise_var: float, kl: Float[Array, '']) -> Float[Array, '']` ¶

`variational_elbo_mc(log_likelihood_fn: Callable[[Float[Array, ' N']], Float[Array, '']], f_samples: Float[Array, 'S N'], kl: Float[Array, '']) -> Float[Array, '']` ¶

`collapsed_elbo(y: Float[Array, ' N'], K_xx_diag: Float[Array, ' N'], K_xz: Float[Array, 'N M'], K_zz: Float[Array, 'M M'] | lx.AbstractLinearOperator, noise_var: float, *, jitter: float = 1e-06, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']` ¶

`gauss_kl(q_mu: Float[Array, 'M R'], q_sqrt: Float[Array, 'R M M'] | Float[Array, 'M R'], K: Float[Array, 'M M'] | lx.AbstractLinearOperator | None = None, *, solver: AbstractSolverStrategy | None = None) -> Float[Array, '']` ¶

`LOVECache` ¶

`LOOResult` ¶

`love_cache(K_op: lx.AbstractLinearOperator, lanczos_order: int = 50, key: jax.Array | None = None, *, initial_vector: Float[Array, ' N'] | None = None) -> LOVECache` ¶

`love_variance(cache: LOVECache, K_star_row: Float[Array, ' N']) -> Float[Array, '']` ¶

`love_residual(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N']) -> Float[Array, '']` ¶

`love_variance_error_bound(cache: LOVECache, K_op: lx.AbstractLinearOperator, K_star_row: Float[Array, ' N'], lambda_min: float | Float[Array, '']) -> Float[Array, '']` ¶

`leave_one_out_cv(operator: lx.AbstractLinearOperator, y: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None, diag_inv_method: str = 'solve', diag_inv_num_probes: int = 30, diag_inv_key: jax.Array | None = None) -> LOOResult` ¶

`oilmm_project(Y: Float[Array, 'N P'], W: Float[Array, 'P L'], noise_var: Float[Array, ' P'] | float, *, check_orthonormal: bool = False, orthonormal_tol: float | None = None) -> tuple[Float[Array, 'N L'], Float[Array, ' L']]` ¶

`oilmm_back_project(f_means: Float[Array, 'N L'], f_vars: Float[Array, 'N L'], W: Float[Array, 'P L']) -> tuple[Float[Array, 'N P'], Float[Array, 'N P']]` ¶