Primitives

Layer 0: pure functions over lineax.AbstractLinearOperator with structural dispatch — each primitive inspects the operator (diagonal, Kronecker, block-diagonal, low-rank, block-tridiagonal, …) and routes to the cheapest exact algorithm, falling back to a dense computation only when no structured path exists. Where a structured operator has to be materialized anyway — cholesky of a SumOfKroneckers — a DenseFallbackWarning names the matrix-free alternative.

Solve, logdet & Cholesky

The workhorses behind Gaussian densities: \(A^{-1}b\), \(\log|A|\), and \(A = LL^\top\). cholesky returns a lazy lower-triangular operator that preserves structure (the Cholesky of a Kronecker is a Kronecker of Cholesky factors); cholesky_logdet turns an existing factor into \(\log|A| = 2\sum_i \log L_{ii}\) for free.

solve(A, b, solver=...) takes either a lineax solver, such as lineax.CG(...), which replaces the dense fallback and is threaded into the structural rules (per Kronecker factor, per block), or a gaussx strategy, such as CGSolver(), which then owns the whole solve. See Solvers.

Structured linear algebra and Gaussian primitives for JAX.

solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n'], *, solver: lx.AbstractLinearSolver | AbstractSolveStrategy | None = None) -> Float[Array, ' n']

Solve A x = b with structural dispatch.

Operators with no structural rule take a dense fallback. With the default solver=None that fallback has numpy semantics: a singular or numerically singular A gives a non-finite solution (nan or inf), in the value and in its gradients, rather than an error, as gaussx.logdet already has. A log-density built on it is then non-finite too, which a sampler or optimiser can reject; an exception would abort the whole run. An explicit solver keeps lineax's behaviour instead, which raises on a singular system or an iterative solver that does not converge.

Structured rules keep their own semantics for singular operators: a singular gaussx.DiagonalizedOperator, for example, returns the minimum-norm pseudo-inverse solution.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	The linear operator A.	required
vector	Float[Array, ' n']	The right-hand side b.	required
solver	AbstractLinearSolver | AbstractSolveStrategy | None	A lineax solver, which overrides the dense fallback and is threaded into the structural rules (e.g. used per Kronecker factor); or a gaussx strategy such as gaussx.CGSolver, which then owns the whole solve (gh-376). None uses a dense direct solve chosen from the operator's tags (Cholesky for positive semi-definite, triangular, else LU).	None

Returns:

Type	Description
Float[Array, ' n']	The solution x. On the default dense fallback, non-finite when
Float[Array, ' n']	A is singular.

Raises:

Type	Description
TypeError	If solver is neither a lineax solver nor a gaussx strategy.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
>>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
>>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
>>> x = gaussx.solve(K, jnp.ones(4))  # one solve per factor
>>> [round(float(v), 4) for v in x]
[0.25, 0.2, 0.125, 0.1]

Source code in src/gaussx/_primitives/_solve.py

def solve(
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
    *,
    solver: lx.AbstractLinearSolver | AbstractSolveStrategy | None = None,
) -> Float[Array, " n"]:
    """Solve ``A x = b`` with structural dispatch.

    Operators with no structural rule take a dense fallback. With the
    default ``solver=None`` that fallback has numpy semantics: a singular or
    numerically singular ``A`` gives a non-finite solution (``nan`` or
    ``inf``), in the value and in its gradients, rather than an error, as
    `gaussx.logdet` already has. A log-density built on it is then
    non-finite too, which a sampler or optimiser can reject; an exception
    would abort the whole run. An explicit ``solver`` keeps lineax's
    behaviour instead, which raises on a singular system or an iterative
    solver that does not converge.

    Structured rules keep their own semantics for singular operators: a
    singular `gaussx.DiagonalizedOperator`, for example, returns the
    minimum-norm pseudo-inverse solution.

    Args:
        operator: The linear operator A.
        vector: The right-hand side b.
        solver: A lineax solver, which overrides the dense fallback and is
            threaded into the structural rules (e.g. used per Kronecker
            factor); or a gaussx strategy such as `gaussx.CGSolver`, which
            then owns the whole solve (gh-376). ``None`` uses a dense direct
            solve chosen from the operator's tags (Cholesky for positive
            semi-definite, triangular, else LU).

    Returns:
        The solution x. On the default dense fallback, non-finite when
        ``A`` is singular.

    Raises:
        TypeError: If *solver* is neither a lineax solver nor a gaussx
            strategy.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
        >>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
        >>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
        >>> x = gaussx.solve(K, jnp.ones(4))  # one solve per factor
        >>> [round(float(v), 4) for v in x]
        [0.25, 0.2, 0.125, 0.1]
    """
    if solver is not None and not isinstance(solver, lx.AbstractLinearSolver):
        # lazy import, cycle: _strategies._base -> _strategies -> _strategies._dense ->
        #   _primitives._solve
        from gaussx._strategies._base import AbstractSolveStrategy

        if isinstance(solver, AbstractSolveStrategy):
            return solver.solve(operator, vector)
        raise TypeError(
            "solver must be a lineax.AbstractLinearSolver (e.g. lineax.CG(...)) "
            "or a gaussx.AbstractSolveStrategy (e.g. gaussx.CGSolver()); got "
            f"{type(solver).__name__}."
        )
    if isinstance(operator, lx.IdentityLinearOperator):
        return vector
    if isinstance(operator, InverseOperator):
        # (A⁻¹)⁻¹ b = A b: one matvec, no factorisation (gh-349).
        return operator.original.mv(vector)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _solve_diagonal(operator, vector)
    if isinstance(operator, DiagonalizedOperator):
        return _solve_diagonalised(operator, vector)
    if isinstance(operator, MaskedOperator) and operator.capacitance is not None:
        return _solve_masked(operator, vector)
    if isinstance(operator, BlockDiag):
        return _solve_block_diag(operator, vector, solver)
    if isinstance(operator, Kronecker):
        return _solve_kronecker(operator, vector, solver)
    if isinstance(operator, LowRankUpdate):
        if operator.rank == 0:
            return solve(operator.base, vector, solver=solver)
        c = orthonormal_scaled_identity(operator)
        if c is not None:
            return _solve_low_rank_orthonormal(operator, c, vector)
        return _solve_low_rank(operator, vector, solver)
    if isinstance(operator, SumOfKroneckers):
        return _solve_sum_of_kroneckers(operator, vector, solver)
    if isinstance(operator, KroneckerSum):
        return _solve_kronecker_sum(operator, vector, solver)
    if isinstance(operator, KroneckerSumSqrt):
        return operator.solve(vector)
    if isinstance(operator, SpectralFunction):
        return operator.solve(vector)
    if isinstance(operator, BlockTriDiag) and operator.symmetric:
        # Non-symmetric diagonal blocks take the dense LU fallback (gh-344).
        return _solve_block_tridiag(operator, vector)
    if isinstance(operator, SparseOperator):
        return _solve_sparse(operator, vector, solver)
    if isinstance(operator, LowerBlockTriDiag):
        return _solve_lower_block_tridiag(operator, vector)
    if isinstance(operator, UpperBlockTriDiag):
        return _solve_upper_block_tridiag(operator, vector)
    if isinstance(operator, lx.TaggedLinearOperator):
        return _solve_tagged(operator, vector, solver)
    if isinstance(operator, lx.MulLinearOperator):
        return solve(operator.operator, vector, solver=solver) / operator.scalar
    if isinstance(operator, lx.DivLinearOperator):
        return solve(operator.operator, vector, solver=solver) * operator.scalar
    if isinstance(operator, lx.NegLinearOperator):
        return -solve(operator.operator, vector, solver=solver)
    if isinstance(operator, lx.ComposedLinearOperator):
        # (A B) x = b  =>  x = B^{-1} (A^{-1} b)
        y = solve(operator.operator1, vector, solver=solver)
        return solve(operator.operator2, y, solver=solver)
    if isinstance(operator, lx.AddLinearOperator):
        # ``sum_operator`` builds native lineax sums rather than a
        # `SumOfKroneckers`, so the same reduction has to be reachable here.
        x = _sum_of_kroneckers_solve(operator, vector)
        if x is not None:
            return x
    return _solve_fallback(operator, vector, solver)

logdet(operator: lx.AbstractLinearOperator) -> Float[Array, '']

Compute log |det(A)| with structural dispatch.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	The linear operator A.	required

Returns:

Type	Description
Float[Array, '']	Scalar log |det(A)|.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
>>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
>>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
>>> round(float(gaussx.logdet(K)), 4)  # n_B logdet(A) + n_A logdet(B)
7.3778

Source code in src/gaussx/_primitives/_logdet.py

def logdet(operator: lx.AbstractLinearOperator) -> Float[Array, ""]:
    """Compute log |det(A)| with structural dispatch.

    Args:
        operator: The linear operator A.

    Returns:
        Scalar log |det(A)|.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
        >>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
        >>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
        >>> round(float(gaussx.logdet(K)), 4)  # n_B logdet(A) + n_A logdet(B)
        7.3778
    """
    if isinstance(operator, lx.IdentityLinearOperator):
        return jnp.zeros((), dtype=operator.in_structure().dtype)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _logdet_diagonal(operator)
    if isinstance(operator, DiagonalizedOperator):
        return _logdet_diagonalised(operator)
    if isinstance(operator, BlockDiag):
        return _logdet_block_diag(operator)
    if isinstance(operator, Kronecker):
        return _logdet_kronecker(operator)
    if isinstance(operator, LowRankUpdate):
        if operator.rank == 0:
            return logdet(operator.base)
        c = orthonormal_scaled_identity(operator)
        if c is not None:
            return _logdet_low_rank_orthonormal(operator, c)
        return _logdet_low_rank(operator)
    if isinstance(operator, SumOfKroneckers):
        return _logdet_sum_of_kroneckers(operator)
    if isinstance(operator, KroneckerSum):
        return _logdet_kronecker_sum(operator)
    if isinstance(operator, SpectralFunction):
        return operator.logdet()
    if isinstance(operator, BlockTriDiag) and operator.symmetric:
        # Non-symmetric diagonal blocks take the dense slogdet (gh-344).
        return _logdet_block_tridiag(operator)
    if isinstance(operator, LowerBlockTriDiag | UpperBlockTriDiag):
        return _logdet_block_bidiagonal(operator)
    if isinstance(operator, SparseOperator):
        return _logdet_sparse(operator)
    if isinstance(operator, lx.TaggedLinearOperator):
        return logdet(operator.operator)
    if isinstance(operator, lx.MulLinearOperator):
        n = operator.out_size()
        return n * jnp.log(jnp.abs(operator.scalar)) + logdet(operator.operator)
    if isinstance(operator, lx.DivLinearOperator):
        n = operator.out_size()
        return logdet(operator.operator) - n * jnp.log(jnp.abs(operator.scalar))
    if isinstance(operator, lx.NegLinearOperator):
        # log|det(-A)| = log|det(A)| since |(-1)^n| = 1.
        return logdet(operator.operator)
    if isinstance(operator, lx.ComposedLinearOperator) and (
        operator.operator1.in_size() == operator.operator1.out_size()
        and operator.operator2.in_size() == operator.operator2.out_size()
    ):
        return logdet(operator.operator1) + logdet(operator.operator2)
    if isinstance(operator, lx.AddLinearOperator):
        factorization = _sum_of_kroneckers_eigen(operator)
        if factorization is not None:
            return factorization.logdet()
    return _logdet_dense(operator)

cholesky(operator: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator | SparseCholeskyFactor

cholesky(operator: SparseOperator) -> SparseCholeskyFactor

cholesky(operator: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator

Compute Cholesky factor L such that A = L L^T.

Returns a linear operator (not a raw array). For structured operators, the result preserves structure. A SparseOperator returns a SparseCholeskyFactor of the fill-reducing permutation, P A Pᵀ = L Lᵀ (gaussx.sparse_cholesky on the cached symbolic analysis of its pattern), with solve, logdet, solve_lower_transpose and diag_inv.

A scalar multiple c · A or A / c (c > 0) factors A and scales the factor by √c, keeping its structured type. A concretely negative multiple such as -A is not PSD and raises.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A positive-definite linear operator.	required

Returns:

Type	Description
AbstractLinearOperator | SparseCholeskyFactor	Lower-triangular operator L, or a SparseCholeskyFactor.

Raises:

Type	Description
ValueError	If operator is a concretely negative multiple of a non-NSD operator, such as -A.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> psd = lx.positive_semidefinite_tag
>>> A = lx.MatrixLinearOperator(jnp.array([[2.0, 1.0], [1.0, 2.0]]), psd)
>>> B = lx.MatrixLinearOperator(jnp.array([[3.0, 0.5], [0.5, 1.0]]), psd)
>>> K = gaussx.Kronecker(A, B)
>>> L = gaussx.cholesky(K)  # Kronecker(chol(A), chol(B))
>>> type(L).__name__
'Kronecker'
>>> bool(jnp.allclose((L @ L.T).as_matrix(), K.as_matrix(), atol=1e-5))
True

Source code in src/gaussx/_primitives/_cholesky.py

def cholesky(
    operator: lx.AbstractLinearOperator,
) -> lx.AbstractLinearOperator | SparseCholeskyFactor:
    """Compute Cholesky factor L such that A = L L^T.

    Returns a linear operator (not a raw array). For structured
    operators, the result preserves structure. A `SparseOperator` returns a
    `SparseCholeskyFactor` of the fill-reducing permutation, ``P A Pᵀ = L Lᵀ``
    (`gaussx.sparse_cholesky` on the cached symbolic analysis of its
    pattern), with ``solve``, ``logdet``, ``solve_lower_transpose`` and
    ``diag_inv``.

    A scalar multiple ``c · A`` or ``A / c`` (``c > 0``) factors ``A`` and
    scales the factor by ``√c``, keeping its structured type. A concretely
    negative multiple such as ``-A`` is not PSD and raises.

    Args:
        operator: A positive-definite linear operator.

    Returns:
        Lower-triangular operator L, or a `SparseCholeskyFactor`.

    Raises:
        ValueError: If ``operator`` is a concretely negative multiple of a
            non-NSD operator, such as ``-A``.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> psd = lx.positive_semidefinite_tag
        >>> A = lx.MatrixLinearOperator(jnp.array([[2.0, 1.0], [1.0, 2.0]]), psd)
        >>> B = lx.MatrixLinearOperator(jnp.array([[3.0, 0.5], [0.5, 1.0]]), psd)
        >>> K = gaussx.Kronecker(A, B)
        >>> L = gaussx.cholesky(K)  # Kronecker(chol(A), chol(B))
        >>> type(L).__name__
        'Kronecker'
        >>> bool(jnp.allclose((L @ L.T).as_matrix(), K.as_matrix(), atol=1e-5))
        True
    """
    if isinstance(operator, lx.IdentityLinearOperator):
        return operator
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _cholesky_diagonal(operator)
    if isinstance(operator, BlockDiag):
        return _cholesky_block_diag(operator)
    if isinstance(operator, Kronecker):
        return _cholesky_kronecker(operator)
    if isinstance(operator, BlockTriDiag):
        return _cholesky_block_tridiag(operator)
    if isinstance(operator, SumOfKroneckers):
        return _cholesky_sum_kronecker(operator)
    if isinstance(operator, SparseOperator):
        return _cholesky_sparse(operator)
    if isinstance(operator, DiagonalizedOperator):
        warn_dense_fallback(
            "cholesky(DiagonalizedOperator) materialises the operator and runs "
            "a dense O(n^3) Cholesky. gaussx.sqrt(A) returns a symmetric root S "
            "with S S^T = A in O(n log n), which is enough for sampling."
        )
        return _cholesky_dense(operator)
    if isinstance(operator, lx.TaggedLinearOperator):
        return cholesky(operator.operator)
    if isinstance(
        operator, lx.MulLinearOperator | lx.DivLinearOperator | lx.NegLinearOperator
    ):
        return scaled_root(operator, cholesky, _cholesky_dense, "cholesky")
    return _cholesky_dense(operator)

cholesky_logdet(L: Float[Array, 'N N']) -> Float[Array, '']

Compute log|A| from Cholesky factor L where A = L Lᵀ.

Parameters:

Name	Type	Description	Default
L	Float[Array, 'N N']	Lower-triangular Cholesky factor, shape (N, N).	required

Returns:

Type	Description
Float[Array, '']	Scalar log-determinant.

Source code in src/gaussx/_primitives/_logdet.py

def cholesky_logdet(L: Float[Array, "N N"]) -> Float[Array, ""]:
    """Compute log|A| from Cholesky factor L where A = L Lᵀ.

    Args:
        L: Lower-triangular Cholesky factor, shape ``(N, N)``.

    Returns:
        Scalar log-determinant.
    """
    return 2.0 * jnp.sum(jnp.log(jnp.diag(L)))

Pseudo-determinant

pseudo_logdet is \(\log|A|_+ = \sum_{\lambda_i>0}\log\lambda_i\) for a symmetric PSD \(A\): half of it is the normalising constant of an intrinsic GMRF (Besag / ICAR, RW1, RW2), whose structure matrix is singular. Each structure takes its cheapest exact path:

Operator	Path
structure="laplacian" on a SparseOperator or a 1 × 1-block BlockTriDiag (Besag, RW1)	Matrix-tree theorem: \(\log\operatorname{pdet}(L) = \sum_c(\log n_c + \log\lvert L_c^{(-k_c)}\rvert)\) over connected components, one sparse (or banded) Cholesky in all, sparsity kept
null_space=B (any basis of \(\ker A\))	\(\log\lvert A + BB^\top\rvert - \log\lvert B^\top B\rvert\); the sum is matvec-only, so strategy=SLQLogdet() estimates it matrix-free
KroneckerSum (a grid Laplacian)	Factor eigenvalues, all pairwise sums \(\lambda^H_i + \lambda^W_j\); no factorisation
anything else	dense eigvalsh

The eigenvalue paths drop \(\lambda \le\) rcond \(\cdot\,\lambda_{\max}\). The matrix-tree theorem says every principal minor of a connected weighted Laplacian equals the weighted spanning-tree count \(\tau_w(G)\) and \(\operatorname{pdet}(L) = N\,\tau_w(G)\). \(\log|R|_+\) does not depend on the precision scale \(\tau\) (\(\log|\tau R|_+ = \operatorname{rank}(R)\log\tau + \log|R|_+\)), so compute it once and cache it.

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

# ICAR normalising constant for model comparison: computed once, cached
R = gaussx.besag_structure(laplacian)  # a SparseOperator graph Laplacian
half_log_pdet = 0.5 * gaussx.pseudo_logdet(R, structure="laplacian")

# On a grid: no factorisation, sum the logs of the non-zero λ^H_i + λ^W_j
L_H = lx.MatrixLinearOperator(gaussx.rw1_structure(512).as_matrix(), lx.symmetric_tag)
L_W = lx.MatrixLinearOperator(gaussx.rw1_structure(512).as_matrix(), lx.symmetric_tag)
half_log_pdet_grid = 0.5 * gaussx.pseudo_logdet(gaussx.KroneckerSum(L_H, L_W))

Structured linear algebra and Gaussian primitives for JAX.

pseudo_logdet(operator: lx.AbstractLinearOperator, *, null_space: Float[ArrayLike, 'N c'] | Float[ArrayLike, ' N'] | None = None, structure: Literal['laplacian'] | None = None, rcond: float | None = None, strategy: AbstractLogdetStrategy | None = None) -> Float[Array, '']

Log pseudo-determinant: the log of the product of the non-zero eigenvalues.

For a symmetric positive-semidefinite A, \(\log|A|_+ = \sum_{\lambda_i > 0}\log\lambda_i\). Intrinsic GMRFs (Besag / ICAR, RW1, RW2) have singular structure matrices R, and \(\tfrac12\log|R|_+\) is their normalising constant. Paths, cheapest first:

• structure="laplacian" (a weighted graph Laplacian, e.g. Besag or RW1): the matrix-tree theorem. Every principal \((n_c-1)\)-minor of a connected component's Laplacian \(L_c\) equals its weighted spanning-tree count, and \(\operatorname{pdet}(L_c) = n_c\) times it, so \(\log|L|_+ = \sum_c (\log n_c + \log|L_c^{(-k_c)}|)\). The minors are taken by replacing one node per component with a unit diagonal, which keeps the sparsity pattern, so it is one sparse Cholesky (SparseCholeskySolver by default; G4) for a SparseOperator, or one banded Cholesky for a BlockTriDiag with 1 × 1 blocks (a path, as from rw1_structure). Components are found from the sparsity pattern on the host (once per pattern), so every stored off-diagonal weight must be non-zero. Gradients are exact. A KroneckerSum (a grid Laplacian) takes the eigenvalue path below instead.
• null_space given: \(\operatorname{pdet}(A) = \det(A + NN^\top)\) for orthonormal \(N\) spanning \(\ker A\) (each zero eigenvalue becomes 1). Any basis \(B\) of the kernel works, since \(\det(A + BB^\top) = \operatorname{pdet}(A)\det(B^\top B)\) and the second factor is subtracted. A + BBᵀ is a matvec-only lineax sum, so strategy=gaussx.SLQLogdet() estimates it matrix-free; the default is gaussx.logdet of it (dense). The matrix determinant lemma does not apply because A is singular.
• A KroneckerSum (recursively, with dense, diagonal or DiagonalizedOperator factors) or a DiagonalizedOperator: all pairwise sums of the factor eigenvalues, with no factorisation of the full operator.
• Anything else: dense eigvalsh.

The eigenvalue paths drop eigenvalues ≤ rcond · λ_max.

Constants are constant. log|τR|_+ = rank(R) log τ + log|R|_+, so log|R|_+ never depends on hyperparameters: compute it once, outside any θ loop, and keep it.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Symmetric positive-semidefinite operator A.	required
null_space	Float[ArrayLike, 'N c'] | Float[ArrayLike, ' N'] | None	A basis of ker A, shape (N, c) (or (N,) for c = 1); for a connected graph Laplacian, 1/√N. kernellib's graph_null_space gives one column per connected component.	None
structure	Literal['laplacian'] | None	"laplacian" to use the matrix-tree theorem on a weighted graph Laplacian (SparseOperator, 1 × 1-block BlockTriDiag or KroneckerSum).	None
rcond	float | None	Relative cut-off for the eigenvalue paths. Defaults to n · eps of the dtype, with n the largest dense eigenproblem solved (a factor's size for a KroneckerSum), as in numpy.linalg.matrix_rank.	None
strategy	AbstractLogdetStrategy | None	Log-determinant strategy for the null-space path (of A + BBᵀ) and the Laplacian path (of the minor). Not used by the eigenvalue paths.	None

Returns:

Type	Description
Float[Array, '']	Scalar \(\log|A|_+\).

Raises:

Type	Description
ValueError	If both null_space and structure are given, if strategy is given for an eigenvalue path, if structure is unknown, or if null_space has the wrong number of rows.
TypeError	If structure="laplacian" is given an operator it has no Laplacian path for.

Examples:

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

# ICAR normalising constant on the path graph 0 - 1 - 2 - 3 (computed once)
n = 4
R = gaussx.SparseOperator.from_coo(
    np.r_[np.arange(n), np.arange(1, n)],
    np.r_[np.arange(n), np.arange(n - 1)],
    jnp.r_[jnp.array([1.0, 2.0, 2.0, 1.0]), -jnp.ones(n - 1)],
    (n, n),
    symmetric=True,
)
half_log_pdet = 0.5 * gaussx.pseudo_logdet(R, structure="laplacian")
# pdet(R) = n · (one spanning tree) = 4
assert jnp.allclose(2 * half_log_pdet, jnp.log(4.0))

# The same through its null space, constants / √n
ones = jnp.ones(n) / jnp.sqrt(n)
assert jnp.allclose(gaussx.pseudo_logdet(R, null_space=ones), jnp.log(4.0))

# On a grid there is no factorisation: Σ log of the non-zero λ^H_i + λ^W_j
L = gaussx.rw1_structure(5).as_matrix()
grid = gaussx.KroneckerSum(
    lx.MatrixLinearOperator(L, lx.symmetric_tag),
    lx.MatrixLinearOperator(L, lx.symmetric_tag),
)
half_log_pdet_grid = 0.5 * gaussx.pseudo_logdet(grid)

Source code in src/gaussx/_primitives/_logdet.py

def pseudo_logdet(
    operator: lx.AbstractLinearOperator,
    *,
    null_space: Float[ArrayLike, "N c"] | Float[ArrayLike, " N"] | None = None,
    structure: Literal["laplacian"] | None = None,
    rcond: float | None = None,
    strategy: AbstractLogdetStrategy | None = None,
) -> Float[Array, ""]:
    r"""Log pseudo-determinant: the log of the product of the non-zero eigenvalues.

    For a symmetric positive-semidefinite ``A``,
    $\log|A|_+ = \sum_{\lambda_i > 0}\log\lambda_i$. Intrinsic GMRFs (Besag /
    ICAR, RW1, RW2) have singular structure matrices ``R``, and
    $\tfrac12\log|R|_+$ is their normalising constant. Paths, cheapest first:

    - ``structure="laplacian"`` (a weighted graph Laplacian, e.g. Besag or
      RW1): the matrix-tree theorem. Every principal $(n_c-1)$-minor of a
      connected component's Laplacian $L_c$ equals its weighted spanning-tree
      count, and $\operatorname{pdet}(L_c) = n_c$ times it, so
      $\log|L|_+ = \sum_c (\log n_c + \log|L_c^{(-k_c)}|)$. The minors are
      taken by replacing one node per component with a unit diagonal, which
      keeps the sparsity pattern, so it is **one** sparse Cholesky
      (`SparseCholeskySolver` by default; G4) for a `SparseOperator`, or one
      banded Cholesky for a `BlockTriDiag` with ``1 × 1`` blocks (a path, as
      from `rw1_structure`). Components are found from the sparsity pattern
      on the host (once per pattern), so every stored off-diagonal weight
      must be non-zero. Gradients are exact. A `KroneckerSum` (a grid
      Laplacian) takes the eigenvalue path below instead.
    - ``null_space`` given: $\operatorname{pdet}(A) = \det(A + NN^\top)$
      for orthonormal $N$ spanning $\ker A$ (each zero eigenvalue becomes
      1). Any basis $B$ of the kernel works, since
      $\det(A + BB^\top) = \operatorname{pdet}(A)\det(B^\top B)$ and the
      second factor is subtracted. ``A + BBᵀ`` is a matvec-only lineax sum,
      so ``strategy=gaussx.SLQLogdet()`` estimates it matrix-free; the
      default is `gaussx.logdet` of it (dense). The matrix determinant lemma
      does not apply because ``A`` is singular.
    - A `KroneckerSum` (recursively, with dense, diagonal or
      `DiagonalizedOperator` factors) or a `DiagonalizedOperator`: all
      pairwise sums of the factor eigenvalues, with no factorisation of the
      full operator.
    - Anything else: dense ``eigvalsh``.

    The eigenvalue paths drop eigenvalues ``≤ rcond · λ_max``.

    **Constants are constant.** ``log|τR|_+ = rank(R) log τ + log|R|_+``, so
    ``log|R|_+`` never depends on hyperparameters: compute it once, outside
    any θ loop, and keep it.

    Args:
        operator: Symmetric positive-semidefinite operator ``A``.
        null_space: A basis of ``ker A``, shape ``(N, c)`` (or ``(N,)`` for
            ``c = 1``); for a connected graph Laplacian, ``1/√N``. kernellib's
            ``graph_null_space`` gives one column per connected component.
        structure: ``"laplacian"`` to use the matrix-tree theorem on a
            weighted graph Laplacian (`SparseOperator`, ``1 × 1``-block
            `BlockTriDiag` or `KroneckerSum`).
        rcond: Relative cut-off for the eigenvalue paths. Defaults to
            ``n · eps`` of the dtype, with ``n`` the largest dense
            eigenproblem solved (a factor's size for a `KroneckerSum`), as in
            ``numpy.linalg.matrix_rank``.
        strategy: Log-determinant strategy for the null-space path (of
            ``A + BBᵀ``) and the Laplacian path (of the minor). Not used by
            the eigenvalue paths.

    Returns:
        Scalar $\log|A|_+$.

    Raises:
        ValueError: If both ``null_space`` and ``structure`` are given, if
            ``strategy`` is given for an eigenvalue path, if ``structure`` is
            unknown, or if ``null_space`` has the wrong number of rows.
        TypeError: If ``structure="laplacian"`` is given an operator it has
            no Laplacian path for.

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

        # ICAR normalising constant on the path graph 0 - 1 - 2 - 3 (computed once)
        n = 4
        R = gaussx.SparseOperator.from_coo(
            np.r_[np.arange(n), np.arange(1, n)],
            np.r_[np.arange(n), np.arange(n - 1)],
            jnp.r_[jnp.array([1.0, 2.0, 2.0, 1.0]), -jnp.ones(n - 1)],
            (n, n),
            symmetric=True,
        )
        half_log_pdet = 0.5 * gaussx.pseudo_logdet(R, structure="laplacian")
        # pdet(R) = n · (one spanning tree) = 4
        assert jnp.allclose(2 * half_log_pdet, jnp.log(4.0))

        # The same through its null space, constants / √n
        ones = jnp.ones(n) / jnp.sqrt(n)
        assert jnp.allclose(gaussx.pseudo_logdet(R, null_space=ones), jnp.log(4.0))

        # On a grid there is no factorisation: Σ log of the non-zero λ^H_i + λ^W_j
        L = gaussx.rw1_structure(5).as_matrix()
        grid = gaussx.KroneckerSum(
            lx.MatrixLinearOperator(L, lx.symmetric_tag),
            lx.MatrixLinearOperator(L, lx.symmetric_tag),
        )
        half_log_pdet_grid = 0.5 * gaussx.pseudo_logdet(grid)
        ```
    """
    if structure not in (None, "laplacian"):
        raise ValueError(f"structure must be None or 'laplacian', got {structure!r}.")
    if null_space is not None and structure is not None:
        raise ValueError("Pass either null_space or structure, not both.")
    if isinstance(operator, lx.TaggedLinearOperator):
        operator = operator.operator
    if null_space is not None:
        return _pseudo_logdet_null_space(operator, null_space, strategy)
    if structure == "laplacian":
        if isinstance(operator, SparseOperator):
            return _pseudo_logdet_laplacian_sparse(operator, strategy)
        if isinstance(operator, BlockTriDiag):
            return _pseudo_logdet_laplacian_banded(operator, strategy)
        if not isinstance(operator, KroneckerSum | DiagonalizedOperator):
            raise TypeError(
                "structure='laplacian' needs a SparseOperator, a BlockTriDiag "
                "with 1 x 1 blocks or a KroneckerSum, got "
                f"{type(operator).__name__}; pass null_space instead."
            )
    if strategy is not None:
        raise ValueError(
            "strategy is used by the null_space and structure='laplacian' "
            "paths only; the eigenvalue paths are exact."
        )
    eigenvalues, size = _spectrum(operator)
    if rcond is None:
        rcond = size * float(jnp.finfo(eigenvalues.dtype).eps)
    keep = eigenvalues > rcond * jnp.max(eigenvalues)
    logs = jnp.log(jnp.where(keep, eigenvalues, 1))
    return jnp.sum(jnp.where(keep, logs, 0))

Trace & diagonal

Exact where structure allows; stochastic (Hutchinson / Hutch++ / XTrace probing) for matrix-free operators. trace_and_diag shares one probe pass between both estimates.

Hutch++ (algorithm="hutchpp", Meyer et al., 2021) and XTrace (algorithm="xtrace", Epperly et al., 2024) deflate a randomized low-rank part, \(\operatorname{tr}(A) = \operatorname{tr}(Q^\top AQ) + \operatorname{tr}\big((I-QQ^\top)A(I-QQ^\top)\big)\), and probe only the remainder. On a decaying spectrum they reach relative error \(\varepsilon\) in \(O(1/\varepsilon)\) matvecs, against Hutchinson's \(O(1/\varepsilon^2)\). Examples are traces of kernel matrices, of \(A^{-1}B\) in GP hyperparameter gradients, or of Hessians.

Structured linear algebra and Gaussian primitives for JAX.

trace(operator: lx.AbstractLinearOperator, *, stochastic: bool = False, num_probes: int = 20, key: jax.Array | None = None, sampler: SamplerName | None = None, algorithm: Literal['hutchinson', 'hutchpp', 'xtrace'] = 'hutchinson') -> Float[Array, '']

Compute the trace of an operator.

When stochastic=True, uses a stochastic estimator — only requires matvec access, no materialization. Hutchinson averages \(z^\top A z\) over probes, with variance \(O(\|A\|_F^2 / m)\) after \(m\) matvecs. The variance-reduced estimators deflate a low-rank part first,

\[ \operatorname{tr}(A) = \operatorname{tr}(Q^\top A Q) + \operatorname{tr}\big((I - QQ^\top) A (I - QQ^\top)\big), \]

with \(Q\) from a randomized range finder: the first term is exact and only the (small) remainder is probed. For PSD \(A\) with a decaying spectrum that reaches relative error \(\varepsilon\) in \(O(1/\varepsilon)\) matvecs instead of \(O(1/\varepsilon^2)\).

• "hutchpp" (Hutch++, Meyer, Musco, Musco & Woodruff, 2021): num_probes is the total matvec budget \(m \ge 3\), a third each for the range sketch, the exact projection and Hutchinson on the remainder.
• "xtrace" (Epperly, Tropp & Webber, 2024; matfree's leave_one_out_xtrace): every probe both builds \(Q\) and estimates the remainder, by leave-one-out; num_probes probes cost 2 * num_probes matvecs. Usually the most accurate per matvec.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A square linear operator.	required
stochastic	bool	If True, use stochastic trace estimation.	False
num_probes	int	Number of probe vectors for stochastic mode.	20
key	Array | None	PRNG key for stochastic mode.	None
sampler	SamplerName | None	Probe distribution ("signs", "normal", "sphere"). Defaults to "signs" for Hutchinson and "sphere" for XTrace (which requires a rotationally invariant distribution).	None
algorithm	Literal['hutchinson', 'hutchpp', 'xtrace']	"hutchinson" (Monte-Carlo), "hutchpp" (Hutch++, low-rank deflation plus Hutchinson) or "xtrace" (leave-one-out, Epperly et al. 2024 — much lower variance for the same number of matvecs).	'hutchinson'

Returns:

Type	Description
Float[Array, '']	Scalar trace value (exact or estimated).

Raises:

Type	Description
ValueError	For an unknown algorithm, algorithm="xtrace" with sampler="signs", or algorithm="hutchpp" with num_probes < 3.

References

Meyer, R. A., Musco, C., Musco, C. & Woodruff, D. P. (2021). Hutch++: optimal stochastic trace estimation. SOSA, 142-155.

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

Examples:

>>> import einx, jax.numpy as jnp, jax.random as jr, lineax as lx
>>> import gaussx as gx
>>> U, _ = jnp.linalg.qr(jr.normal(jr.key(0), (200, 200)))
>>> lam = 1.0 / jnp.arange(1.0, 201.0) ** 2  # decaying spectrum
>>> U_lam = einx.multiply("i k, k -> i k", U, lam)
>>> A = lx.MatrixLinearOperator(einx.dot("i k, j k -> i j", U_lam, U))
>>> est = gx.trace(A, stochastic=True, num_probes=30, algorithm="hutchpp")
>>> bool(jnp.abs(est - jnp.sum(lam)) / jnp.sum(lam) < 0.05)
True

Structured operators take an exact fast path:

>>> import gaussx
>>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
>>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
>>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
>>> float(gaussx.trace(K))  # trace(A) · trace(B)
27.0

Source code in src/gaussx/_primitives/_trace.py

def trace(
    operator: lx.AbstractLinearOperator,
    *,
    stochastic: bool = False,
    num_probes: int = 20,
    key: jax.Array | None = None,
    sampler: SamplerName | None = None,
    algorithm: Literal["hutchinson", "hutchpp", "xtrace"] = "hutchinson",
) -> Float[Array, ""]:
    r"""Compute the trace of an operator.

    When ``stochastic=True``, uses a stochastic estimator — only requires
    matvec access, no materialization. Hutchinson averages $z^\top A z$
    over probes, with variance $O(\|A\|_F^2 / m)$ after $m$ matvecs. The
    variance-reduced estimators deflate a low-rank part first,

    $$
    \operatorname{tr}(A) = \operatorname{tr}(Q^\top A Q)
    + \operatorname{tr}\big((I - QQ^\top) A (I - QQ^\top)\big),
    $$

    with $Q$ from a randomized range finder: the first term is exact and
    only the (small) remainder is probed. For PSD $A$ with a decaying
    spectrum that reaches relative error $\varepsilon$ in
    $O(1/\varepsilon)$ matvecs instead of $O(1/\varepsilon^2)$.

    - ``"hutchpp"`` (Hutch++, Meyer, Musco, Musco & Woodruff, 2021):
      ``num_probes`` is the total matvec budget $m \ge 3$, a third each
      for the range sketch, the exact projection and Hutchinson on the
      remainder.
    - ``"xtrace"`` (Epperly, Tropp & Webber, 2024; matfree's
      ``leave_one_out_xtrace``): every probe both builds $Q$ and estimates
      the remainder, by leave-one-out; ``num_probes`` probes cost
      ``2 * num_probes`` matvecs. Usually the most accurate per matvec.

    Args:
        operator: A square linear operator.
        stochastic: If ``True``, use stochastic trace estimation.
        num_probes: Number of probe vectors for stochastic mode.
        key: PRNG key for stochastic mode.
        sampler: Probe distribution (``"signs"``, ``"normal"``,
            ``"sphere"``). Defaults to ``"signs"`` for Hutchinson and
            ``"sphere"`` for XTrace (which requires a rotationally
            invariant distribution).
        algorithm: ``"hutchinson"`` (Monte-Carlo), ``"hutchpp"``
            (Hutch++, low-rank deflation plus Hutchinson) or ``"xtrace"``
            (leave-one-out, Epperly et al. 2024 — much lower variance
            for the same number of matvecs).

    Returns:
        Scalar trace value (exact or estimated).

    Raises:
        ValueError: For an unknown ``algorithm``, ``algorithm="xtrace"``
            with ``sampler="signs"``, or ``algorithm="hutchpp"`` with
            ``num_probes < 3``.

    References:
        Meyer, R. A., Musco, C., Musco, C. & Woodruff, D. P. (2021).
        Hutch++: optimal stochastic trace estimation. *SOSA*, 142-155.

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

    Examples:
        >>> import einx, jax.numpy as jnp, jax.random as jr, lineax as lx
        >>> import gaussx as gx
        >>> U, _ = jnp.linalg.qr(jr.normal(jr.key(0), (200, 200)))
        >>> lam = 1.0 / jnp.arange(1.0, 201.0) ** 2  # decaying spectrum
        >>> U_lam = einx.multiply("i k, k -> i k", U, lam)
        >>> A = lx.MatrixLinearOperator(einx.dot("i k, j k -> i j", U_lam, U))
        >>> est = gx.trace(A, stochastic=True, num_probes=30, algorithm="hutchpp")
        >>> bool(jnp.abs(est - jnp.sum(lam)) / jnp.sum(lam) < 0.05)
        True

        Structured operators take an exact fast path:

        >>> import gaussx
        >>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
        >>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
        >>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
        >>> float(gaussx.trace(K))  # trace(A) · trace(B)
        27.0
    """

    # Every recursive call forwards the estimator options, so a wrapped or
    # structured matrix-free operator is never materialised (gh-320).
    def rec(op: lx.AbstractLinearOperator, k: jax.Array | None = key) -> Array:
        return trace(
            op,
            stochastic=stochastic,
            num_probes=num_probes,
            key=k,
            sampler=sampler,
            algorithm=algorithm,
        )

    def rec_all(ops) -> list[Array]:
        ops = tuple(ops)
        keys = split_keys(key, len(ops))
        return [rec(op, k) for op, k in zip(ops, keys, strict=True)]

    if isinstance(operator, lx.IdentityLinearOperator):
        return jnp.asarray(operator.in_size(), dtype=operator.in_structure().dtype)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return jnp.sum(lx.diagonal(operator))
    if isinstance(operator, DiagonalizedOperator):
        total = jnp.sum(operator.eigenvalues)
        return jnp.real(total) if operator.real_output else total
    if isinstance(operator, Toeplitz):
        # Constant diagonal c[0] (gh-373).
        return operator.in_size() * operator.column[0]
    if isinstance(operator, BlockDiag):
        return ft.reduce(jnp.add, rec_all(operator.operators))
    if isinstance(operator, Kronecker):
        # trace(A ⊗ B) = trace(A) · trace(B).
        return ft.reduce(jnp.multiply, rec_all(operator.operators))
    if isinstance(operator, BlockTriDiag | LowerBlockTriDiag | UpperBlockTriDiag):
        return _trace_block_tridiag(operator)
    if isinstance(operator, LowRankUpdate):
        return _trace_low_rank(operator, rec(operator.base))
    if isinstance(operator, KroneckerSum):
        # trace(A ⊕ B) = n_b · trace(A) + n_a · trace(B).
        trace_a, trace_b = rec_all((operator.A, operator.B))
        return operator.B.out_size() * trace_a + operator.A.out_size() * trace_b
    if isinstance(operator, SumOfKroneckers):
        return ft.reduce(jnp.add, rec_all(operator.operators))
    if isinstance(operator, lx.TaggedLinearOperator):
        return rec(operator.operator)
    if isinstance(operator, lx.AddLinearOperator):
        first, second = rec_all((operator.operator1, operator.operator2))
        return first + second
    if isinstance(operator, lx.MulLinearOperator):
        return operator.scalar * rec(operator.operator)
    if isinstance(operator, lx.DivLinearOperator):
        return rec(operator.operator) / operator.scalar
    if isinstance(operator, lx.NegLinearOperator):
        return -rec(operator.operator)
    if stochastic:
        return _trace_stochastic(operator, num_probes, key, sampler, algorithm)
    return jnp.trace(operator.as_matrix())

diag(operator: lx.AbstractLinearOperator, *, stochastic: bool = False, num_probes: int = 20, key: jax.Array | None = None, sampler: SamplerName = 'signs') -> Float[Array, ' n']

Extract the diagonal of an operator as a 1D array.

When stochastic=True, uses Hutchinson's diagonal estimator via matfree — only requires matvec access, no materialization.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A linear operator.	required
stochastic	bool	If True, use stochastic diagonal estimation.	False
num_probes	int	Number of probe vectors for stochastic mode.	20
key	Array | None	PRNG key for stochastic mode.	None
sampler	SamplerName	Probe distribution for stochastic mode ("signs", "normal", "sphere").	'signs'

Returns:

Type	Description
Float[Array, ' n']	1D array of diagonal entries (exact or estimated).

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
>>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
>>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
>>> [float(v) for v in gaussx.diag(K)]  # diag(A) ⊗ diag(B)
[4.0, 5.0, 8.0, 10.0]

Source code in src/gaussx/_primitives/_diag.py

def diag(
    operator: lx.AbstractLinearOperator,
    *,
    stochastic: bool = False,
    num_probes: int = 20,
    key: jax.Array | None = None,
    sampler: SamplerName = "signs",
) -> Float[Array, " n"]:
    """Extract the diagonal of an operator as a 1D array.

    When ``stochastic=True``, uses Hutchinson's diagonal estimator
    via matfree — only requires matvec access, no materialization.

    Args:
        operator: A linear operator.
        stochastic: If ``True``, use stochastic diagonal estimation.
        num_probes: Number of probe vectors for stochastic mode.
        key: PRNG key for stochastic mode.
        sampler: Probe distribution for stochastic mode (``"signs"``,
            ``"normal"``, ``"sphere"``).

    Returns:
        1D array of diagonal entries (exact or estimated).

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
        >>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
        >>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
        >>> [float(v) for v in gaussx.diag(K)]  # diag(A) ⊗ diag(B)
        [4.0, 5.0, 8.0, 10.0]
    """

    # Every recursive call forwards the estimator options, so a wrapped or
    # structured matrix-free operator is never materialised (gh-320).
    def rec(op: lx.AbstractLinearOperator, k: jax.Array | None = key) -> Array:
        return diag(
            op, stochastic=stochastic, num_probes=num_probes, key=k, sampler=sampler
        )

    def rec_all(ops) -> list[Array]:
        ops = tuple(ops)
        return [
            rec(op, k) for op, k in zip(ops, split_keys(key, len(ops)), strict=True)
        ]

    if isinstance(operator, lx.IdentityLinearOperator):
        return jnp.ones(operator.in_size(), dtype=operator.in_structure().dtype)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return lx.diagonal(operator)
    if isinstance(operator, BlockDiag):
        return jnp.concatenate(rec_all(operator.operators))
    if isinstance(operator, Kronecker):
        # diag(A ⊗ B) = diag(A) ⊗ diag(B).
        return ft.reduce(jnp.kron, rec_all(operator.operators))
    if isinstance(operator, BlockTriDiag | LowerBlockTriDiag | UpperBlockTriDiag):
        return _diag_block_tridiag(operator)
    if isinstance(operator, LowRankUpdate):
        return _diag_low_rank(operator, rec(operator.base))
    if isinstance(operator, KroneckerSum):
        return _diag_kronecker_sum(*rec_all((operator.A, operator.B)))
    if isinstance(operator, SparseOperator | SpectralFunction):
        return operator.diagonal()
    if isinstance(operator, Toeplitz):
        # A symmetric Toeplitz matrix has the constant diagonal c[0] (gh-373).
        return jnp.full(operator.in_size(), operator.column[0])
    if isinstance(operator, DiagonalizedOperator) and _is_fft_pair(operator):
        # F⁻¹ diag(λ) F has the constant diagonal mean(λ) (gh-373).
        mean = jnp.mean(operator.eigenvalues)
        mean = jnp.real(mean) if operator.real_output else mean
        return jnp.full(operator.in_size(), mean)
    if isinstance(operator, SumOfKroneckers):
        return ft.reduce(jnp.add, rec_all(operator.operators))
    if isinstance(operator, lx.TaggedLinearOperator):
        return rec(operator.operator)
    if isinstance(operator, lx.AddLinearOperator):
        first, second = rec_all((operator.operator1, operator.operator2))
        return first + second
    if isinstance(operator, lx.MulLinearOperator):
        return operator.scalar * rec(operator.operator)
    if isinstance(operator, lx.DivLinearOperator):
        return rec(operator.operator) / operator.scalar
    if isinstance(operator, lx.NegLinearOperator):
        return -rec(operator.operator)
    if stochastic:
        return _diag_stochastic(operator, num_probes, key, sampler)
    if isinstance(operator, DiagonalizedOperator):
        warn_dense_fallback(
            "diag(DiagonalizedOperator) with a non-FFT transform pair "
            "materialises the operator; diag(..., stochastic=True) estimates "
            "it from matvecs."
        )
    return jnp.diag(operator.as_matrix())

trace_and_diag(operator: lx.AbstractLinearOperator, *, num_probes: int = 20, key: jax.Array | None = None, sampler: SamplerName = 'signs') -> tuple[Float[Array, ''], Float[Array, ' n']]

Jointly estimate the trace and diagonal from one probe pass.

Halves the matvec budget relative to calling trace(..., stochastic=True) and diag(..., stochastic=True) separately — both statistics are accumulated from the same A @ probe products.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A square linear operator.	required
num_probes	int	Number of probe vectors.	20
key	Array | None	PRNG key. If None, uses jax.random.PRNGKey(0).	None
sampler	SamplerName	Probe distribution ("signs", "normal", "sphere").	'signs'

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, ' n']]	Tuple (trace_estimate, diagonal_estimate).

Source code in src/gaussx/_primitives/_trace.py

def trace_and_diag(
    operator: lx.AbstractLinearOperator,
    *,
    num_probes: int = 20,
    key: jax.Array | None = None,
    sampler: SamplerName = "signs",
) -> tuple[Float[Array, ""], Float[Array, " n"]]:
    """Jointly estimate the trace and diagonal from one probe pass.

    Halves the matvec budget relative to calling
    ``trace(..., stochastic=True)`` and ``diag(..., stochastic=True)``
    separately — both statistics are accumulated from the same
    ``A @ probe`` products.

    Args:
        operator: A square linear operator.
        num_probes: Number of probe vectors.
        key: PRNG key. If ``None``, uses ``jax.random.PRNGKey(0)``.
        sampler: Probe distribution (``"signs"``, ``"normal"``,
            ``"sphere"``).

    Returns:
        Tuple ``(trace_estimate, diagonal_estimate)``.
    """
    if key is None:
        key = jax.random.PRNGKey(0)

    n = operator.in_size()
    probe_fn = resolve_sampler(
        sampler, n, num_probes, dtype=operator.in_structure().dtype
    )
    integrand = matfree.stochtrace.monte_carlo_trace_and_diagonal()
    estimate = matfree.stochtrace.estimator_monte_carlo(integrand, probe_fn)
    result = estimate(operator.mv, key)
    return result["trace"], result["diagonal"]

Inverse, square root & spectral decompositions

inv and sqrt return lazy operators that route their matvecs through structured solves / Lanczos; eig, eigvals, and svd take an optional rank for partial (Krylov) decompositions.

Structured linear algebra and Gaussian primitives for JAX.

inv(operator: lx.AbstractLinearOperator, *, solver: lx.AbstractLinearSolver | None = None) -> lx.AbstractLinearOperator

Return a lazy inverse operator A^{-1}.

The returned operator computes A^{-1} v via solve(A, v) when mv is called. For structured operators, the inverse preserves structure.

Related to lineax.invert (lineax >= 0.1.1), which wraps lx.linear_solve in a FunctionLinearOperator. The gaussx fallback InverseOperator differs in that its matvec routes through the structured gaussx solve dispatch, and its as_matrix uses a Cholesky path for PSD operators.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	An invertible linear operator.	required
solver	AbstractLinearSolver | None	Optional lineax solver for the fallback InverseOperator.	None

Returns:

Type	Description
AbstractLinearOperator	An operator representing A^{-1}.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
>>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
>>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
>>> K_inv = gaussx.inv(K)  # Kronecker(inv(A), inv(B)), never densified
>>> type(K_inv).__name__
'Kronecker'
>>> [round(float(v), 4) for v in K_inv.mv(jnp.ones(4))]
[0.25, 0.2, 0.125, 0.1]

Source code in src/gaussx/_primitives/_inv.py

def inv(
    operator: lx.AbstractLinearOperator,
    *,
    solver: lx.AbstractLinearSolver | None = None,
) -> lx.AbstractLinearOperator:
    """Return a lazy inverse operator A^{-1}.

    The returned operator computes A^{-1} v via ``solve(A, v)``
    when ``mv`` is called. For structured operators, the inverse
    preserves structure.

    Related to ``lineax.invert`` (lineax >= 0.1.1), which wraps
    ``lx.linear_solve`` in a ``FunctionLinearOperator``. The gaussx
    fallback ``InverseOperator`` differs in that its matvec routes
    through the *structured* gaussx ``solve`` dispatch, and its
    ``as_matrix`` uses a Cholesky path for PSD operators.

    Args:
        operator: An invertible linear operator.
        solver: Optional lineax solver for the fallback InverseOperator.

    Returns:
        An operator representing A^{-1}.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> A = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0]))
        >>> B = lx.DiagonalLinearOperator(jnp.array([4.0, 5.0]))
        >>> K = gaussx.Kronecker(A, B)  # diag(4, 5, 8, 10)
        >>> K_inv = gaussx.inv(K)  # Kronecker(inv(A), inv(B)), never densified
        >>> type(K_inv).__name__
        'Kronecker'
        >>> [round(float(v), 4) for v in K_inv.mv(jnp.ones(4))]
        [0.25, 0.2, 0.125, 0.1]
    """
    if isinstance(operator, lx.IdentityLinearOperator):
        return operator
    if isinstance(operator, InverseOperator):
        # (A⁻¹)⁻¹ = A (gh-349).
        return operator.original
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _inv_diagonal(operator)
    if isinstance(operator, DiagonalizedOperator):
        return _inv_diagonalised(operator)
    if isinstance(operator, BlockDiag):
        return _inv_block_diag(operator)
    if isinstance(operator, Kronecker):
        return _inv_kronecker(operator)
    if isinstance(operator, LowRankUpdate) and _is_square(operator.base):
        if operator.rank == 0:
            return inv(operator.base, solver=solver)
        c = orthonormal_scaled_identity(operator)
        if c is not None:
            return _inv_low_rank_orthonormal(operator, c)
        # Decided from static structure only (gh-328): a value check on
        # ``U == V`` is lost once the operator crosses a jit boundary.
        if lx.is_symmetric(operator) and operator.symmetric_factors:
            return _inv_low_rank_symmetric(operator, solver)
        return _inv_low_rank_general(operator, solver)
    if isinstance(operator, lx.MulLinearOperator):
        return (1.0 / operator.scalar) * inv(operator.operator, solver=solver)
    if isinstance(operator, lx.DivLinearOperator):
        return operator.scalar * inv(operator.operator, solver=solver)
    if isinstance(operator, lx.NegLinearOperator):
        return -inv(operator.operator, solver=solver)
    if isinstance(operator, lx.ComposedLinearOperator) and (
        operator.operator1.in_size() == operator.operator1.out_size()
        and operator.operator2.in_size() == operator.operator2.out_size()
    ):
        # (A B)^{-1} = B^{-1} A^{-1}
        return inv(operator.operator2, solver=solver) @ inv(
            operator.operator1, solver=solver
        )
    return InverseOperator(operator, solver)

sqrt(operator: lx.AbstractLinearOperator, *, lanczos_order: int | None = None) -> lx.AbstractLinearOperator

Compute matrix square root S such that S @ S = A.

Requires A to be positive semi-definite.

When lanczos_order is given, returns a lazy SqrtOperator that computes sqrt(A) @ v via matfree Lanczos without materializing the full square root matrix.

A scalar multiple c · A or A / c (c > 0) takes the root of A and scales it by √c, keeping its structured type. A concretely negative multiple such as -A is not PSD and raises.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A PSD linear operator.	required
lanczos_order	int | None	Order of Lanczos iteration for matrix-free sqrt. If None, uses dense eigendecomposition for most operators, except SumOfKroneckers — where None falls back to a Lanczos sqrt with the module default order (no closed-form sqrt exists for the sum-of-Kronecker structure). Pass an explicit lanczos_order to override the default rank.	None

Returns:

Type	Description
AbstractLinearOperator	Operator S satisfying S @ S = A.

Raises:

Type	Description
ValueError	If operator is a concretely negative multiple of a non-NSD operator, such as -A.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> psd = lx.positive_semidefinite_tag
>>> A = lx.MatrixLinearOperator(jnp.array([[2.0, 1.0], [1.0, 2.0]]), psd)
>>> B = lx.MatrixLinearOperator(jnp.array([[3.0, 0.5], [0.5, 1.0]]), psd)
>>> S = gaussx.sqrt(gaussx.BlockDiag(A, B))  # one root per block
>>> type(S).__name__
'BlockDiag'
>>> M = gaussx.BlockDiag(A, B).as_matrix()
>>> bool(jnp.allclose((S @ S).as_matrix(), M, atol=1e-5))
True

Source code in src/gaussx/_primitives/_sqrt.py

def sqrt(
    operator: lx.AbstractLinearOperator,
    *,
    lanczos_order: int | None = None,
) -> lx.AbstractLinearOperator:
    """Compute matrix square root S such that S @ S = A.

    Requires A to be positive semi-definite.

    When ``lanczos_order`` is given, returns a lazy ``SqrtOperator``
    that computes ``sqrt(A) @ v`` via matfree Lanczos without
    materializing the full square root matrix.

    A scalar multiple ``c · A`` or ``A / c`` (``c > 0``) takes the root of
    ``A`` and scales it by ``√c``, keeping its structured type. A concretely
    negative multiple such as ``-A`` is not PSD and raises.

    Args:
        operator: A PSD linear operator.
        lanczos_order: Order of Lanczos iteration for matrix-free
            sqrt. If ``None``, uses dense eigendecomposition for most
            operators, *except* `SumOfKroneckers` — where ``None``
            falls back to a Lanczos sqrt with the module default order
            (no closed-form sqrt exists for the sum-of-Kronecker
            structure). Pass an explicit ``lanczos_order`` to override
            the default rank.

    Returns:
        Operator S satisfying S @ S = A.

    Raises:
        ValueError: If ``operator`` is a concretely negative multiple of a
            non-NSD operator, such as ``-A``.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> psd = lx.positive_semidefinite_tag
        >>> A = lx.MatrixLinearOperator(jnp.array([[2.0, 1.0], [1.0, 2.0]]), psd)
        >>> B = lx.MatrixLinearOperator(jnp.array([[3.0, 0.5], [0.5, 1.0]]), psd)
        >>> S = gaussx.sqrt(gaussx.BlockDiag(A, B))  # one root per block
        >>> type(S).__name__
        'BlockDiag'
        >>> M = gaussx.BlockDiag(A, B).as_matrix()
        >>> bool(jnp.allclose((S @ S).as_matrix(), M, atol=1e-5))
        True
    """
    if isinstance(operator, lx.IdentityLinearOperator):
        return operator
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _sqrt_diagonal(operator)
    if isinstance(operator, lx.TaggedLinearOperator):
        return sqrt(operator.operator, lanczos_order=lanczos_order)
    if isinstance(
        operator, lx.MulLinearOperator | lx.DivLinearOperator | lx.NegLinearOperator
    ):
        return scaled_root(
            operator,
            lambda base: sqrt(base, lanczos_order=lanczos_order),
            # The unscaled route keeps a requested matrix-free Lanczos root.
            _sqrt_dense
            if lanczos_order is None
            else lambda op: SqrtOperator(op, lanczos_order),
            "sqrt",
        )
    if isinstance(operator, DiagonalizedOperator):
        # S = V⁻¹ √Λ V satisfies S @ S = A for any diagonalisable A.
        return operator.with_eigenvalues(jnp.sqrt(operator.eigenvalues))
    if isinstance(operator, BlockDiag):
        return _sqrt_block_diag(operator)
    if isinstance(operator, Kronecker):
        return _sqrt_kronecker(operator)
    if isinstance(operator, KroneckerSum):
        return _sqrt_kronecker_sum(operator)
    if isinstance(operator, LowRankUpdate) and operator.rank > 0:
        c = orthonormal_scaled_identity(operator)
        if c is not None:
            return _sqrt_low_rank_orthonormal(operator, c)
    if isinstance(operator, SumOfKroneckers):
        return _sqrt_sum_kronecker(
            operator,
            lanczos_order=(
                _DEFAULT_LANCZOS_ORDER if lanczos_order is None else lanczos_order
            ),
        )
    if lanczos_order is not None:
        return SqrtOperator(operator, lanczos_order)
    return _sqrt_dense(operator)

eig(operator: lx.AbstractLinearOperator, *, rank: int | None = None, method: Literal['lanczos', 'randomized'] = 'lanczos', key: jax.Array | None = None) -> tuple[Array, Array]

Compute eigenvalues and eigenvectors.

For symmetric operators returns real eigenvalues via eigh.

rank follows one rule shared with eigvals and gaussx.svd: an operator with a structured decomposition (diagonal, gaussx.BlockDiag, gaussx.Kronecker, gaussx.KroneckerSum) is decomposed exactly and truncated to its rank largest eigenvalues (by real part, descending). Any other operator gets a partial eigendecomposition of a symmetric operator without matrix materialization, by method:

• "lanczos" (default): matfree Lanczos.
• "randomized": randomized_eigh with its defaults (oversample=10, n_power_iter=2, which="largest"); call it directly to tune them. Use n_power_iter >= 2 for slowly decaying spectra.

Randomized methods target the top of the spectrum (the eigenvalues of largest magnitude).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A square linear operator.	required
rank	int | None	Number of eigenvalues to compute. If None, computes the full eigendecomposition.	None
method	Literal['lanczos', 'randomized']	Partial-eig algorithm, "lanczos" or "randomized". Only used when rank is given.	'lanczos'
key	Array | None	PRNG key for the initial random vector (Lanczos) or the Gaussian test matrix (randomized) when using partial eig. If None, uses jax.random.PRNGKey(0).	None

Returns:

Type	Description
Array	Tuple (eigenvalues, eigenvectors) where eigenvalues has
Array	shape (K,) and eigenvectors has shape (N, K).

Raises:

Type	Description
ValueError	If method is invalid, or "randomized" without rank.

Source code in src/gaussx/_primitives/_eig.py

def eig(
    operator: lx.AbstractLinearOperator,
    *,
    rank: int | None = None,
    method: Literal["lanczos", "randomized"] = "lanczos",
    key: jax.Array | None = None,
) -> tuple[Array, Array]:
    """Compute eigenvalues and eigenvectors.

    For symmetric operators returns real eigenvalues via ``eigh``.

    ``rank`` follows one rule shared with `eigvals` and `gaussx.svd`: an
    operator with a structured decomposition (diagonal, `gaussx.BlockDiag`,
    `gaussx.Kronecker`, `gaussx.KroneckerSum`) is decomposed exactly and
    truncated to its ``rank`` largest eigenvalues (by real part,
    descending). Any other operator gets a partial eigendecomposition of a
    symmetric operator without matrix materialization, by ``method``:

    - ``"lanczos"`` (default): matfree Lanczos.
    - ``"randomized"``: `randomized_eigh` with its defaults
      (``oversample=10``, ``n_power_iter=2``, ``which="largest"``); call it
      directly to tune them. Use ``n_power_iter >= 2`` for slowly decaying
      spectra.

    Randomized methods target the **top** of the spectrum (the eigenvalues
    of largest magnitude).

    Args:
        operator: A square linear operator.
        rank: Number of eigenvalues to compute. If ``None``,
            computes the full eigendecomposition.
        method: Partial-eig algorithm, ``"lanczos"`` or ``"randomized"``.
            Only used when ``rank`` is given.
        key: PRNG key for the initial random vector (Lanczos) or the
            Gaussian test matrix (randomized) when using partial eig.
            If ``None``, uses ``jax.random.PRNGKey(0)``.

    Returns:
        Tuple ``(eigenvalues, eigenvectors)`` where eigenvalues has
        shape ``(K,)`` and eigenvectors has shape ``(N, K)``.

    Raises:
        ValueError: If ``method`` is invalid, or ``"randomized"`` without
            ``rank``.
    """
    if method not in ("lanczos", "randomized"):
        raise ValueError(f"method must be 'lanczos' or 'randomized', got {method!r}.")
    if method == "randomized" and rank is None:
        raise ValueError("method='randomized' needs a rank.")
    return _eig(operator, rank, method, key, symmetric=False)

eigvals(operator: lx.AbstractLinearOperator, *, rank: int | None = None, key: jax.Array | None = None) -> Array

Compute eigenvalues only.

When rank is given, returns the top-k eigenvalues: exactly and truncated for an operator with a structured decomposition (as in eig), via matfree Lanczos without matrix materialization otherwise.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A square linear operator.	required
rank	int | None	Number of eigenvalues to compute.	None
key	Array | None	PRNG key for partial eigendecomposition.	None

Returns:

Type	Description
Array	Eigenvalues array of shape (K,).

Source code in src/gaussx/_primitives/_eig.py

def eigvals(
    operator: lx.AbstractLinearOperator,
    *,
    rank: int | None = None,
    key: jax.Array | None = None,
) -> Array:
    """Compute eigenvalues only.

    When ``rank`` is given, returns the top-k eigenvalues: exactly and
    truncated for an operator with a structured decomposition (as in
    `eig`), via matfree Lanczos without matrix materialization otherwise.

    Args:
        operator: A square linear operator.
        rank: Number of eigenvalues to compute.
        key: PRNG key for partial eigendecomposition.

    Returns:
        Eigenvalues array of shape ``(K,)``.
    """
    return _eigvals(operator, rank, key, symmetric=False)

svd(operator: lx.AbstractLinearOperator, *, rank: int | None = None, method: Literal['lanczos', 'randomized'] = 'lanczos', key: jax.Array | None = None) -> tuple[Float[Array, 'm k'], Float[Array, ' k'], Float[Array, 'k n']]

Compute the singular value decomposition A = U diag(s) V^T.

rank follows one rule shared with gaussx.eig and gaussx.eigvals: an operator with a structured SVD (diagonal, gaussx.Kronecker, gaussx.BlockDiag) is decomposed exactly and truncated to its rank largest singular values. Any other operator gets a partial (truncated) SVD without materializing the operator, by method:

• "lanczos" (default): matfree's Golub-Kahan bidiagonalization.
• "randomized": randomized_svd with its defaults (oversample=10, n_power_iter=2); call it directly to tune them. Use n_power_iter >= 2 for slowly decaying spectra.

Both target the top of the spectrum (the largest singular values).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A linear operator.	required
rank	int | None	Number of singular values to compute. If None, computes the full SVD (requires materialization).	None
method	Literal['lanczos', 'randomized']	Partial-SVD algorithm, "lanczos" or "randomized". Only used when rank is given.	'lanczos'
key	Array | None	PRNG key for the initial random vector (Lanczos) or the Gaussian test matrix (randomized) when using partial SVD. If None, uses jax.random.PRNGKey(0).	None

Returns:

Type	Description
Float[Array, 'm k']	Tuple (U, s, Vt) where U has shape (M, K),
Float[Array, ' k']	s has shape (K,), and Vt has shape (K, N).

Raises:

Type	Description
ValueError	If method is invalid, or "randomized" without rank.

Source code in src/gaussx/_primitives/_svd.py

def svd(
    operator: lx.AbstractLinearOperator,
    *,
    rank: int | None = None,
    method: Literal["lanczos", "randomized"] = "lanczos",
    key: jax.Array | None = None,
) -> tuple[Float[Array, "m k"], Float[Array, " k"], Float[Array, "k n"]]:
    """Compute the singular value decomposition ``A = U diag(s) V^T``.

    ``rank`` follows one rule shared with `gaussx.eig` and `gaussx.eigvals`:
    an operator with a structured SVD (diagonal, `gaussx.Kronecker`,
    `gaussx.BlockDiag`) is decomposed exactly and truncated to its ``rank``
    largest singular values. Any other operator gets a partial (truncated)
    SVD without materializing the operator, by ``method``:

    - ``"lanczos"`` (default): matfree's Golub-Kahan bidiagonalization.
    - ``"randomized"``: `randomized_svd` with its defaults
      (``oversample=10``, ``n_power_iter=2``); call it directly to tune
      them. Use ``n_power_iter >= 2`` for slowly decaying spectra.

    Both target the **top** of the spectrum (the largest singular values).

    Args:
        operator: A linear operator.
        rank: Number of singular values to compute. If ``None``,
            computes the full SVD (requires materialization).
        method: Partial-SVD algorithm, ``"lanczos"`` or ``"randomized"``.
            Only used when ``rank`` is given.
        key: PRNG key for the initial random vector (Lanczos) or the
            Gaussian test matrix (randomized) when using partial SVD.
            If ``None``, uses ``jax.random.PRNGKey(0)``.

    Returns:
        Tuple ``(U, s, Vt)`` where U has shape ``(M, K)``,
        s has shape ``(K,)``, and Vt has shape ``(K, N)``.

    Raises:
        ValueError: If ``method`` is invalid, or ``"randomized"`` without
            ``rank``.
    """
    if method not in ("lanczos", "randomized"):
        raise ValueError(f"method must be 'lanczos' or 'randomized', got {method!r}.")
    if method == "randomized" and rank is None:
        raise ValueError("method='randomized' needs a rank.")
    if isinstance(operator, lx.TaggedLinearOperator):
        return svd(operator.operator, rank=rank, method=method, key=key)
    structured = _svd_structured(operator)
    if structured is not None:
        # Exact structured SVD first, then the top ``rank`` (gh-383).
        if rank is None:
            return structured
        U, s, Vt = _sort_svd_descending(*structured)
        return U[:, :rank], s[:rank], Vt[:rank, :]
    if rank is not None:
        if method == "randomized":
            return randomized_svd(operator, rank, key=key)
        return _svd_partial(operator, rank, key)
    return _svd_dense(operator)

frobenius_norm(operator: lx.AbstractLinearOperator, *, stochastic: bool = False, num_probes: int = 20, key: jax.Array | None = None, sampler: SamplerName = 'signs') -> Float[Array, '']

Compute the Frobenius norm ||A||_F with structural dispatch.

Structured operators avoid materialization:

• Diagonal: vector 2-norm of the diagonal.
• BlockDiag: root of the sum of squared block norms.
• Kronecker: ||A (x) B||_F = ||A||_F * ||B||_F.
• Scaled/negated/tagged operators delegate to the wrapped operator.

When stochastic=True, estimates ||A||_F^2 = tr(A^T A) via matfree's Hutchinson estimator — matvec access only.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A linear operator.	required
stochastic	bool	If True, use stochastic estimation.	False
num_probes	int	Number of probe vectors for stochastic mode.	20
key	Array | None	PRNG key for stochastic mode.	None
sampler	SamplerName	Probe distribution for stochastic mode ("signs", "normal", "sphere").	'signs'

Returns:

Type	Description
Float[Array, '']	Scalar Frobenius norm (exact or estimated).

Source code in src/gaussx/_primitives/_frobenius.py

def frobenius_norm(
    operator: lx.AbstractLinearOperator,
    *,
    stochastic: bool = False,
    num_probes: int = 20,
    key: jax.Array | None = None,
    sampler: SamplerName = "signs",
) -> Float[Array, ""]:
    """Compute the Frobenius norm ``||A||_F`` with structural dispatch.

    Structured operators avoid materialization:

    - Diagonal: vector 2-norm of the diagonal.
    - BlockDiag: root of the sum of squared block norms.
    - Kronecker: ``||A (x) B||_F = ||A||_F * ||B||_F``.
    - Scaled/negated/tagged operators delegate to the wrapped operator.

    When ``stochastic=True``, estimates ``||A||_F^2 = tr(A^T A)`` via
    matfree's Hutchinson estimator — matvec access only.

    Args:
        operator: A linear operator.
        stochastic: If ``True``, use stochastic estimation.
        num_probes: Number of probe vectors for stochastic mode.
        key: PRNG key for stochastic mode.
        sampler: Probe distribution for stochastic mode (``"signs"``,
            ``"normal"``, ``"sphere"``).

    Returns:
        Scalar Frobenius norm (exact or estimated).
    """

    # Every recursive call forwards the estimator options, so a wrapped or
    # structured matrix-free operator is never materialised (gh-320).
    def rec(op: lx.AbstractLinearOperator, k: jax.Array | None = key) -> Array:
        return frobenius_norm(
            op, stochastic=stochastic, num_probes=num_probes, key=k, sampler=sampler
        )

    def rec_all(ops) -> list[Array]:
        ops = tuple(ops)
        keys = split_keys(key, len(ops))
        return [rec(op, k) for op, k in zip(ops, keys, strict=True)]

    if isinstance(operator, lx.IdentityLinearOperator):
        return jnp.sqrt(jnp.asarray(float(operator.in_size())))
    if isinstance(operator, lx.DiagonalLinearOperator):
        d = lx.diagonal(operator)
        return jnp.sqrt(jnp.sum(jnp.abs(d) ** 2))
    if isinstance(operator, BlockDiag):
        norms = jnp.stack(rec_all(operator.operators))
        return jnp.sqrt(jnp.sum(norms * norms))
    if isinstance(operator, Kronecker):
        return ft.reduce(jnp.multiply, rec_all(operator.operators))
    if isinstance(operator, lx.TaggedLinearOperator):
        return rec(operator.operator)
    if isinstance(operator, lx.MulLinearOperator):
        return jnp.abs(operator.scalar) * rec(operator.operator)
    if isinstance(operator, lx.DivLinearOperator):
        return rec(operator.operator) / jnp.abs(operator.scalar)
    if isinstance(operator, lx.NegLinearOperator):
        return rec(operator.operator)
    if stochastic:
        return _frobenius_stochastic(operator, num_probes, key, sampler)
    # |a|², so a complex operator gets a real norm.
    mat = operator.as_matrix()
    return jnp.sqrt(jnp.sum(jnp.abs(mat) ** 2))

submatrix(operator: lx.AbstractLinearOperator, row_idx: Int[Array, ' R'], col_idx: Int[Array, ' C']) -> Float[Array, 'R C']

Extract A[row_idx, col_idx] without forming the full matrix.

For structured operators, exploits the structure to avoid materializing the full (N, N) matrix when only a sub-block is needed (e.g., the conditional Gaussian extracts Sigma_AA, Sigma_AB, Sigma_BB from a joint covariance).

Currently dispatches on:

• lineax.DiagonalLinearOperator
• lineax.IdentityLinearOperator
• lineax.TaggedLinearOperator / MulLinearOperator / DivLinearOperator / NegLinearOperator / AddLinearOperator (recurse into the wrapped operators)
• gaussx.BlockDiag
• gaussx.Kronecker
• gaussx.LowRankUpdate (base sub-block plus U[r] diag(d) V[c]ᵀ)

Falls back to operator.as_matrix()[ix_(row_idx, col_idx)] for other operators.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator A, shape (N, N).	required
row_idx	Int[Array, ' R']	Row indices, shape (R,).	required
col_idx	Int[Array, ' C']	Column indices, shape (C,).	required

Returns:

Type	Description
Float[Array, 'R C']	Dense sub-matrix A[ix_(row_idx, col_idx)] of shape (R, C).

Source code in src/gaussx/_primitives/_submatrix.py

def submatrix(
    operator: lx.AbstractLinearOperator,
    row_idx: Int[Array, " R"],
    col_idx: Int[Array, " C"],
) -> Float[Array, "R C"]:
    """Extract ``A[row_idx, col_idx]`` without forming the full matrix.

    For structured operators, exploits the structure to avoid
    materializing the full ``(N, N)`` matrix when only a sub-block is
    needed (e.g., the conditional Gaussian extracts ``Sigma_AA``,
    ``Sigma_AB``, ``Sigma_BB`` from a joint covariance).

    Currently dispatches on:

    - `lineax.DiagonalLinearOperator`
    - `lineax.IdentityLinearOperator`
    - `lineax.TaggedLinearOperator` / ``MulLinearOperator`` /
      ``DivLinearOperator`` / ``NegLinearOperator`` / ``AddLinearOperator``
      (recurse into the wrapped operators)
    - `gaussx.BlockDiag`
    - `gaussx.Kronecker`
    - `gaussx.LowRankUpdate` (base sub-block plus ``U[r] diag(d) V[c]ᵀ``)

    Falls back to ``operator.as_matrix()[ix_(row_idx, col_idx)]`` for
    other operators.

    Args:
        operator: Linear operator A, shape ``(N, N)``.
        row_idx: Row indices, shape ``(R,)``.
        col_idx: Column indices, shape ``(C,)``.

    Returns:
        Dense sub-matrix ``A[ix_(row_idx, col_idx)]`` of shape ``(R, C)``.
    """
    n = operator.in_size()
    row_idx = _normalize_indices(row_idx, n)
    col_idx = _normalize_indices(col_idx, n)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return _submatrix_diagonal(operator, row_idx, col_idx)
    if isinstance(operator, BlockDiag):
        return _submatrix_block_diag(operator, row_idx, col_idx)
    if isinstance(operator, Kronecker):
        return _submatrix_kronecker(operator, row_idx, col_idx)
    dtype = _matrix_dtype(operator)
    if row_idx.shape[0] == 0 or col_idx.shape[0] == 0:
        # einx rejects zero-length axes; an empty block needs no structure.
        return jnp.zeros((row_idx.shape[0], col_idx.shape[0]), dtype=dtype)
    if isinstance(operator, LowRankUpdate) and _is_square(operator.base):
        return _submatrix_low_rank(operator, row_idx, col_idx)
    if isinstance(operator, lx.IdentityLinearOperator) and _is_square(operator):
        match = einx.equal("r, c -> r c", row_idx, col_idx)
        return match.astype(dtype)
    if isinstance(operator, lx.TaggedLinearOperator):
        return submatrix(operator.operator, row_idx, col_idx)
    if isinstance(operator, lx.MulLinearOperator):
        return operator.scalar * submatrix(operator.operator, row_idx, col_idx)
    if isinstance(operator, lx.DivLinearOperator):
        return submatrix(operator.operator, row_idx, col_idx) / operator.scalar
    if isinstance(operator, lx.NegLinearOperator):
        return -submatrix(operator.operator, row_idx, col_idx)
    if isinstance(operator, lx.AddLinearOperator):
        return submatrix(operator.operator1, row_idx, col_idx) + submatrix(
            operator.operator2, row_idx, col_idx
        )
    return operator.as_matrix()[jnp.ix_(row_idx, col_idx)]

Generalised eigenproblems

eigh_generalized solves \(A v = \lambda B v\) for symmetric \(A\) and PSD \(B\), returning \(B\)-orthonormal eigenvectors (the minimisers of \(\operatorname{tr}(Y^\top A Y)\) s.t. \(Y^\top B Y = I\) behind Laplacian eigenmaps, LPP and manifold alignment). It dispatches on \(B\): a positive diagonal stays matrix-free (Lanczos with rank=), a tagged positive-definite \(B\) is Cholesky-whitened, and a singular \(B\) has its null space eliminated by a Schur complement.

Structured linear algebra and Gaussian primitives for JAX.

eigh_generalized(A: lx.AbstractLinearOperator, B: lx.AbstractLinearOperator, *, rank: int | None = None, which: Literal['smallest', 'largest'] = 'smallest', rcond: float | None = None, key: PRNGKeyArray | None = None) -> tuple[Float[Array, ' K'], Float[Array, 'N K']]

Solve \(A v = \lambda B v\) for symmetric \(A\) and symmetric PSD \(B\).

The eigenvectors are \(B\)-orthonormal, \(V^\top B V = I\), so the smallest \(K\) of them minimise \(\operatorname{tr}(Y^\top A Y)\) subject to \(Y^\top B Y = I\) (graph embeddings, LPP, manifold alignment). Dispatch is on the structure of \(B\):

• lineax.DiagonalLinearOperator with positive entries: scale, \(S = B^{-1/2} A B^{-1/2}\), \(v = B^{-1/2} u\). This is matrix-free in \(A\): with rank= it runs Lanczos on \(x \mapsto B^{-1/2} A B^{-1/2} x\) and never materialises \(A\). A diagonal with zero entries (only checked outside jit) is sent to the singular path below.
• Positive definite, i.e. tagged lineax.positive_semidefinite_tag (lineax's tag for Cholesky-factorisable operators): Cholesky whitening \(C^{-1} A C^{-\top}\) with \(B = C C^\top\).
• Otherwise \(B\) is treated as PSD and possibly singular. \(B\) is eigendecomposed, \(B = U_+ S_+ U_+^\top\), and directions with eigenvalue \(\le\) rcond \(\cdot \max\) span \(\ker B\). Writing \(A_{\cdot\cdot}\) for the blocks of \(U^\top A U\), the \(\ker B\) rows give \(v_0 = -A_{00}^{-1} A_{0+} v_+\), and the finite eigenpairs solve the Schur-complement pencil \((A_{++} - A_{+0} A_{00}^{-1} A_{0+}) v_+ = \lambda S_+ v_+\). \(A_{00}\) must be positive definite, otherwise the trace minimisation is unbounded along \(\ker B\) and a ValueError is raised. At most \(\operatorname{rank}(B)\) finite eigenpairs exist; asking for more warns and returns fewer. The numerical rank decides output shapes, so this path cannot run under jax.jit (tag \(B\) positive definite for a traceable path).

Krylov subspaces are shift-invariant, so the Lanczos path needs no spectral shift for which="smallest": it runs an oversampled Krylov space of dimension min(N, max(2 * rank + 1, rank + 20)) and keeps the rank extreme Ritz pairs at the requested end.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator	Symmetric (N, N) operator.	required
B	AbstractLinearOperator	Symmetric positive semidefinite (N, N) operator.	required
rank	int | None	Number of eigenpairs \(K\). None returns all finite ones.	None
which	Literal['smallest', 'largest']	"smallest" or "largest" eigenvalues.	'smallest'
rcond	float | None	Relative threshold on \(B\)'s eigenvalues below which a direction counts as \(\ker B\) (singular path only). None means N * eps of \(B\)'s dtype.	None
key	PRNGKeyArray | None	PRNG key for the Lanczos start vector (diagonal path with rank= only). None means jax.random.PRNGKey(0).	None

Returns:

Type	Description
Float[Array, ' K']	(eigenvalues, eigenvectors) of shapes (K,) and (N, K),
Float[Array, 'N K']	eigenvalues in ascending order, eigenvectors \(B\)-orthonormal.

Raises:

Type	Description
ValueError	If which is invalid, rank < 1, \(B\) is zero, or \(A_{00}\) is not positive definite.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> # Laplacian eigenmaps on a path graph: L y = λ D y.
>>> # Its eigenvalues are 1 - cos(πk/4); rank= runs Lanczos.
>>> W = jnp.diag(jnp.ones(4), 1) + jnp.diag(jnp.ones(4), -1)
>>> degree = W @ jnp.ones(5)
>>> L = lx.MatrixLinearOperator(jnp.diag(degree) - W, lx.symmetric_tag)
>>> D = lx.DiagonalLinearOperator(degree)
>>> lam, Y = gaussx.eigh_generalized(L, D, rank=2, which="smallest")
>>> bool(jnp.allclose(lam, 1 - jnp.cos(jnp.pi * jnp.arange(2) / 4), atol=1e-5))
True
>>> # A singular B: one finite eigenpair fewer than N.
>>> B = lx.MatrixLinearOperator(jnp.diag(jnp.array([1.0, 2.0, 0.0])))
>>> A = lx.MatrixLinearOperator(jnp.eye(3) + 0.5 * jnp.ones((3, 3)))
>>> lam, V = gaussx.eigh_generalized(A, B)
>>> lam.shape
(2,)

Source code in src/gaussx/_primitives/_eig.py

def eigh_generalized(
    A: lx.AbstractLinearOperator,
    B: lx.AbstractLinearOperator,
    *,
    rank: int | None = None,
    which: Literal["smallest", "largest"] = "smallest",
    rcond: float | None = None,
    key: PRNGKeyArray | None = None,
) -> tuple[Float[Array, " K"], Float[Array, "N K"]]:
    r"""Solve $A v = \lambda B v$ for symmetric $A$ and symmetric PSD $B$.

    The eigenvectors are $B$-orthonormal, $V^\top B V = I$, so the
    smallest $K$ of them minimise $\operatorname{tr}(Y^\top A Y)$ subject
    to $Y^\top B Y = I$ (graph embeddings, LPP, manifold alignment).
    Dispatch is on the structure of $B$:

    - `lineax.DiagonalLinearOperator` with positive entries: scale,
      $S = B^{-1/2} A B^{-1/2}$, $v = B^{-1/2} u$. This is matrix-free in
      $A$: with `rank=` it runs Lanczos on $x \mapsto B^{-1/2} A B^{-1/2} x$
      and never materialises $A$. A diagonal with zero entries (only
      checked outside `jit`) is sent to the singular path below.
    - Positive definite, i.e. tagged `lineax.positive_semidefinite_tag`
      (lineax's tag for Cholesky-factorisable operators): Cholesky
      whitening $C^{-1} A C^{-\top}$ with $B = C C^\top$.
    - Otherwise $B$ is treated as PSD and possibly singular. $B$ is
      eigendecomposed, $B = U_+ S_+ U_+^\top$, and directions with
      eigenvalue $\le$ `rcond` $\cdot \max$ span $\ker B$. Writing
      $A_{\cdot\cdot}$ for the blocks of $U^\top A U$, the $\ker B$ rows give
      $v_0 = -A_{00}^{-1} A_{0+} v_+$, and the finite eigenpairs solve the
      Schur-complement pencil
      $(A_{++} - A_{+0} A_{00}^{-1} A_{0+}) v_+ = \lambda S_+ v_+$.
      $A_{00}$ must be positive definite, otherwise the trace minimisation
      is unbounded along $\ker B$ and a `ValueError` is raised. At most
      $\operatorname{rank}(B)$ finite eigenpairs exist; asking for more
      warns and returns fewer. The numerical rank decides output shapes,
      so this path cannot run under `jax.jit` (tag $B$ positive definite
      for a traceable path).

    Krylov subspaces are shift-invariant, so the Lanczos path needs no
    spectral shift for `which="smallest"`: it runs an oversampled Krylov
    space of dimension `min(N, max(2 * rank + 1, rank + 20))` and keeps
    the `rank` extreme Ritz pairs at the requested end.

    Args:
        A: Symmetric `(N, N)` operator.
        B: Symmetric positive semidefinite `(N, N)` operator.
        rank: Number of eigenpairs $K$. `None` returns all finite ones.
        which: `"smallest"` or `"largest"` eigenvalues.
        rcond: Relative threshold on $B$'s eigenvalues below which a
            direction counts as $\ker B$ (singular path only). `None`
            means `N * eps` of $B$'s dtype.
        key: PRNG key for the Lanczos start vector (diagonal path with
            `rank=` only). `None` means `jax.random.PRNGKey(0)`.

    Returns:
        `(eigenvalues, eigenvectors)` of shapes `(K,)` and `(N, K)`,
        eigenvalues in ascending order, eigenvectors $B$-orthonormal.

    Raises:
        ValueError: If `which` is invalid, `rank < 1`, $B$ is zero, or
            $A_{00}$ is not positive definite.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> # Laplacian eigenmaps on a path graph: L y = λ D y.
        >>> # Its eigenvalues are 1 - cos(πk/4); rank= runs Lanczos.
        >>> W = jnp.diag(jnp.ones(4), 1) + jnp.diag(jnp.ones(4), -1)
        >>> degree = W @ jnp.ones(5)
        >>> L = lx.MatrixLinearOperator(jnp.diag(degree) - W, lx.symmetric_tag)
        >>> D = lx.DiagonalLinearOperator(degree)
        >>> lam, Y = gaussx.eigh_generalized(L, D, rank=2, which="smallest")
        >>> bool(jnp.allclose(lam, 1 - jnp.cos(jnp.pi * jnp.arange(2) / 4), atol=1e-5))
        True
        >>> # A singular B: one finite eigenpair fewer than N.
        >>> B = lx.MatrixLinearOperator(jnp.diag(jnp.array([1.0, 2.0, 0.0])))
        >>> A = lx.MatrixLinearOperator(jnp.eye(3) + 0.5 * jnp.ones((3, 3)))
        >>> lam, V = gaussx.eigh_generalized(A, B)
        >>> lam.shape
        (2,)
    """
    if which not in ("smallest", "largest"):
        msg = f"which must be 'smallest' or 'largest', got {which!r}."
        raise ValueError(msg)
    if rank is not None and rank < 1:
        msg = f"rank must be a positive integer, got {rank}."
        raise ValueError(msg)
    if A.in_size() != B.in_size() or A.in_size() != A.out_size():
        msg = (
            f"A and B must be square and of the same size, got A "
            f"{A.out_size()}x{A.in_size()} and B {B.out_size()}x{B.in_size()}."
        )
        raise ValueError(msg)

    if isinstance(B, lx.DiagonalLinearOperator):
        d = lx.diagonal(B)
        if isinstance(d, Tracer) or bool(jnp.min(d) > _null_tol(d, rcond, d.shape[0])):
            return _eigh_generalized_diagonal(A, d, rank, which, key)
    elif lx.is_positive_semidefinite(B):
        return _eigh_generalized_cholesky(A, B, rank, which)
    return _eigh_generalized_singular(A, B, rank, which, rcond)

Matrix-free square-root products

\(A^{\pm 1/2}b\) for an operator too large to factorise, via the Hale–Higham–Trefethen contour-integral quadrature: a weighted sum of ~15 shifted solves \((A + \sigma_j I)^{-1}b\), each routed back through solve. Accuracy depends on the condition number only logarithmically, and estimate_spectral_bounds supplies the contour parameters when they are not known ahead of time.

Structured linear algebra and Gaussian primitives for JAX.

sqrt_matmul(operator: lx.AbstractLinearOperator, rhs: Float[Array, 'N C'], *, num_quadrature: int = _DEFAULT_NUM_QUADRATURE, max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER, spectral_bounds: tuple[float, float] | None = None, solver: lx.AbstractLinearSolver | None = None) -> Float[Array, 'N C']

Compute \(A^{1/2} b\) via contour-integral quadrature.

Shares the quadrature of sqrt_inv_matmul through the identity \(A^{1/2} b = A \, (A^{-1/2} b)\), which costs one extra matvec per column.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Square symmetric positive-definite operator \(A\).	required
rhs	Float[Array, 'N C']	Right-hand side of shape (N, C).	required
num_quadrature	int	Number of contour-quadrature nodes \(J\).	_DEFAULT_NUM_QUADRATURE
max_lanczos_iter	int	Lanczos iterations used to estimate the spectral bounds when spectral_bounds is not given.	_DEFAULT_MAX_LANCZOS_ITER
spectral_bounds	tuple[float, float] | None	Optional (lam_min, lam_max) bracketing the spectrum.	None
solver	AbstractLinearSolver | None	Optional lineax solver used for the shifted solves.	None

Returns:

Type	Description
Float[Array, 'N C']	The array \(A^{1/2} b\) of shape (N, C).

Source code in src/gaussx/_primitives/_sqrt_matmul.py

def sqrt_matmul(
    operator: lx.AbstractLinearOperator,
    rhs: Float[Array, "N C"],
    *,
    num_quadrature: int = _DEFAULT_NUM_QUADRATURE,
    max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER,
    spectral_bounds: tuple[float, float] | None = None,
    solver: lx.AbstractLinearSolver | None = None,
) -> Float[Array, "N C"]:
    r"""Compute $A^{1/2} b$ via contour-integral quadrature.

    Shares the quadrature of `sqrt_inv_matmul` through the identity
    $A^{1/2} b = A \, (A^{-1/2} b)$, which costs one extra matvec per column.

    Args:
        operator: Square symmetric positive-definite operator $A$.
        rhs: Right-hand side of shape ``(N, C)``.
        num_quadrature: Number of contour-quadrature nodes $J$.
        max_lanczos_iter: Lanczos iterations used to estimate the spectral
            bounds when ``spectral_bounds`` is not given.
        spectral_bounds: Optional ``(lam_min, lam_max)`` bracketing the
            spectrum.
        solver: Optional lineax solver used for the shifted solves.

    Returns:
        The array $A^{1/2} b$ of shape ``(N, C)``.
    """
    whitened = sqrt_inv_matmul(
        operator,
        rhs,
        num_quadrature=num_quadrature,
        max_lanczos_iter=max_lanczos_iter,
        spectral_bounds=spectral_bounds,
        solver=solver,
    )
    return jax.vmap(operator.mv, in_axes=1, out_axes=1)(whitened)

sqrt_inv_matmul(operator: lx.AbstractLinearOperator, rhs: Float[Array, 'N C'], *, num_quadrature: int = _DEFAULT_NUM_QUADRATURE, max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER, spectral_bounds: tuple[float, float] | None = None, solver: lx.AbstractLinearSolver | None = None) -> Float[Array, 'N C']

Compute \(A^{-1/2} b\) via contour-integral quadrature.

Uses the integral representation

\[A^{-1/2} = \frac{2}{\pi} \int_0^{\infty} (A + t^2 I)^{-1} \, dt\]

discretised by the Hale-Higham-Trefethen elliptic-integral rule, which turns it into \(J\) shifted solves

\[A^{-1/2} b \approx \sum_{j=1}^{J} w_j \, (A + \sigma_j I)^{-1} b, \qquad \sigma_j \ge 0 .\]

Each shifted operator is positive definite, so every solve goes through the ordinary gaussx.solve dispatch. Accuracy improves geometrically in \(J\) with a rate that depends only logarithmically on the condition number \(\kappa = \lambda_{\max} / \lambda_{\min}\): \(J = 15\) reaches 1e-13 at \(\kappa = 10^3\) and 1e-7 at \(\kappa = 10^6\).

The quadrature nodes and weights are treated as constants by jax.grad — they parameterise the contour, not the function value, and the rule is designed to be insensitive to them.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Square symmetric positive-definite operator \(A\).	required
rhs	Float[Array, 'N C']	Right-hand side of shape (N, C).	required
num_quadrature	int	Number of contour-quadrature nodes \(J\).	_DEFAULT_NUM_QUADRATURE
max_lanczos_iter	int	Lanczos iterations used to estimate the spectral bounds when spectral_bounds is not given.	_DEFAULT_MAX_LANCZOS_ITER
spectral_bounds	tuple[float, float] | None	Optional (lam_min, lam_max) bracketing the spectrum. Supplying known bounds skips the Lanczos estimate.	None
solver	AbstractLinearSolver | None	Optional lineax solver used for the shifted solves. Pass lineax.CG(...) for large matrix-free operators.	None

Returns:

Type	Description
Float[Array, 'N C']	The array \(A^{-1/2} b\) of shape (N, C).

Source code in src/gaussx/_primitives/_sqrt_matmul.py

def sqrt_inv_matmul(
    operator: lx.AbstractLinearOperator,
    rhs: Float[Array, "N C"],
    *,
    num_quadrature: int = _DEFAULT_NUM_QUADRATURE,
    max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER,
    spectral_bounds: tuple[float, float] | None = None,
    solver: lx.AbstractLinearSolver | None = None,
) -> Float[Array, "N C"]:
    r"""Compute $A^{-1/2} b$ via contour-integral quadrature.

    Uses the integral representation

    $$A^{-1/2} = \frac{2}{\pi} \int_0^{\infty} (A + t^2 I)^{-1} \, dt$$

    discretised by the Hale-Higham-Trefethen elliptic-integral rule, which
    turns it into $J$ shifted solves

    $$A^{-1/2} b \approx \sum_{j=1}^{J} w_j \, (A + \sigma_j I)^{-1} b,
    \qquad \sigma_j \ge 0 .$$

    Each shifted operator is positive definite, so every solve goes through
    the ordinary `gaussx.solve` dispatch. Accuracy improves geometrically in
    $J$ with a rate that depends only *logarithmically* on the condition
    number $\kappa = \lambda_{\max} / \lambda_{\min}$: $J = 15$ reaches
    ``1e-13`` at $\kappa = 10^3$ and ``1e-7`` at $\kappa = 10^6$.

    The quadrature nodes and weights are treated as constants by
    `jax.grad` — they parameterise the contour, not the function value, and
    the rule is designed to be insensitive to them.

    Args:
        operator: Square symmetric positive-definite operator $A$.
        rhs: Right-hand side of shape ``(N, C)``.
        num_quadrature: Number of contour-quadrature nodes $J$.
        max_lanczos_iter: Lanczos iterations used to estimate the spectral
            bounds when ``spectral_bounds`` is not given.
        spectral_bounds: Optional ``(lam_min, lam_max)`` bracketing the
            spectrum. Supplying known bounds skips the Lanczos estimate.
        solver: Optional lineax solver used for the shifted solves. Pass
            ``lineax.CG(...)`` for large matrix-free operators.

    Returns:
        The array $A^{-1/2} b$ of shape ``(N, C)``.
    """
    shifts, weights = _contour_rule(
        operator,
        num_quadrature=num_quadrature,
        max_lanczos_iter=max_lanczos_iter,
        spectral_bounds=spectral_bounds,
        rhs=rhs,
    )
    solves = jax.vmap(lambda shift: _shifted_solve(operator, shift, rhs, solver))(
        shifts
    )
    return einsum(weights, solves, "j, j n c -> n c")

estimate_spectral_bounds(operator: lx.AbstractLinearOperator, *, max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER, safety: float = _DEFAULT_SAFETY, key: jax.Array | None = None) -> tuple[Float[Array, ''], Float[Array, '']]

Estimate \((\lambda_{\min}, \lambda_{\max})\) of a symmetric operator.

Structured operators (diagonal, identity, Kronecker, block-diagonal, Kronecker-sum) return their exact spectra, including when wrapped in lineax.TaggedLinearOperator or scaled by a scalar -- gaussx.eigvals only recognises the bare classes. Everything else runs a partial Lanczos decomposition with full reorthogonalisation, which stops at the first breakdown: an operator such as cI + UUᵀ exhausts its Krylov space after rank(U) + 1 steps, and continuing past that point only normalises rounding noise into spurious Ritz values.

Ritz values interlace the true spectrum, so a partial Lanczos run brackets it from the inside — and the smallest eigenvalue is the slowest one to converge, so the gap can be a factor of several. The estimate is therefore widened by safety in each direction whenever it came from a partial run. Widening is close to free for the contour quadrature that consumes these bounds: its node count grows only with \(\log \kappa\).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A square symmetric positive-definite operator.	required
max_lanczos_iter	int	Number of Lanczos iterations (clamped to the operator size).	_DEFAULT_MAX_LANCZOS_ITER
safety	float	Factor by which a partial-Lanczos bracket is widened in each direction. Ignored when the full spectrum is available.	_DEFAULT_SAFETY
key	Array | None	PRNG key for the Lanczos start vector. Defaults to jax.random.PRNGKey(0).	None

Returns:

Type	Description
Float[Array, '']	Tuple (lam_min, lam_max) of scalar eigenvalue estimates, floored
Float[Array, '']	at a positive value so downstream square roots stay finite.

Source code in src/gaussx/_primitives/_sqrt_matmul.py

def estimate_spectral_bounds(
    operator: lx.AbstractLinearOperator,
    *,
    max_lanczos_iter: int = _DEFAULT_MAX_LANCZOS_ITER,
    safety: float = _DEFAULT_SAFETY,
    key: jax.Array | None = None,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Estimate $(\lambda_{\min}, \lambda_{\max})$ of a symmetric operator.

    Structured operators (diagonal, identity, Kronecker, block-diagonal,
    Kronecker-sum) return their exact spectra, including when wrapped in
    `lineax.TaggedLinearOperator` or scaled by a scalar -- `gaussx.eigvals`
    only recognises the bare classes. Everything else runs a partial Lanczos
    decomposition with full reorthogonalisation, which stops at the first
    breakdown: an operator such as ``cI + UUᵀ`` exhausts its Krylov space
    after ``rank(U) + 1`` steps, and continuing past that point only
    normalises rounding noise into spurious Ritz values.

    Ritz values interlace the true spectrum, so a partial Lanczos run brackets
    it from the *inside* — and the smallest eigenvalue is the slowest one to
    converge, so the gap can be a factor of several. The estimate is therefore
    widened by ``safety`` in each direction whenever it came from a partial
    run. Widening is close to free for the contour quadrature that consumes
    these bounds: its node count grows only with $\log \kappa$.

    Args:
        operator: A square symmetric positive-definite operator.
        max_lanczos_iter: Number of Lanczos iterations (clamped to the
            operator size).
        safety: Factor by which a partial-Lanczos bracket is widened in each
            direction. Ignored when the full spectrum is available.
        key: PRNG key for the Lanczos start vector. Defaults to
            ``jax.random.PRNGKey(0)``.

    Returns:
        Tuple ``(lam_min, lam_max)`` of scalar eigenvalue estimates, floored
        at a positive value so downstream square roots stay finite.
    """
    if operator.in_size() != operator.out_size():
        raise ValueError("spectral bounds require a square operator")
    if safety < 1.0:
        raise ValueError("safety must be at least 1")
    lam_min, lam_max = _spectral_bracket(operator, max_lanczos_iter, safety, key)
    floor = jnp.finfo(jnp.result_type(lam_min)).tiny
    lam_min = jnp.maximum(lam_min, floor)
    return lam_min, jnp.maximum(lam_max, lam_min)

Joint inverse-quadratic and log-determinant

inv_quad_logdet returns \(\mathrm{tr}(R^\top A^{-1}R)\) and \(\log|A|\) from a single modified-batched-CG pass — the two halves of a Gaussian log-density at roughly the matvec budget of one. Supplying logdet_preconditioner= \(P \approx A\) switches on the Artemev et al. variance reduction, estimating only the near-identity residual \(\log(P^{-1}A)\) stochastically. \(P\) approximates \(A\) itself, unlike the preconditioner= (\(M^{-1} \approx A^{-1}\)) of the CG solvers. The old keyword preconditioner= is a deprecated alias.

Structured linear algebra and Gaussian primitives for JAX.

inv_quad_logdet(operator: lx.AbstractLinearOperator, rhs: Float[Array, 'N C'], *, strategy: AbstractSolverStrategy | None = None, reduce_inv_quad: bool = True, logdet_preconditioner: lx.AbstractLinearOperator | None = None, preconditioner: lx.AbstractLinearOperator | None = UNSET) -> tuple[Float[Array, ' *C'], Float[Array, '']]

Joint inverse-quadratic and log-determinant via shared BBMM CG.

Returns the pair

\[ \mathrm{inv\_quad}(A, R) = \mathrm{tr}(R^{\top} A^{-1} R) = \sum_{c} r_c^{\top} A^{-1} r_c, \qquad \mathrm{logdet}(A) = \log |A| , \]

which together with the \(-\tfrac{N}{2}\log 2\pi\) constant make up the Gaussian log-density. For a structured operator with an exact log-determinant the default is the exact structural pair. Otherwise, with gaussx.BBMMSolver, both come out of one modified-batched-CG pass over [rhs | probes], so a marginal-likelihood step costs roughly half the matvecs of a separate gaussx.solve plus gaussx.logdet.

Passing a logdet_preconditioner \(P \approx A\) applies the variance reduction of Artemev et al. (2021): the same CG pass then produces the Lanczos tridiagonal of \(P^{-1} A\), so the stochastic part of

\[\log |A| = \log |P| + \mathrm{tr}\bigl[\log (P^{-1} A)\bigr]\]

only has to estimate a trace of a near-identity operator.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Square symmetric positive-definite operator \(A\).	required
rhs	Float[Array, 'N C']	Right-hand side \(R\) of shape (N, C).	required
strategy	AbstractSolverStrategy | None	Solver strategy. None (the default) takes the exact structural path (gaussx.solve per column plus gaussx.logdet) when the operator has one (diagonal, Kronecker, BlockDiag, LowRankUpdate, BlockTriDiag, KroneckerSum, ... and scalar or tagged wrappers of them), and gaussx.BBMMSolver otherwise. gaussx.BBMMSolver takes the shared-work path; any other strategy falls back to its own solve per column plus its own logdet.	None
reduce_inv_quad	bool	Whether to sum the per-column quadratic forms into the trace. When False the first return value is the (C,) vector of \(r_c^{\top} A^{-1} r_c\).	True
logdet_preconditioner	AbstractLinearOperator | None	Optional operator \(P \approx A\) with a cheap log-determinant, used for variance reduction as above. This is the approximation to \(A\) itself, not an approximate inverse like the preconditioner= of gaussx.CGSolver and gaussx.linear_solve. Ignored when the strategy's log-determinant is already exact (e.g. gaussx.DenseSolver), since there is no variance to reduce.	None
preconditioner	AbstractLinearOperator | None	Deprecated alias of logdet_preconditioner (gh-405), renamed because everywhere else in gaussx preconditioner= means \(M^{-1} \approx A^{-1}\).	UNSET

Returns:

Type	Description
Float[Array, ' *C']	Tuple (inv_quad, logdet). inv_quad is a scalar when
Float[Array, '']	reduce_inv_quad is True and a (C,) vector otherwise.

Source code in src/gaussx/_primitives/_inv_quad_logdet.py

def inv_quad_logdet(
    operator: lx.AbstractLinearOperator,
    rhs: Float[Array, "N C"],
    *,
    strategy: AbstractSolverStrategy | None = None,
    reduce_inv_quad: bool = True,
    logdet_preconditioner: lx.AbstractLinearOperator | None = None,
    preconditioner: lx.AbstractLinearOperator | None = UNSET,
) -> tuple[Float[Array, " *C"], Float[Array, ""]]:
    r"""Joint inverse-quadratic and log-determinant via shared BBMM CG.

    Returns the pair

    $$
    \mathrm{inv\_quad}(A, R) = \mathrm{tr}(R^{\top} A^{-1} R)
      = \sum_{c} r_c^{\top} A^{-1} r_c,
    \qquad
    \mathrm{logdet}(A) = \log |A| ,
    $$

    which together with the $-\tfrac{N}{2}\log 2\pi$ constant make up the
    Gaussian log-density. For a structured operator with an exact
    log-determinant the default is the exact structural pair. Otherwise, with
    `gaussx.BBMMSolver`, both come out of one modified-batched-CG pass over
    ``[rhs | probes]``, so a marginal-likelihood step costs roughly half the
    matvecs of a separate `gaussx.solve` plus `gaussx.logdet`.

    Passing a ``logdet_preconditioner`` $P \approx A$ applies the variance reduction
    of Artemev et al. (2021): the same CG pass then produces the Lanczos
    tridiagonal of $P^{-1} A$, so the stochastic part of

    $$\log |A| = \log |P| + \mathrm{tr}\bigl[\log (P^{-1} A)\bigr]$$

    only has to estimate a trace of a near-identity operator.

    Args:
        operator: Square symmetric positive-definite operator $A$.
        rhs: Right-hand side $R$ of shape ``(N, C)``.
        strategy: Solver strategy. ``None`` (the default) takes the exact
            structural path (`gaussx.solve` per column plus `gaussx.logdet`)
            when the operator has one (diagonal, `Kronecker`, `BlockDiag`,
            `LowRankUpdate`, `BlockTriDiag`, `KroneckerSum`, ... and scalar
            or tagged wrappers of them), and `gaussx.BBMMSolver` otherwise.
            `gaussx.BBMMSolver` takes the shared-work path; any other
            strategy falls back to its own `solve` per column plus its own
            `logdet`.
        reduce_inv_quad: Whether to sum the per-column quadratic forms into
            the trace. When ``False`` the first return value is the ``(C,)``
            vector of $r_c^{\top} A^{-1} r_c$.
        logdet_preconditioner: Optional operator $P \approx A$ with a cheap
            log-determinant, used for variance reduction as above. This is the
            approximation to $A$ itself, **not** an approximate inverse like
            the ``preconditioner=`` of `gaussx.CGSolver` and
            `gaussx.linear_solve`. Ignored when the strategy's
            log-determinant is already exact (e.g. `gaussx.DenseSolver`),
            since there is no variance to reduce.
        preconditioner: Deprecated alias of ``logdet_preconditioner``
            (gh-405), renamed because everywhere else in gaussx
            ``preconditioner=`` means $M^{-1} \approx A^{-1}$.

    Returns:
        Tuple ``(inv_quad, logdet)``. ``inv_quad`` is a scalar when
        ``reduce_inv_quad`` is ``True`` and a ``(C,)`` vector otherwise.
    """
    preconditioner = default(
        renamed(
            "inv_quad_logdet",
            "preconditioner",
            preconditioner,
            "logdet_preconditioner",
            logdet_preconditioner if logdet_preconditioner is not None else UNSET,
        ),
        None,
    )
    if operator.in_size() != operator.out_size():
        raise ValueError("inv_quad_logdet requires a square operator")
    if jnp.ndim(rhs) != 2:
        raise ValueError(
            f"rhs must have shape (N, C); got array with {jnp.ndim(rhs)} axes"
        )
    if rhs.shape[0] != operator.in_size():
        raise ValueError(
            f"rhs has {rhs.shape[0]} rows but operator has size {operator.in_size()}"
        )
    if strategy is None and (
        as_diagonalized(operator) is not None or _has_structural_logdet(operator)
    ):
        # Exact and cheap: a structural solve per column plus the structural
        # log-determinant, instead of BBMM's stochastic one (gh-340).
        columns = jnp.sum(
            rhs * jax.vmap(ft.partial(_solve, operator), 1, 1)(rhs), axis=0
        )
        inv_quad = jnp.sum(columns) if reduce_inv_quad else columns
        return inv_quad, _logdet(operator)
    if strategy is None:
        strategy = BBMMSolver()

    if isinstance(strategy, BBMMSolver):
        columns, logdet_value = _shared_work((operator, rhs), strategy, preconditioner)
    else:
        columns, logdet_value = _separate_work(operator, rhs, strategy, preconditioner)

    inv_quad = jnp.sum(columns) if reduce_inv_quad else columns
    return inv_quad, logdet_value

Root decompositions

Tall-factor approximations \(RR^\top \approx A\) (and \(R^- (R^-)^\top \approx A^{-1}\)) via Cholesky, pivoted Cholesky, Lanczos, or truncated SVD — the building block for low-rank posterior sampling and BBMM-style solvers.

Structured linear algebra and Gaussian primitives for JAX.

RootDecomposition

Bases: Module

Tall factor R with R Rᵀ ≈ A semantics.

Attributes:

Name	Type	Description
root	Float[Array, 'N k']	Tall factor with shape (N, k).

Source code in src/gaussx/_primitives/_root.py

class RootDecomposition(eqx.Module):
    r"""Tall factor ``R`` with ``R Rᵀ ≈ A`` semantics.

    Attributes:
        root: Tall factor with shape ``(N, k)``.
    """

    root: Float[Array, "N k"]

    @property
    def rank(self) -> int:
        """Number of retained root directions."""
        return self.root.shape[1]

    def matmul(self, x: Float[Array, "*B k"]) -> Float[Array, "*B N"]:
        """Right-multiply by the root factor: ``x Rᵀ``."""
        return x @ self.root.T

rank: int property

Number of retained root directions.

matmul(x: Float[Array, '*B k']) -> Float[Array, '*B N']

Right-multiply by the root factor: x Rᵀ.

Source code in src/gaussx/_primitives/_root.py

def matmul(self, x: Float[Array, "*B k"]) -> Float[Array, "*B N"]:
    """Right-multiply by the root factor: ``x Rᵀ``."""
    return x @ self.root.T

root_decomposition(operator: lx.AbstractLinearOperator, rank: int = 50, method: RootMethod = 'lanczos', key: jax.Array | None = None) -> RootDecomposition

Compute a tall factor R such that R Rᵀ ≈ A.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Square symmetric positive-definite operator A.	required
rank	int	Number of retained directions. Ignored by "cholesky", which returns the exact full-rank factor.	50
method	RootMethod	Decomposition method: "lanczos", "cholesky", "pivoted_cholesky", or "svd".	'lanczos'
key	Array | None	PRNG key for random-start methods.	None

Returns:

Type	Description
RootDecomposition	A RootDecomposition with root shape (N, k).

Source code in src/gaussx/_primitives/_root.py

def root_decomposition(
    operator: lx.AbstractLinearOperator,
    rank: int = 50,
    method: RootMethod = "lanczos",
    key: jax.Array | None = None,
) -> RootDecomposition:
    r"""Compute a tall factor ``R`` such that ``R Rᵀ ≈ A``.

    Args:
        operator: Square symmetric positive-definite operator ``A``.
        rank: Number of retained directions. Ignored by ``"cholesky"``,
            which returns the exact full-rank factor.
        method: Decomposition method: ``"lanczos"``, ``"cholesky"``,
            ``"pivoted_cholesky"``, or ``"svd"``.
        key: PRNG key for random-start methods.

    Returns:
        A `RootDecomposition` with root shape ``(N, k)``.
    """
    _n, rank = _validate_square_rank(operator, rank, method=method)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return RootDecomposition(_diagonal_root(operator, rank, method, inverse=False))
    if method == "cholesky":
        return RootDecomposition(_cholesky_root(operator))
    if method == "lanczos":
        return RootDecomposition(_eig_root(operator, rank, key, inverse=False))
    if method == "svd":
        return RootDecomposition(_svd_root(operator, rank, key, inverse=False))
    if method == "pivoted_cholesky":
        mat = _symmetric_operator_matrix(operator)
        return RootDecomposition(_pivoted_cholesky_root(mat, rank))
    raise ValueError(f"Unknown root decomposition method: {method!r}")

root_inv_decomposition(operator: lx.AbstractLinearOperator, rank: int = 50, method: RootMethod = 'lanczos', key: jax.Array | None = None) -> RootDecomposition

Compute a tall factor R⁻ such that R⁻ (R⁻)ᵀ ≈ A⁻¹.

method chooses the approximation target, not only the algorithm (gh-413):

• "lanczos" / "svd": R⁻ (R⁻)ᵀ is A⁻¹ restricted to the dominant rank-dimensional eigenspace of A (its largest eigenvalues, inverted) — GPyTorch's convention, used by gaussx.love_cache.
• "pivoted_cholesky": R⁻ (R⁻)ᵀ is a rank-rank approximation of A⁻¹ itself, which favours the dominant subspace of A⁻¹ (the smallest eigenvalues of A); on a diagonal A it keeps the rank largest 1/d, the Frobenius-optimal rank-rank approximation of A⁻¹.
• "cholesky": exact (rank is ignored).

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Square symmetric positive-definite operator A.	required
rank	int	Number of retained directions. Ignored by "cholesky", which returns the exact full-rank inverse-root factor.	50
method	RootMethod	Decomposition method: "lanczos", "cholesky", "pivoted_cholesky", or "svd".	'lanczos'
key	Array | None	PRNG key for random-start methods.	None

Returns:

Type	Description
RootDecomposition	A RootDecomposition with inverse-root shape (N, k).

Source code in src/gaussx/_primitives/_root.py

def root_inv_decomposition(
    operator: lx.AbstractLinearOperator,
    rank: int = 50,
    method: RootMethod = "lanczos",
    key: jax.Array | None = None,
) -> RootDecomposition:
    r"""Compute a tall factor ``R⁻`` such that ``R⁻ (R⁻)ᵀ ≈ A⁻¹``.

    ``method`` chooses the approximation target, not only the algorithm
    (gh-413):

    - ``"lanczos"`` / ``"svd"``: ``R⁻ (R⁻)ᵀ`` is ``A⁻¹`` restricted to the
      dominant ``rank``-dimensional eigenspace of ``A`` (its *largest*
      eigenvalues, inverted) — GPyTorch's convention, used by
      `gaussx.love_cache`.
    - ``"pivoted_cholesky"``: ``R⁻ (R⁻)ᵀ`` is a rank-``rank`` approximation
      of ``A⁻¹`` itself, which favours the dominant subspace of ``A⁻¹`` (the
      *smallest* eigenvalues of ``A``); on a diagonal ``A`` it keeps the
      ``rank`` largest ``1/d``, the Frobenius-optimal rank-``rank``
      approximation of ``A⁻¹``.
    - ``"cholesky"``: exact (``rank`` is ignored).

    Args:
        operator: Square symmetric positive-definite operator ``A``.
        rank: Number of retained directions. Ignored by ``"cholesky"``,
            which returns the exact full-rank inverse-root factor.
        method: Decomposition method: ``"lanczos"``, ``"cholesky"``,
            ``"pivoted_cholesky"``, or ``"svd"``.
        key: PRNG key for random-start methods.

    Returns:
        A `RootDecomposition` with inverse-root shape ``(N, k)``.
    """
    n, rank = _validate_square_rank(operator, rank, method=method)
    if isinstance(operator, lx.DiagonalLinearOperator):
        return RootDecomposition(_diagonal_root(operator, rank, method, inverse=True))
    if method == "cholesky":
        L = _cholesky_root(operator)
        identity = jnp.eye(n, dtype=L.dtype)
        L_inv = jax.scipy.linalg.solve_triangular(L, identity, lower=True)
        return RootDecomposition(L_inv.T)
    if method == "lanczos":
        return RootDecomposition(_eig_root(operator, rank, key, inverse=True))
    if method == "svd":
        return RootDecomposition(_svd_root(operator, rank, key, inverse=True))
    if method == "pivoted_cholesky":
        # Dense fallback: forms A^{-1} explicitly via the lazy inverse
        # then runs pivoted Cholesky on it — O(N^3) time, O(N^2) memory.
        # Prefer ``method="lanczos"`` for large operators where structural
        # solves/matvecs are available.
        inv_mat = _symmetric_operator_matrix(inv(operator))
        return RootDecomposition(_pivoted_cholesky_root(inv_mat, rank))
    raise ValueError(f"Unknown inverse-root decomposition method: {method!r}")

Quantiles

mixture_quantile inverts a CDF \(F\) at levels \(q\) by solving \(F(x) - q = 0\) on a bracket \([\ell, u]\) — the standard remedy for a mixture \(F_{\text{mix}} = \tfrac{1}{E}\sum_e F_e\), which has no closed-form inverse. It runs optimistix's root_find with Chandrupatla, a bracketing solver that mixes inverse quadratic interpolation with bisection (Chandrupatla, 1997), so it needs no derivatives, and differentiates through the implicit adjoint \(\partial x_q / \partial\theta = -\partial_\theta F / p(x_q)\). mixture_quantile_gaussian_approx is the cheap moment-matched Gaussian alternative for Gaussian ensembles (Lakshminarayanan et al., 2017).

import jax.numpy as jnp
from jax.scipy.stats import norm
import gaussx

means = jnp.array([-1.0, 1.0])  # two-component mixture
cdf = lambda x: 0.5 * (norm.cdf(x - means[0]) + norm.cdf(x - means[1]))
gaussx.mixture_quantile(cdf, jnp.array([0.025, 0.5, 0.975]), -10.0, 10.0)

Structured linear algebra and Gaussian primitives for JAX.

Chandrupatla

Bases: AbstractRootFinder

Chandrupatla's hybrid inverse-quadratic / bisection root finder.

A bracketing method for \(f(x) = 0\) that keeps the robustness of bisection and the superlinear convergence of inverse quadratic interpolation (Chandrupatla, 1997). The bracket \([a, b]\) with \(f(a) f(b) \le 0\) and the previous point \(c\) are updated each step with the new point \(x_t = (1 - t)\,a + t\,b\). The next \(t\) is the inverse quadratic interpolant through \((a, b, c)\),

\[ t = \frac{f_a}{f_b - f_a}\frac{f_c}{f_b - f_c} + \frac{c - a}{b - a}\frac{f_a}{f_c - f_a}\frac{f_b}{f_c - f_b}, \]

accepted only when \(\xi = (a - b)/(c - b)\) and \(\Phi = (f_a - f_b)/(f_c - f_b)\) satisfy \(\Phi^2 < \xi\) and \((1 - \Phi)^2 < 1 - \xi\), i.e. when the interpolant is monotone on the bracket; otherwise \(t = 1/2\) (bisection). \(t\) is clipped to \([t_\ell, 1 - t_\ell]\) with \(t_\ell = \tfrac12\tau / |b - c|\) and \(\tau = \mathrm{rtol}\,|x_m| + \mathrm{atol}\), where \(x_m\) is whichever of \(a, b\) has the smaller \(|f|\). The solve stops when \(t_\ell > 1/2\), i.e. \(|b - c| < \tau\), which bounds the bracket \(|b - a| \le |b - c|\) and so puts \(x_m\) within \(\tau\) of the root; or when \(f(x_m) = 0\); or when the bracket has shrunk to adjacent floats and \(x_t\) can no longer move. It returns \(x_m\). (Scherer's formulation stops at \(2\tau\); halving keeps the requested tolerance.)

The solver is elementwise: y may be an array, fn must act on each entry independently (as a CDF evaluated pointwise does), and every entry carries its own bracket and stops on its own. The loop ends when all entries have converged. Pass tags=frozenset( {lineax.diagonal_tag}) to optimistix.root_find so the implicit adjoint uses the diagonal Jacobian.

Requires options=dict(lower=..., upper=...) (broadcastable to y), as optimistix.Bisection does. The bracket and every iterate keep y's dtype; only the function values take fn's. An entry whose bracket has no sign change (or a NaN endpoint value) is not iterated: it returns the endpoint with the smaller \(|f|\), so the caller decides whether that is an error. With has_aux=True the auxiliary output is that of fn at the returned root (one extra evaluation).

Pseudocode (per entry, Scherer 2010, §6.1):

b, a, c = lower, upper, upper;  t = 1/2
repeat:
    x_t = (1 − t) a + t b;  f_t = f(x_t)
    if sign f_t = sign f_a:  c ← a          else:  c ← b; b ← a
    a ← x_t
    x_m = argmin over a, b of |f|;  t_lim = (rtol |x_m| + atol) / (2 |b − c|)
    stop if f(x_m) = 0 or t_lim > 1/2 or x_t ∈ {a, b} before the update
    t = IQI(a, b, c) if Φ² < ξ and (1 − Φ)² < 1 − ξ else 1/2
    t = clip(t, t_lim, 1 − t_lim)

Attributes:

Name	Type	Description
rtol	float	Relative tolerance on the root.
atol	float	Absolute tolerance on the root.

References

Chandrupatla, T. R. (1997). A new hybrid quadratic/bisection algorithm for finding the zero of a nonlinear function without using derivatives. Advances in Engineering Software 28(3), 145-149.

Scherer, P. O. J. (2010). Computational Physics: Simulation of Classical and Quantum Systems. Springer, §6.1.

Examples:

>>> import jax.numpy as jnp
>>> import optimistix as optx
>>> import gaussx
>>> sol = optx.root_find(
...     lambda x, _: x**3 - 2.0,
...     gaussx.Chandrupatla(rtol=1e-10, atol=1e-12),
...     jnp.array(1.0),
...     options=dict(lower=0.0, upper=2.0),
... )
>>> round(float(sol.value), 6)
1.259921

Source code in src/gaussx/_primitives/_chandrupatla.py

class Chandrupatla(optx.AbstractRootFinder):
    r"""Chandrupatla's hybrid inverse-quadratic / bisection root finder.

    A bracketing method for $f(x) = 0$ that keeps the robustness of
    bisection and the superlinear convergence of inverse quadratic
    interpolation (Chandrupatla, 1997). The bracket $[a, b]$ with
    $f(a) f(b) \le 0$ and the previous point $c$ are updated each step
    with the new point $x_t = (1 - t)\,a + t\,b$. The next $t$ is the inverse
    quadratic interpolant through $(a, b, c)$,

    $$
    t = \frac{f_a}{f_b - f_a}\frac{f_c}{f_b - f_c}
      + \frac{c - a}{b - a}\frac{f_a}{f_c - f_a}\frac{f_b}{f_c - f_b},
    $$

    accepted only when $\xi = (a - b)/(c - b)$ and
    $\Phi = (f_a - f_b)/(f_c - f_b)$ satisfy $\Phi^2 < \xi$ and
    $(1 - \Phi)^2 < 1 - \xi$, i.e. when the interpolant is monotone on the
    bracket; otherwise $t = 1/2$ (bisection). $t$ is clipped to
    $[t_\ell, 1 - t_\ell]$ with $t_\ell = \tfrac12\tau / |b - c|$ and
    $\tau = \mathrm{rtol}\,|x_m| + \mathrm{atol}$, where $x_m$ is
    whichever of $a, b$ has the smaller $|f|$. The solve stops when
    $t_\ell > 1/2$, i.e. $|b - c| < \tau$, which bounds the bracket
    $|b - a| \le |b - c|$ and so puts $x_m$ within $\tau$ of the root; or
    when $f(x_m) = 0$; or when the bracket has shrunk to adjacent floats
    and $x_t$ can no longer move. It returns $x_m$. (Scherer's
    formulation stops at $2\tau$; halving keeps the requested tolerance.)

    The solver is **elementwise**: ``y`` may be an array, ``fn`` must act
    on each entry independently (as a CDF evaluated pointwise does), and
    every entry carries its own bracket and stops on its own. The loop
    ends when all entries have converged. Pass ``tags=frozenset(
    {lineax.diagonal_tag})`` to `optimistix.root_find` so the implicit
    adjoint uses the diagonal Jacobian.

    Requires ``options=dict(lower=..., upper=...)`` (broadcastable to
    ``y``), as `optimistix.Bisection` does. The bracket and every iterate
    keep ``y``'s dtype; only the function values take ``fn``'s. An entry
    whose bracket has no sign change (or a NaN endpoint value) is not
    iterated: it returns the endpoint with the smaller $|f|$, so the caller
    decides whether that is an error. With ``has_aux=True`` the auxiliary
    output is that of ``fn`` at the returned root (one extra evaluation).

    Pseudocode (per entry, Scherer 2010, §6.1):

        b, a, c = lower, upper, upper;  t = 1/2
        repeat:
            x_t = (1 − t) a + t b;  f_t = f(x_t)
            if sign f_t = sign f_a:  c ← a          else:  c ← b; b ← a
            a ← x_t
            x_m = argmin over a, b of |f|;  t_lim = (rtol |x_m| + atol) / (2 |b − c|)
            stop if f(x_m) = 0 or t_lim > 1/2 or x_t ∈ {a, b} before the update
            t = IQI(a, b, c) if Φ² < ξ and (1 − Φ)² < 1 − ξ else 1/2
            t = clip(t, t_lim, 1 − t_lim)

    Attributes:
        rtol: Relative tolerance on the root.
        atol: Absolute tolerance on the root.

    References:
        Chandrupatla, T. R. (1997). A new hybrid quadratic/bisection
        algorithm for finding the zero of a nonlinear function without
        using derivatives. *Advances in Engineering Software* 28(3),
        145-149.

        Scherer, P. O. J. (2010). *Computational Physics: Simulation of
        Classical and Quantum Systems*. Springer, §6.1.

    Examples:
        >>> import jax.numpy as jnp
        >>> import optimistix as optx
        >>> import gaussx
        >>> sol = optx.root_find(
        ...     lambda x, _: x**3 - 2.0,
        ...     gaussx.Chandrupatla(rtol=1e-10, atol=1e-12),
        ...     jnp.array(1.0),
        ...     options=dict(lower=0.0, upper=2.0),
        ... )
        >>> round(float(sol.value), 6)
        1.259921
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = optx.max_norm

    def init(
        self,
        fn: Callable[[Array, Any], tuple[Array, Any]],
        y: Array,
        args: PyTree,
        options: dict[str, Any],
        f_struct: Any,
        aux_struct: Any,
        tags: frozenset[object],
    ) -> _ChandrupatlaState:
        del aux_struct, tags
        if not isinstance(
            f_struct, jax.ShapeDtypeStruct
        ) or f_struct.shape != jnp.shape(y):
            raise ValueError(
                "Chandrupatla needs an elementwise function: the output must be "
                "a single array with the same shape as y."
            )
        shape, dtype = f_struct.shape, jnp.result_type(y)
        lower = jnp.broadcast_to(jnp.asarray(options["lower"], dtype), shape)
        upper = jnp.broadcast_to(jnp.asarray(options["upper"], dtype), shape)
        f_lower, _ = fn(lower, args)
        f_upper, _ = fn(upper, args)
        # Scherer's initialisation: b = lower, a = c = upper.
        a, b, c = upper, lower, upper
        fa, fb, fc = f_upper, f_lower, f_upper
        # NaN-aware: a finite residual always beats a NaN one.
        a_better = (jnp.abs(fa) < jnp.abs(fb)) | jnp.isnan(fb)
        x_best = jnp.where(a_better, a, b)
        f_best = jnp.where(a_better, fa, fb)
        # An infinite endpoint makes the convex-combination step NaN, so it
        # counts as unbracketed (finite brackets only).
        bracketed = (
            (jnp.sign(fa) * jnp.sign(fb) <= 0) & jnp.isfinite(a) & jnp.isfinite(b)
        )
        tol = self.rtol * jnp.abs(x_best) + self.atol
        done = (~bracketed) | (f_best == 0) | (jnp.abs(b - a) < tol)
        return _ChandrupatlaState(
            a=a,
            b=b,
            c=c,
            fa=fa,
            fb=fb,
            fc=fc,
            t=jnp.full(shape, 0.5, dtype),
            x_best=x_best,
            done=done,
        )

    def step(
        self,
        fn: Callable[[Array, Any], tuple[Array, Any]],
        y: Array,
        args: PyTree,
        options: dict[str, Any],
        state: _ChandrupatlaState,
        tags: frozenset[object],
    ) -> tuple[Array, _ChandrupatlaState, Any]:
        del y, options, tags
        s = state
        # The convex combination cannot overflow for finite a, b, unlike
        # a + t (b - a) on a bracket wider than the largest float.
        x_t = (1 - s.t) * s.a + s.t * s.b
        stalled = (x_t == s.a) | (x_t == s.b)
        f_t, aux = fn(x_t, args)

        same = jnp.sign(f_t) == jnp.sign(s.fa)
        c = jnp.where(same, s.a, s.b)
        fc = jnp.where(same, s.fa, s.fb)
        b = jnp.where(same, s.b, s.a)
        fb = jnp.where(same, s.fb, s.fa)
        a, fa = x_t, f_t

        a_better = (jnp.abs(fa) < jnp.abs(fb)) | jnp.isnan(fb)
        x_best = jnp.where(a_better, a, b)
        f_best = jnp.where(a_better, fa, fb)
        tol = self.rtol * jnp.abs(x_best) + self.atol
        t_lim = 0.5 * tol / jnp.abs(b - c)
        converged = (f_best == 0) | (t_lim > 0.5) | stalled

        # The divisions below may produce inf/nan on degenerate entries;
        # those entries either fail the IQI test (nan comparisons are
        # False) and bisect, or have already converged.
        xi = (a - b) / (c - b)
        phi = (fa - fb) / (fc - fb)
        use_iqi = (phi**2 < xi) & ((1 - phi) ** 2 < 1 - xi)
        t_iqi = fa / (fb - fa) * fc / (fb - fc) + (c - a) / (b - a) * fa / (
            fc - fa
        ) * fb / (fc - fb)
        t = jnp.clip(jnp.where(use_iqi, t_iqi, 0.5), t_lim, 1 - t_lim)

        new = _ChandrupatlaState(
            a=a,
            b=b,
            c=c,
            fa=fa,
            fb=fb,
            fc=fc,
            t=t,
            x_best=x_best,
            done=converged,
        )
        # Freeze entries that had already converged.
        new = jax.tree.map(lambda old, nw: jnp.where(s.done, old, nw), s, new)
        return new.x_best, new, aux

    def terminate(
        self,
        fn: Callable[[Array, Any], tuple[Array, Any]],
        y: Array,
        args: PyTree,
        options: dict[str, Any],
        state: _ChandrupatlaState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], optx.RESULTS]:
        del fn, y, args, options, tags
        return jnp.all(state.done), optx.RESULTS.successful

    def postprocess(
        self,
        fn: Callable[[Array, Any], tuple[Array, Any]],
        y: Array,
        aux: Any,
        args: PyTree,
        options: dict[str, Any],
        state: _ChandrupatlaState,
        tags: frozenset[object],
        result: optx.RESULTS,
    ) -> tuple[Array, Any, dict[str, Any]]:
        del y, aux, options, tags, result
        # Re-evaluate so ``aux`` belongs to the returned point; the step's
        # aux is that of the last trial point, which may be the other end.
        _, aux = fn(state.x_best, args)
        return state.x_best, aux, {}

mixture_quantile(cdf_fn: Callable[[Float[Array, ...]], Float[Array, ...]], q: Float[ArrayLike, ...], lower: Float[ArrayLike, ...], upper: Float[ArrayLike, ...], *, solver: optx.AbstractRootFinder | None = None, rtol: float = 0.0001, atol: float = 1e-06, max_steps: int = 64, throw: bool = False) -> Float[Array, ...]

Invert a CDF at one or more quantile levels by bracketed root finding.

The \(q\)-quantile of a continuous CDF \(F\) is the root of \(g_q(x) = F(x) - q\). \(F\) is monotone, so \(g_q\) changes sign exactly once on any bracket \([\ell, u]\) with \(F(\ell) \le q \le F(u)\), and a bracketing method needs neither derivatives nor a good initial guess. A mixture or ensemble CDF

\[ F_{\text{mix}}(x) = \frac{1}{E}\sum_{e=1}^{E} F_e(x) \]

has no closed-form inverse even when every \(F_e\) does; this is what the function is for. The default solver is Chandrupatla, run through optimistix.root_find.

Shapes. q, lower and upper broadcast together (NumPy rules) to the output shape, and cdf_fn is called on arrays of that shape. It must act elementwise — entry i of its output may depend only on entry i of its input — e.g. a scalar CDF, or a per-test-point mixture CDF with a trailing quantile axis. With Q levels q of shape (Q,) and a batch of brackets, give the brackets a trailing axis, (*batch, 1), to get (*batch, Q).

Gradients go through optimistix.ImplicitAdjoint: for a root \(x_q\) of \(F(x; \theta) - q\),

\[ \frac{\partial x_q}{\partial \theta} = -\frac{\partial_\theta F(x_q; \theta)}{p(x_q; \theta)}, \qquad \frac{\partial x_q}{\partial q} = \frac{1}{p(x_q; \theta)}, \]

with \(p = \partial_x F\) the density, evaluated with one Jacobian-vector product (the Jacobian is tagged diagonal). Parameters closed over by cdf_fn are differentiated too. The brackets are not.

Pseudocode:

F_lo, F_hi = cdf(lower) − q, cdf(upper) − q
if throw: error unless F_lo · F_hi ≤ 0 everywhere
x = root_find(x ↦ cdf(x) − q, Chandrupatla, bracket [lower, upper],
              adjoint = implicit, Jacobian tag = diagonal)

Parameters:

Name	Type	Description	Default
cdf_fn	Callable[[Float[Array, ...]], Float[Array, ...]]	Elementwise CDF, mapping an array to an array of the same shape. It must be strictly increasing where it crosses q: on a flat segment at level q the solver returns some point of the segment, not the generalized inverse inf{x : F(x) >= q}. Gaussian-mixture CDFs always qualify.	required
q	Float[ArrayLike, ...]	Quantile levels in \([0, 1]\).	required
lower	Float[ArrayLike, ...]	Finite lower bracket, with cdf_fn(lower) <= q.	required
upper	Float[ArrayLike, ...]	Finite upper bracket, with cdf_fn(upper) >= q. Infinite brackets are rejected; for a Gaussian mixture, means ± 40 standard deviations bracket every level.	required
solver	AbstractRootFinder | None	An optimistix root finder that takes options=dict(lower=..., upper=...). Defaults to Chandrupatla(rtol=rtol, atol=atol). optimistix.Bisection also fits, for scalar problems only.	None
rtol	float	Relative tolerance on the quantile (default solver only).	0.0001
atol	float	Absolute tolerance on the quantile (default solver only).	1e-06
max_steps	int	Maximum number of solver iterations.	64
throw	bool	If True, raise when a level is not bracketed (cdf_fn(lower) <= q <= cdf_fn(upper) fails, including a NaN endpoint value, an infinite endpoint or a reversed bracket) or the solver does not converge. If False, an unbracketed level returns the bracket endpoint whose CDF is closer to \(q\) (its gradient is then meaningless), and a non-converged one the best estimate so far.	False

Returns:

Type	Description
Float[Array, ...]	Quantiles, with the broadcast shape of q, lower and
Float[Array, ...]	upper.

References

Chandrupatla, T. R. (1997). A new hybrid quadratic/bisection algorithm for finding the zero of a nonlinear function without using derivatives. Advances in Engineering Software 28(3), 145-149.

Examples:

>>> import jax.numpy as jnp
>>> from jax.scipy.stats import norm
>>> import gaussx
>>> cdf = lambda x: 0.5 * (norm.cdf(x + 1.0) + norm.cdf(x - 1.0))
>>> x = gaussx.mixture_quantile(cdf, jnp.array([0.5, 0.975]), -10.0, 10.0)
>>> [round(float(v), 3) for v in x]
[0.0, 2.646]

Source code in src/gaussx/_primitives/_quantile.py

def mixture_quantile(
    cdf_fn: Callable[[Float[Array, ...]], Float[Array, ...]],
    q: Float[ArrayLike, ...],
    lower: Float[ArrayLike, ...],
    upper: Float[ArrayLike, ...],
    *,
    solver: optx.AbstractRootFinder | None = None,
    rtol: float = 1e-4,
    atol: float = 1e-6,
    max_steps: int = 64,
    throw: bool = False,
) -> Float[Array, ...]:
    r"""Invert a CDF at one or more quantile levels by bracketed root finding.

    The $q$-quantile of a continuous CDF $F$ is the root of
    $g_q(x) = F(x) - q$. $F$ is monotone, so $g_q$ changes sign exactly
    once on any bracket $[\ell, u]$ with $F(\ell) \le q \le F(u)$, and a
    bracketing method needs neither derivatives nor a good initial guess.
    A mixture or ensemble CDF

    $$
    F_{\text{mix}}(x) = \frac{1}{E}\sum_{e=1}^{E} F_e(x)
    $$

    has no closed-form inverse even when every $F_e$ does; this is what
    the function is for. The default solver is `Chandrupatla`, run through
    `optimistix.root_find`.

    **Shapes.** ``q``, ``lower`` and ``upper`` broadcast together (NumPy
    rules) to the output shape, and ``cdf_fn`` is called on arrays of that
    shape. It must act **elementwise** — entry ``i`` of its output may
    depend only on entry ``i`` of its input — e.g. a scalar CDF, or a
    per-test-point mixture CDF with a trailing quantile axis. With ``Q``
    levels ``q`` of shape ``(Q,)`` and a batch of brackets, give the
    brackets a trailing axis, ``(*batch, 1)``, to get ``(*batch, Q)``.

    **Gradients** go through `optimistix.ImplicitAdjoint`: for a root
    $x_q$ of $F(x; \theta) - q$,

    $$
    \frac{\partial x_q}{\partial \theta}
        = -\frac{\partial_\theta F(x_q; \theta)}{p(x_q; \theta)},
    \qquad \frac{\partial x_q}{\partial q} = \frac{1}{p(x_q; \theta)},
    $$

    with $p = \partial_x F$ the density, evaluated with one Jacobian-vector
    product (the Jacobian is tagged diagonal). Parameters closed over by
    ``cdf_fn`` are differentiated too. The brackets are not.

    Pseudocode:

        F_lo, F_hi = cdf(lower) − q, cdf(upper) − q
        if throw: error unless F_lo · F_hi ≤ 0 everywhere
        x = root_find(x ↦ cdf(x) − q, Chandrupatla, bracket [lower, upper],
                      adjoint = implicit, Jacobian tag = diagonal)

    Args:
        cdf_fn: Elementwise CDF, mapping an array to an array of the same
            shape. It must be strictly increasing where it crosses ``q``: on
            a flat segment at level ``q`` the solver returns some point of
            the segment, not the generalized inverse
            ``inf{x : F(x) >= q}``. Gaussian-mixture CDFs always qualify.
        q: Quantile levels in $[0, 1]$.
        lower: Finite lower bracket, with ``cdf_fn(lower) <= q``.
        upper: Finite upper bracket, with ``cdf_fn(upper) >= q``. Infinite
            brackets are rejected; for a Gaussian mixture, means ± 40
            standard deviations bracket every level.
        solver: An optimistix root finder that takes
            ``options=dict(lower=..., upper=...)``. Defaults to
            ``Chandrupatla(rtol=rtol, atol=atol)``. `optimistix.Bisection`
            also fits, for scalar problems only.
        rtol: Relative tolerance on the quantile (default solver only).
        atol: Absolute tolerance on the quantile (default solver only).
        max_steps: Maximum number of solver iterations.
        throw: If ``True``, raise when a level is not bracketed
            (``cdf_fn(lower) <= q <= cdf_fn(upper)`` fails, including a NaN
            endpoint value, an infinite endpoint or a reversed bracket) or
            the solver does not
            converge. If ``False``, an unbracketed level
            returns the bracket endpoint whose CDF is closer to $q$ (its
            gradient is then meaningless), and a non-converged one the
            best estimate so far.

    Returns:
        Quantiles, with the broadcast shape of ``q``, ``lower`` and
        ``upper``.

    References:
        Chandrupatla, T. R. (1997). A new hybrid quadratic/bisection
        algorithm for finding the zero of a nonlinear function without
        using derivatives. *Advances in Engineering Software* 28(3),
        145-149.

    Examples:
        >>> import jax.numpy as jnp
        >>> from jax.scipy.stats import norm
        >>> import gaussx
        >>> cdf = lambda x: 0.5 * (norm.cdf(x + 1.0) + norm.cdf(x - 1.0))
        >>> x = gaussx.mixture_quantile(cdf, jnp.array([0.5, 0.975]), -10.0, 10.0)
        >>> [round(float(v), 3) for v in x]
        [0.0, 2.646]
    """
    q = jnp.asarray(q)
    dtype = jnp.result_type(q, lower, upper, float)
    q = q.astype(dtype)
    shape = jnp.broadcast_shapes(q.shape, jnp.shape(lower), jnp.shape(upper))
    lower = jnp.broadcast_to(jnp.asarray(lower, dtype), shape)
    upper = jnp.broadcast_to(jnp.asarray(upper, dtype), shape)
    if solver is None:
        solver = Chandrupatla(rtol=rtol, atol=atol)

    def fn(x: Float[Array, ...], level: Float[Array, ...]) -> Float[Array, ...]:
        return cdf_fn(x) - level

    if throw:
        # Directed checks: they also reject NaN endpoint values and a
        # reversed bracket, which a sign-product test would let through.
        g_lo, g_hi = fn(lower, q), fn(upper, q)
        lower = eqx.error_if(
            lower,
            ~jnp.all(
                (g_lo <= 0) & (g_hi >= 0) & jnp.isfinite(lower) & jnp.isfinite(upper)
            ),
            "mixture_quantile: some levels q are not bracketed by finite "
            "endpoints, i.e. cdf_fn(lower) <= q <= cdf_fn(upper) fails or a "
            "bracket is infinite.",
        )
    sol = optx.root_find(
        fn,
        solver,
        0.5 * lower + 0.5 * upper,  # no overflow for wide finite brackets
        args=jnp.broadcast_to(q, shape),
        options=dict(lower=lower, upper=upper),
        max_steps=max_steps,
        adjoint=optx.ImplicitAdjoint(),
        throw=throw,
        tags=frozenset({lx.diagonal_tag}),
    )
    return sol.value

mixture_quantile_gaussian_approx(means: Float[Array, '*batch E'], stds: Float[Array, '*batch E'], q: Float[ArrayLike, ' Q']) -> Float[Array, '*batch Q']

Quantiles of the Gaussian matched to an ensemble's first two moments.

For an equal-weight ensemble of Gaussians \(\mathcal{N}(\mu_e, \sigma_e^2)\) the mixture has mean and variance

\[ \hat\mu = \frac{1}{E}\sum_e \mu_e, \qquad \hat\sigma^2 = \frac{1}{E}\sum_e \big(\sigma_e^2 + (\mu_e - \hat\mu)^2\big), \]

and its quantiles are approximated by those of \(\mathcal{N}(\hat\mu, \hat\sigma^2)\), \(\hat x_q = \hat\mu + \hat\sigma\,\Phi^{-1}(q)\) — the deep-ensemble predictive of Lakshminarayanan, Pritzel & Blundell (2017). It is exact for a single component and cheap otherwise: a diagnostic, or a seed for the brackets of mixture_quantile. The variance is accumulated from deviations about \(\hat\mu\), which equals the raw-moment form \(\frac{1}{E}\sum_e(\sigma_e^2 + \mu_e^2) - \hat\mu^2\) without its cancellation.

Parameters:

Name	Type	Description	Default
means	Float[Array, '*batch E']	Component means, shape (*batch, E).	required
stds	Float[Array, '*batch E']	Component standard deviations, shape (*batch, E).	required
q	Float[ArrayLike, ' Q']	Quantile levels, shape (Q,).	required

Returns:

Type	Description
Float[Array, '*batch Q']	Approximate quantiles, shape (*batch, Q).

References

Lakshminarayanan, B., Pritzel, A. & Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. NeurIPS 30.

Examples:

>>> import jax.numpy as jnp
>>> import gaussx
>>> x = gaussx.mixture_quantile_gaussian_approx(
...     jnp.array([[0.0]]), jnp.array([[2.0]]), jnp.array([0.5, 0.8413447])
... )
>>> [round(float(v), 4) for v in x[0]]
[0.0, 2.0]

Source code in src/gaussx/_primitives/_quantile.py

def mixture_quantile_gaussian_approx(
    means: Float[Array, "*batch E"],
    stds: Float[Array, "*batch E"],
    q: Float[ArrayLike, " Q"],
) -> Float[Array, "*batch Q"]:
    r"""Quantiles of the Gaussian matched to an ensemble's first two moments.

    For an equal-weight ensemble of Gaussians $\mathcal{N}(\mu_e,
    \sigma_e^2)$ the mixture has mean and variance

    $$
    \hat\mu = \frac{1}{E}\sum_e \mu_e, \qquad
    \hat\sigma^2 = \frac{1}{E}\sum_e \big(\sigma_e^2 + (\mu_e - \hat\mu)^2\big),
    $$

    and its quantiles are approximated by those of
    $\mathcal{N}(\hat\mu, \hat\sigma^2)$,
    $\hat x_q = \hat\mu + \hat\sigma\,\Phi^{-1}(q)$ — the deep-ensemble
    predictive of Lakshminarayanan, Pritzel & Blundell (2017). It is exact
    for a single component and cheap otherwise: a diagnostic, or a seed
    for the brackets of `mixture_quantile`. The variance is accumulated
    from deviations about $\hat\mu$, which equals the raw-moment form
    $\frac{1}{E}\sum_e(\sigma_e^2 + \mu_e^2) - \hat\mu^2$ without its
    cancellation.

    Args:
        means: Component means, shape ``(*batch, E)``.
        stds: Component standard deviations, shape ``(*batch, E)``.
        q: Quantile levels, shape ``(Q,)``.

    Returns:
        Approximate quantiles, shape ``(*batch, Q)``.

    References:
        Lakshminarayanan, B., Pritzel, A. & Blundell, C. (2017). Simple and
        scalable predictive uncertainty estimation using deep ensembles.
        *NeurIPS 30*.

    Examples:
        >>> import jax.numpy as jnp
        >>> import gaussx
        >>> x = gaussx.mixture_quantile_gaussian_approx(
        ...     jnp.array([[0.0]]), jnp.array([[2.0]]), jnp.array([0.5, 0.8413447])
        ... )
        >>> [round(float(v), 4) for v in x[0]]
        [0.0, 2.0]
    """
    dtype = jnp.result_type(means, stds, q, float)
    means = jnp.asarray(means, dtype)
    stds = jnp.asarray(stds, dtype)
    q = jnp.atleast_1d(jnp.asarray(q, dtype))
    mu = reduce(means, "... e -> ...", "mean")
    dev = einx.subtract("... e, ... -> ... e", means, mu)
    # Rescaled norm: stds**2 would overflow for stds near sqrt(max float).
    big = jnp.maximum(
        reduce(jnp.abs(stds), "... e -> ...", "max"),
        reduce(jnp.abs(dev), "... e -> ...", "max"),
    )
    big = jnp.where(big > 0, big, jnp.ones_like(big))
    s_r = einx.divide("... e, ... -> ... e", stds, big)
    d_r = einx.divide("... e, ... -> ... e", dev, big)
    sigma = big * jnp.sqrt(reduce(s_r**2 + d_r**2, "... e -> ...", "mean"))
    z = jax.scipy.special.ndtri(q)
    return einx.add(
        "..., ... q -> ... q", mu, einx.multiply("..., q -> ... q", sigma, z)
    )

Support types

Structured linear algebra and Gaussian primitives for JAX.

DenseFallbackWarning

Bases: UserWarning

Warning emitted when a structured primitive materialises an operator.

Source code in src/gaussx/_primitives/_cholesky.py

class DenseFallbackWarning(UserWarning):
    """Warning emitted when a structured primitive materialises an operator."""