Operators & Tags

Layer 1: structured linear operators extending lineax.AbstractLinearOperator. All are immutable equinox.Module pytrees, so they compose freely with jit, grad, and vmap. The primitives dispatch on these types: a solve against a Kronecker factorizes per Kronecker factor, a logdet of a BlockDiag sums per block, a LowRankUpdate solve applies Woodbury.

Structured products & sums

The Kronecker product \(A_1 \otimes A_2 \otimes \cdots\) gives \(O(\sum_i n_i^3)\) solves on a \(\prod_i n_i\) grid; the Kronecker sum \(A \otimes I + I \otimes B\) diagonalises in the joint eigenbasis with eigenvalues \(\lambda_i + \mu_j\).

KroneckerSum and SumOfKroneckers are different operators

They are not word-order variants of one idea:

	Math	Eigendecomposition
KroneckerSum(A, B)	\(A \otimes I + I \otimes B\)	Closed form — eigenvalues \(\lambda_i + \mu_j\), eigenvectors \(V_A \otimes V_B\)
SumOfKroneckers(K_1, K_2, …)	\(\sum_k A_k \otimes B_k\)	None in general. solve / logdet reduce the two-term case with one term positive definite — \(B \otimes C + \sigma^2 I\) and \(B_1 \otimes C_1 + B_2 \otimes C_2\) — to one eigendecomposition per factor, \(O(n_c^3 + n_d^3)\); eigendecompose covers the same two-term symmetric case more generally by densifying a \((n_c n_d)^2\) block

SumOfKroneckers was called SumKronecker until gh-136. The old name still imports and subclasses the new one — so isinstance checks keep working — but emits a DeprecationWarning on construction and will be removed in gaussx 0.7.0. Its square root and sampler followed in gh-297: SumKroneckerSqrt → SumOfKroneckersSqrt and sumkronecker_sample → sum_of_kroneckers_sample. The old names warn and will be removed in gaussx 0.7.0. Each square-root class sits next to its operator: KroneckerSumSqrt is the closed-form root of a KroneckerSum, and SumOfKroneckersSqrt is the lazy Lanczos root that sqrt returns for a SumOfKroneckers.

Structured linear algebra and Gaussian primitives for JAX.

Kronecker

Bases: AbstractLinearOperator

Kronecker product operator A₁ ⊗ A₂ ⊗ … ⊗ Aₖ.

Matvec uses Roth's column lemma for efficient computation without materializing the full Kronecker product. For two factors A (m x n) and B (p x q), the product (A kron B) vec(X) is computed as vec(B X A^T) where X is reshaped to (q, n).

Complexity: O(sum n_i^3) instead of O((prod n_i)^2) for the naive approach.

Parameters:

Name	Type	Description	Default
*operators	AbstractLinearOperator	Two or more lineax.AbstractLinearOperator instances.	()

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)  # 4 x 4, stored as two 2 x 2 factors
>>> K.in_size()
4
>>> bool(jnp.allclose(K.as_matrix(), jnp.kron(A.as_matrix(), B.as_matrix())))
True
>>> type(gaussx.cholesky(K)).__name__  # primitives keep the structure
'Kronecker'

Source code in src/gaussx/_operators/_kronecker.py

class Kronecker(lx.AbstractLinearOperator):
    """Kronecker product operator ``A₁ ⊗ A₂ ⊗ … ⊗ Aₖ``.

    Matvec uses Roth's column lemma for efficient computation without
    materializing the full Kronecker product. For two factors A (m x n)
    and B (p x q), the product (A kron B) vec(X) is computed as
    vec(B X A^T) where X is reshaped to (q, n).

    Complexity: O(sum n_i^3) instead of O((prod n_i)^2) for the naive approach.

    Args:
        *operators: Two or more ``lineax.AbstractLinearOperator`` instances.

    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)  # 4 x 4, stored as two 2 x 2 factors
        >>> K.in_size()
        4
        >>> bool(jnp.allclose(K.as_matrix(), jnp.kron(A.as_matrix(), B.as_matrix())))
        True
        >>> type(gaussx.cholesky(K)).__name__  # primitives keep the structure
        'Kronecker'
    """

    operators: tuple[lx.AbstractLinearOperator, ...]
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        *operators: lx.AbstractLinearOperator,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        if len(operators) < 2:
            raise ValueError("Kronecker requires at least two operators.")
        self.operators = operators
        self._in_size = _prod(op.in_size() for op in operators)
        self._out_size = _prod(op.out_size() for op in operators)
        self._dtype = _resolve_dtype(*operators)

        self.tags = _to_frozenset(tags) | {kronecker_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " m"]:
        return _kronecker_mv(self.operators, vector)

    def as_matrix(self) -> Float[Array, "m n"]:
        result = self.operators[0].as_matrix()
        for op in self.operators[1:]:
            result = jnp.kron(result, op.as_matrix())
        return result

    def transpose(self) -> Kronecker:
        return Kronecker(
            *(op.T for op in self.operators),
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

BlockDiag

Bases: AbstractLinearOperator

Block diagonal operator diag(A₁, A₂, …, Aₖ).

Each sub-operator acts on its own slice of the input vector. Matvec, transpose, logdet, solve, and cholesky all decompose per-block.

Parameters:

Name	Type	Description	Default
*operators	AbstractLinearOperator	One or more lineax.AbstractLinearOperator instances forming the diagonal blocks.	()

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)
>>> D = gaussx.BlockDiag(A, B)  # diag(A, B), 4 x 4
>>> D.as_matrix().shape
(4, 4)
>>> logdet_blocks = gaussx.logdet(A) + gaussx.logdet(B)
>>> bool(jnp.allclose(gaussx.logdet(D), logdet_blocks))
True

Source code in src/gaussx/_operators/_block_diag.py

class BlockDiag(lx.AbstractLinearOperator):
    """Block diagonal operator ``diag(A₁, A₂, …, Aₖ)``.

    Each sub-operator acts on its own slice of the input vector.
    Matvec, transpose, logdet, solve, and cholesky all decompose
    per-block.

    Args:
        *operators: One or more ``lineax.AbstractLinearOperator`` instances
            forming the diagonal blocks.

    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)
        >>> D = gaussx.BlockDiag(A, B)  # diag(A, B), 4 x 4
        >>> D.as_matrix().shape
        (4, 4)
        >>> logdet_blocks = gaussx.logdet(A) + gaussx.logdet(B)
        >>> bool(jnp.allclose(gaussx.logdet(D), logdet_blocks))
        True
    """

    operators: tuple[lx.AbstractLinearOperator, ...]
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        *operators: lx.AbstractLinearOperator,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        if len(operators) == 0:
            raise ValueError("BlockDiag requires at least one operator.")
        self.operators = operators
        self._in_size = sum(op.in_size() for op in operators)
        self._out_size = sum(op.out_size() for op in operators)
        self._dtype = _resolve_dtype(*operators)

        self.tags = _to_frozenset(tags) | {block_diagonal_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " m"]:
        # Use jax.lax.dynamic_slice with static offsets for JIT compatibility
        results = []
        offset = 0
        for op in self.operators:
            size = op.in_size()
            block = jax.lax.dynamic_slice(vector, (offset,), (size,))
            results.append(op.mv(block))
            offset += size
        return jnp.concatenate(results)

    def as_matrix(self) -> Float[Array, "m n"]:
        matrices = [op.as_matrix() for op in self.operators]
        return jax.scipy.linalg.block_diag(*matrices)

    def transpose(self) -> BlockDiag:
        return BlockDiag(
            *(op.T for op in self.operators),
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

KroneckerSum

Bases: AbstractLinearOperator

Kronecker sum A \oplus B = A \otimes I_b + I_a \otimes B.

Appears in separable PDEs, graph Laplacians, and space-time GPs. If A = Q_A \Lambda_A Q_A^T and B = Q_B \Lambda_B Q_B^T, the Kronecker sum has eigenvectors Q_A \otimes Q_B with eigenvalues \lambda^A_i + \lambda^B_j.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator	First operator, shape (n_a, n_a).	required
B	AbstractLinearOperator	Second operator, shape (n_b, n_b).	required

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.KroneckerSum(A, B)  # A ⊗ I + I ⊗ B
>>> I2 = jnp.eye(2)
>>> dense = jnp.kron(A.as_matrix(), I2) + jnp.kron(I2, B.as_matrix())
>>> bool(jnp.allclose(S.as_matrix(), dense))
True
>>> sign, expected = jnp.linalg.slogdet(dense)
>>> bool(jnp.allclose(gaussx.logdet(S), expected, atol=1e-5))  # λ_i + μ_j
True

Source code in src/gaussx/_operators/_kronecker_sum.py

class KroneckerSum(lx.AbstractLinearOperator):
    r"""Kronecker sum ``A \oplus B = A \otimes I_b + I_a \otimes B``.

    Appears in separable PDEs, graph Laplacians, and space-time GPs.
    If ``A = Q_A \Lambda_A Q_A^T`` and ``B = Q_B \Lambda_B Q_B^T``,
    the Kronecker sum has eigenvectors ``Q_A \otimes Q_B`` with
    eigenvalues ``\lambda^A_i + \lambda^B_j``.

    Args:
        A: First operator, shape ``(n_a, n_a)``.
        B: Second operator, shape ``(n_b, n_b)``.

    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.KroneckerSum(A, B)  # A ⊗ I + I ⊗ B
        >>> I2 = jnp.eye(2)
        >>> dense = jnp.kron(A.as_matrix(), I2) + jnp.kron(I2, B.as_matrix())
        >>> bool(jnp.allclose(S.as_matrix(), dense))
        True
        >>> sign, expected = jnp.linalg.slogdet(dense)
        >>> bool(jnp.allclose(gaussx.logdet(S), expected, atol=1e-5))  # λ_i + μ_j
        True
    """

    A: lx.AbstractLinearOperator
    B: lx.AbstractLinearOperator
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _n_a: int = eqx.field(static=True)
    _n_b: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        A: lx.AbstractLinearOperator,
        B: lx.AbstractLinearOperator,
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        if A.in_size() != A.out_size():
            raise ValueError(
                f"A must be square, got in_size={A.in_size()}, out_size={A.out_size()}."
            )
        if B.in_size() != B.out_size():
            raise ValueError(
                f"B must be square, got in_size={B.in_size()}, out_size={B.out_size()}."
            )
        self.A = A
        self.B = B
        n_a = A.in_size()
        n_b = B.in_size()
        self._n_a = n_a
        self._n_b = n_b
        self._in_size = n_a * n_b
        self._out_size = n_a * n_b
        self._dtype = _resolve_dtype(A, B)

        self.tags = _to_frozenset(tags) | {kronecker_sum_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        # (A (x) I_b + I_a (x) B) vec(X) = vec(B X + X A^T)
        # where X is (n_b, n_a)
        X = rearrange(vector, "(a b) -> b a", a=self._n_a, b=self._n_b)
        # I_a (x) B: apply B to each column of X
        BX = jax.vmap(self.B.mv, in_axes=1, out_axes=1)(X)
        # A (x) I_b: apply A^T to each row of X (= apply A to rows of X^T)
        XAt = jax.vmap(self.A.mv)(X)  # (n_b, n_a): apply A to rows
        result = BX + XAt
        return rearrange(result, "b a -> (a b)")

    def as_matrix(self) -> Float[Array, "n n"]:
        A_mat = self.A.as_matrix()
        B_mat = self.B.as_matrix()
        I_a = jnp.eye(self._n_a, dtype=jnp.dtype(self._dtype))
        I_b = jnp.eye(self._n_b, dtype=jnp.dtype(self._dtype))
        return jnp.kron(A_mat, I_b) + jnp.kron(I_a, B_mat)

    def transpose(self) -> KroneckerSum:
        return KroneckerSum(
            self.A.T,
            self.B.T,
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

    def eigendecompose(
        self,
    ) -> tuple[Float[Array, " n"], Float[Array, "n n"]]:
        """Symmetric eigendecomposition via per-factor decomposition.

        Assumes both factors are symmetric so the returned
        ``Q = Q_A ⊗ Q_B`` is orthonormal — callers rely on
        ``self == Q @ diag(eigenvalues) @ Q.T``. Diagonal factors get a
        structural shortcut; other operators are materialized and
        decomposed via ``jnp.linalg.eigh``. We deliberately avoid
        routing untagged factors through `gaussx.eig` because
        that primitive falls back to ``jnp.linalg.eig`` for untagged
        operators and would return general (non-orthonormal)
        eigenvectors — breaking the ``Q.T == Q^{-1}`` contract for the
        common case of numerically symmetric matrices wrapped as plain
        `lineax.MatrixLinearOperator`.

        Returns:
            Tuple ``(eigenvalues, Q)`` where ``Q = Q_A ⊗ Q_B`` and the
            eigenvalues are ``lambda^A_i + lambda^B_j`` for all pairs.
        """

        evals_a, evecs_a = _eigh_factor(self.A)
        evals_b, evecs_b = _eigh_factor(self.B)
        # Eigenvalues: lambda_a_i + lambda_b_j for all (i, j) pairs
        eigenvalues = rearrange(evals_a[:, None] + evals_b[None, :], "a b -> (a b)")
        # Eigenvectors: Q_A (x) Q_B
        Q = jnp.kron(evecs_a, evecs_b)
        return eigenvalues, Q

eigendecompose() -> tuple[Float[Array, ' n'], Float[Array, 'n n']]

Symmetric eigendecomposition via per-factor decomposition.

Assumes both factors are symmetric so the returned Q = Q_A ⊗ Q_B is orthonormal — callers rely on self == Q @ diag(eigenvalues) @ Q.T. Diagonal factors get a structural shortcut; other operators are materialized and decomposed via jnp.linalg.eigh. We deliberately avoid routing untagged factors through gaussx.eig because that primitive falls back to jnp.linalg.eig for untagged operators and would return general (non-orthonormal) eigenvectors — breaking the Q.T == Q^{-1} contract for the common case of numerically symmetric matrices wrapped as plain lineax.MatrixLinearOperator.

Returns:

Type	Description
Float[Array, ' n']	Tuple (eigenvalues, Q) where Q = Q_A ⊗ Q_B and the
Float[Array, 'n n']	eigenvalues are lambda^A_i + lambda^B_j for all pairs.

Source code in src/gaussx/_operators/_kronecker_sum.py

def eigendecompose(
    self,
) -> tuple[Float[Array, " n"], Float[Array, "n n"]]:
    """Symmetric eigendecomposition via per-factor decomposition.

    Assumes both factors are symmetric so the returned
    ``Q = Q_A ⊗ Q_B`` is orthonormal — callers rely on
    ``self == Q @ diag(eigenvalues) @ Q.T``. Diagonal factors get a
    structural shortcut; other operators are materialized and
    decomposed via ``jnp.linalg.eigh``. We deliberately avoid
    routing untagged factors through `gaussx.eig` because
    that primitive falls back to ``jnp.linalg.eig`` for untagged
    operators and would return general (non-orthonormal)
    eigenvectors — breaking the ``Q.T == Q^{-1}`` contract for the
    common case of numerically symmetric matrices wrapped as plain
    `lineax.MatrixLinearOperator`.

    Returns:
        Tuple ``(eigenvalues, Q)`` where ``Q = Q_A ⊗ Q_B`` and the
        eigenvalues are ``lambda^A_i + lambda^B_j`` for all pairs.
    """

    evals_a, evecs_a = _eigh_factor(self.A)
    evals_b, evecs_b = _eigh_factor(self.B)
    # Eigenvalues: lambda_a_i + lambda_b_j for all (i, j) pairs
    eigenvalues = rearrange(evals_a[:, None] + evals_b[None, :], "a b -> (a b)")
    # Eigenvectors: Q_A (x) Q_B
    Q = jnp.kron(evecs_a, evecs_b)
    return eigenvalues, Q

KroneckerSumSqrt

Bases: AbstractLinearOperator

Symmetric square root of A \oplus B via per-factor eigenvectors.

Represents the symmetric matrix S with S @ S = A \oplus B (where \oplus is the Kronecker sum A ⊗ I + I ⊗ B). The square root is never materialized: mv and solve apply S and S^{-1} matrix-free using the per-factor eigendecompositions, so the cost is governed by the factor sizes rather than the full n_a · n_b dimension.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator	Symmetric PSD factor, shape (n_a, n_a).	required
B	AbstractLinearOperator	Symmetric PSD factor, shape (n_b, n_b).	required

Raises:

Type	Description
ValueError	If either factor is non-square or untagged as symmetric.
EquinoxRuntimeError	If A \oplus B is not positive semidefinite (checked with equinox.error_if, so also under jax.jit).

Source code in src/gaussx/_operators/_kronecker_sum.py

class KroneckerSumSqrt(lx.AbstractLinearOperator):
    r"""Symmetric square root of ``A \oplus B`` via per-factor eigenvectors.

    Represents the symmetric matrix ``S`` with ``S @ S = A \oplus B``
    (where ``\oplus`` is the Kronecker sum ``A ⊗ I + I ⊗ B``). The square
    root is never materialized: `mv` and `solve` apply ``S`` and
    ``S^{-1}`` matrix-free using the per-factor eigendecompositions, so the
    cost is governed by the factor sizes rather than the full ``n_a · n_b``
    dimension.

    Args:
        A: Symmetric PSD factor, shape ``(n_a, n_a)``.
        B: Symmetric PSD factor, shape ``(n_b, n_b)``.

    Raises:
        ValueError: If either factor is non-square or untagged as symmetric.
        EquinoxRuntimeError: If ``A \oplus B`` is not positive semidefinite
            (checked with `equinox.error_if`, so also under ``jax.jit``).
    """

    a_factor: Float[Array, "a a"] | Float[Array, " a"]
    b_factor: Float[Array, "b b"] | Float[Array, " b"]
    eigenvectors_a: Float[Array, "a a"]
    eigenvectors_b: Float[Array, "b b"]
    sqrt_eigenvalues: Float[Array, "a b"]
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _n_a: int = eqx.field(static=True)
    _n_b: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)

    def __init__(
        self,
        A: lx.AbstractLinearOperator,
        B: lx.AbstractLinearOperator,
    ) -> None:
        if A.in_size() != A.out_size():
            raise ValueError(
                f"A must be square, got in_size={A.in_size()}, out_size={A.out_size()}."
            )
        if B.in_size() != B.out_size():
            raise ValueError(
                f"B must be square, got in_size={B.in_size()}, out_size={B.out_size()}."
            )
        # The symmetric sqrt is well-defined only when A and B are
        # symmetric PSD. Without these tags, ``jnp.linalg.eigh`` would
        # silently use only the lower triangle and return wrong
        # eigenvectors for non-symmetric inputs.
        if not lx.is_symmetric(A) or not lx.is_symmetric(B):
            raise ValueError(
                "KroneckerSumSqrt requires both factors to be symmetric "
                "(tag them with lx.symmetric_tag or lx.positive_semidefinite_tag)."
            )
        # The factors (a matrix, or the diagonal of a diagonal factor, which
        # is never materialised) are the differentiable leaves; the cached
        # eigendecomposition is a constant for autodiff. `mv` and `solve` go
        # through custom JVPs written in terms of the factors, because the
        # eigenvector derivative divides by eigenvalue gaps and is NaN at a
        # repeated eigenvalue (e.g. an isotropic factor ``s I``), gh-295.
        a_factor = _factor_array(A)
        b_factor = _factor_array(B)
        evals_a, evecs_a = _factor_eigh(jax.lax.stop_gradient(a_factor))
        evals_b, evecs_b = _factor_eigh(jax.lax.stop_gradient(b_factor))
        eigenvalues = _checked_kronecker_sum_spectrum(evals_a, evals_b)
        sqrt_eigenvalues = jnp.sqrt(jnp.maximum(eigenvalues, 0.0))

        self.a_factor = a_factor
        self.b_factor = b_factor
        self.eigenvectors_a = evecs_a
        self.eigenvectors_b = evecs_b
        self.sqrt_eigenvalues = sqrt_eigenvalues
        self._n_a = A.in_size()
        self._n_b = B.in_size()
        self._in_size = self._n_a * self._n_b
        self._out_size = self._in_size
        self._dtype = str(jnp.result_type(evecs_a, evecs_b, sqrt_eigenvalues))

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        Z = rearrange(vector, "(a b) -> 1 a b", a=self._n_a, b=self._n_b)
        result = _kronecker_sum_root_apply(*self._parts(), Z)
        return rearrange(result, "1 a b -> (a b)")

    def solve(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        """Apply the inverse square root ``S^{-1}`` to ``vector``.

        Args:
            vector: Input vector, shape ``(n_a · n_b,)``.

        Returns:
            ``S^{-1} @ vector``, shape ``(n_a · n_b,)``.
        """
        Z = rearrange(vector, "(a b) -> 1 a b", a=self._n_a, b=self._n_b)
        result = _kronecker_sum_root_solve(*self._parts(), Z)
        return rearrange(result, "1 a b -> (a b)")

    def _parts(self):
        return (
            self.a_factor,
            self.b_factor,
            self.eigenvectors_a,
            self.eigenvectors_b,
            self.sqrt_eigenvalues,
        )

    def as_matrix(self) -> Float[Array, "n n"]:
        basis = jnp.eye(self._in_size, dtype=jnp.dtype(self._dtype))
        return jax.vmap(self.mv, in_axes=1, out_axes=1)(basis)

    def transpose(self) -> KroneckerSumSqrt:
        return self

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

solve(vector: Float[Array, ' n']) -> Float[Array, ' n']

Apply the inverse square root S^{-1} to vector.

Parameters:

Name	Type	Description	Default
vector	Float[Array, ' n']	Input vector, shape (n_a · n_b,).	required

Returns:

Type	Description
Float[Array, ' n']	S^{-1} @ vector, shape (n_a · n_b,).

Source code in src/gaussx/_operators/_kronecker_sum.py

def solve(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
    """Apply the inverse square root ``S^{-1}`` to ``vector``.

    Args:
        vector: Input vector, shape ``(n_a · n_b,)``.

    Returns:
        ``S^{-1} @ vector``, shape ``(n_a · n_b,)``.
    """
    Z = rearrange(vector, "(a b) -> 1 a b", a=self._n_a, b=self._n_b)
    result = _kronecker_sum_root_solve(*self._parts(), Z)
    return rearrange(result, "1 a b -> (a b)")

SumOfKroneckers

Bases: AbstractLinearOperator

Sum of Kronecker products Σ_k A_k \otimes B_k.

Appears in multi-output GPs with correlated outputs, e.g. K_task \otimes K_spatial + \sigma^2 I_task \otimes I_spatial.

Matvec is computed as the sum of the Kronecker matvecs.

gaussx.solve and gaussx.logdet dispatch structurally on two terms with one of them positive definite — B ⊗ C + σ² I and the general B₁ ⊗ C₁ + B₂ ⊗ C₂ — by whitening with the positive-definite term and eigendecomposing each factor of the other. That costs O(n_c³ + n_d³) rather than the O((n_c n_d)³) of a dense solve, and the same reduction is reachable through the AddLinearOperator chains sum_operator builds. The factors of the non-anchor term must advertise symmetry, and the anchor's must be diagonal, a multiple of the identity, or tagged positive-semidefinite; anything else keeps the dense fallback.

The anchor must be positive definite

Whitening inverts the anchor term. positive_semidefinite_tag -- the strongest statement lineax can make -- permits a singular operator, and whitening by one of those returns NaN where the dense fallback would have answered. When both terms qualify, the one whose inversion needs no factorization is preferred, so a scalar or diagonal shift always wins over a tagged dense factor; a multiple-of-the-identity anchor is folded in as a shift and is never inverted at all. Beyond that the definiteness is the caller's to guarantee, exactly as it is for gaussx.cholesky.

A diagonal or scaled-identity anchor factor is whitened by its square root, so each entry must be strictly positive. With concrete values a non-positive one is skipped in favour of the other term, or of the dense fallback. A traced one cannot be inspected, so it keeps the diagonal whitening and raises an EquinoxRuntimeError at run time rather than returning NaN.

Three or more terms have no closed form and keep the dense fallback. Drive the structured mv instead with solve(op, b, solver=lineax.CG(...)) and gaussx.SLQLogdet.

eigendecompose is the separate, more general route: a joint eigendecomposition of the second Kronecker pair (requires A_2, B_2 to be symmetric). It forms a dense (n_c n_d) x (n_c n_d) matrix internally, so it is intended for moderate factor sizes (typical for multi-output GPs where the task dimension is small).

Not to be confused with KroneckerSum, which is the standard matrix-theory Kronecker sum A \otimes I + I \otimes B — a different object with a closed-form eigendecomposition. This class is a sum of arbitrary Kronecker products, which has none in general.

Parameters:

Name	Type	Description	Default
kron1	Kronecker	First Kronecker product A_1 \otimes B_1.	required
kron2	Kronecker	Second Kronecker product A_2 \otimes B_2.	required
*krons	Kronecker	Additional two-factor Kronecker products.	()

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)
>>> I2 = lx.IdentityLinearOperator(A.in_structure())
>>> K = gaussx.SumOfKroneckers(gaussx.Kronecker(A, B), gaussx.Kronecker(I2, I2))
>>> dense = jnp.kron(A.as_matrix(), B.as_matrix()) + jnp.eye(4)
>>> bool(jnp.allclose(K.as_matrix(), dense))
True
>>> sign, expected = jnp.linalg.slogdet(dense)
>>> bool(jnp.allclose(gaussx.logdet(K), expected, atol=1e-5))  # per factor
True

Source code in src/gaussx/_operators/_sum_kronecker.py

class SumOfKroneckers(lx.AbstractLinearOperator):
    r"""Sum of Kronecker products ``Σ_k A_k \otimes B_k``.

    Appears in multi-output GPs with correlated outputs, e.g.
    ``K_task \otimes K_spatial + \sigma^2 I_task \otimes I_spatial``.

    Matvec is computed as the sum of the Kronecker matvecs.

    `gaussx.solve` and `gaussx.logdet` dispatch structurally on **two**
    terms with one of them positive definite — ``B ⊗ C + σ² I`` and the
    general ``B₁ ⊗ C₁ + B₂ ⊗ C₂`` — by whitening with the positive-definite
    term and eigendecomposing each factor of the other. That costs
    ``O(n_c³ + n_d³)`` rather than the ``O((n_c n_d)³)`` of a dense solve,
    and the same reduction is reachable through the ``AddLinearOperator``
    chains `sum_operator` builds. The factors of the non-anchor term must
    advertise symmetry, and the anchor's must be diagonal, a multiple of the
    identity, or tagged positive-semidefinite; anything else keeps the dense
    fallback.

    !!! warning "The anchor must be positive *definite*"

        Whitening inverts the anchor term. ``positive_semidefinite_tag`` --
        the strongest statement lineax can make -- permits a singular
        operator, and whitening by one of those returns ``NaN`` where the
        dense fallback would have answered. When both terms qualify, the one
        whose inversion needs no factorization is preferred, so a scalar or
        diagonal shift always wins over a tagged dense factor; a
        multiple-of-the-identity anchor is folded in as a shift and is never
        inverted at all. Beyond that the definiteness is the caller's to
        guarantee, exactly as it is for `gaussx.cholesky`.

        A diagonal or scaled-identity anchor *factor* is whitened by its
        square root, so each entry must be strictly positive. With concrete
        values a non-positive one is skipped in favour of the other term, or
        of the dense fallback. A traced one cannot be inspected, so it keeps
        the diagonal whitening and raises an ``EquinoxRuntimeError`` at run
        time rather than returning ``NaN``.

    Three or more terms have no closed form and keep the dense fallback.
    Drive the structured `mv` instead with
    ``solve(op, b, solver=lineax.CG(...))`` and `gaussx.SLQLogdet`.

    `eigendecompose` is the separate, more general route: a joint
    eigendecomposition of the second Kronecker pair (requires
    ``A_2, B_2`` to be symmetric).  It forms a dense
    ``(n_c n_d) x (n_c n_d)`` matrix internally, so it is
    intended for moderate factor sizes (typical for multi-output GPs
    where the task dimension is small).

    Not to be confused with `KroneckerSum`, which is the standard
    matrix-theory Kronecker *sum* ``A \otimes I + I \otimes B`` — a
    different object with a closed-form eigendecomposition. This class is a
    sum of arbitrary Kronecker *products*, which has none in general.

    Args:
        kron1: First Kronecker product ``A_1 \otimes B_1``.
        kron2: Second Kronecker product ``A_2 \otimes B_2``.
        *krons: Additional two-factor Kronecker products.

    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)
        >>> I2 = lx.IdentityLinearOperator(A.in_structure())
        >>> K = gaussx.SumOfKroneckers(gaussx.Kronecker(A, B), gaussx.Kronecker(I2, I2))
        >>> dense = jnp.kron(A.as_matrix(), B.as_matrix()) + jnp.eye(4)
        >>> bool(jnp.allclose(K.as_matrix(), dense))
        True
        >>> sign, expected = jnp.linalg.slogdet(dense)
        >>> bool(jnp.allclose(gaussx.logdet(K), expected, atol=1e-5))  # per factor
        True
    """

    operators: tuple[Kronecker, ...]
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        kron1: Kronecker,
        kron2: Kronecker,
        *krons: Kronecker,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        operators = (kron1, kron2, *krons)
        if any(len(kron.operators) != 2 for kron in operators):
            raise ValueError("SumOfKroneckers requires two-factor Kronecker products.")
        if any(kron.in_size() != kron1.in_size() for kron in operators[1:]):
            raise ValueError("Kronecker products must have the same size (input size).")
        if any(kron.out_size() != kron1.out_size() for kron in operators[1:]):
            raise ValueError(
                "Kronecker products must have the same size (output size)."
            )
        self.operators = operators
        self._in_size = kron1.in_size()
        self._out_size = kron1.out_size()
        self._dtype = _resolve_dtype(*operators)
        self.tags = _to_frozenset(tags)

    @property
    def kron1(self) -> Kronecker:
        return self.operators[0]

    @property
    def kron2(self) -> Kronecker:
        return self.operators[1]

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " m"]:
        result = self.operators[0].mv(vector)
        for kron in self.operators[1:]:
            result = result + kron.mv(vector)
        return result

    def as_matrix(self) -> Float[Array, "n n"]:
        result = self.operators[0].as_matrix()
        for kron in self.operators[1:]:
            result = result + kron.as_matrix()
        return result

    def transpose(self) -> SumOfKroneckers:
        # ``type(self)`` rather than the literal class so the deprecated
        # `SumKronecker` subclass survives a transpose — external
        # ``isinstance`` checks and ``singledispatch`` registrations keyed on
        # it would otherwise silently stop applying on ``op.T``.
        return type(self)(
            *(
                Kronecker(
                    kron.operators[0].T,
                    kron.operators[1].T,
                    tags=lx.transpose_tags(kron.tags),
                )
                for kron in self.operators
            ),
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

    def eigendecompose(
        self,
    ) -> tuple[Float[Array, " n"], Float[Array, "n n"]]:
        r"""Eigendecompose via joint eigendecomposition of the second pair.

        Decomposes ``A_2 = Q_C \Lambda_C Q_C^T`` and
        ``B_2 = Q_D \Lambda_D Q_D^T``, then transforms the first pair
        into the eigenbasis and diagonalizes the result.

        !!! note

            This forms a dense ``(n_c n_d) x (n_c n_d)`` matrix
            internally and is O((n_c n_d)^3).  Intended for moderate
            factor sizes (e.g. multi-output GPs where task dimension
            is small).

        Raises:
            ValueError: If the factors of ``kron2`` are not symmetric.

        Returns:
            Tuple ``(eigenvalues, Q)`` where
            ``self == Q @ diag(eigenvalues) @ Q^T``.
        """
        if len(self.operators) != 2:
            count = len(self.operators)
            raise ValueError(
                f"eigendecompose requires exactly two Kronecker products, got {count}."
            )
        A2_op, B2_op = self.kron2.operators
        A1_op, B1_op = self.kron1.operators
        # The final ``eigh(transformed)`` call requires ``transformed`` to
        # be symmetric, which in turn requires *both* kron pairs to have
        # symmetric factors (so that ``A1_tilde`` and ``B1_tilde`` stay
        # symmetric under the ``Q^T A Q`` rotation).
        if not lx.is_symmetric(A2_op) or not lx.is_symmetric(B2_op):
            raise ValueError("eigendecompose requires kron2 factors to be symmetric.")
        if not lx.is_symmetric(A1_op) or not lx.is_symmetric(B1_op):
            raise ValueError("eigendecompose requires kron1 factors to be symmetric.")

        A1, B1 = (op.as_matrix() for op in self.kron1.operators)

        # Per-factor eigendecomposition routed through the structural
        # primitive: Diagonal / BlockDiag / nested Kronecker factors of
        # ``A_2`` or ``B_2`` skip materialization here.
        evals_c, Q_C = eig(A2_op)
        evals_d, Q_D = eig(B2_op)

        # Transform first pair into eigenbasis of second pair
        A1_tilde = Q_C.T @ A1 @ Q_C
        B1_tilde = Q_D.T @ B1 @ Q_D

        # kron(A1_tilde, B1_tilde) + diag(evals_c kron evals_d)
        diag_vals = rearrange(
            evals_c[:, None] * evals_d[None, :],
            "a b -> (a b)",
        )
        transformed = jnp.kron(A1_tilde, B1_tilde) + jnp.diag(diag_vals)
        evals, V = jnp.linalg.eigh(transformed)

        Q = jnp.kron(Q_C, Q_D) @ V
        return evals, Q

eigendecompose() -> tuple[Float[Array, ' n'], Float[Array, 'n n']]

Eigendecompose via joint eigendecomposition of the second pair.

Decomposes A_2 = Q_C \Lambda_C Q_C^T and B_2 = Q_D \Lambda_D Q_D^T, then transforms the first pair into the eigenbasis and diagonalizes the result.

Note

This forms a dense (n_c n_d) x (n_c n_d) matrix internally and is O((n_c n_d)^3). Intended for moderate factor sizes (e.g. multi-output GPs where task dimension is small).

Raises:

Type	Description
ValueError	If the factors of kron2 are not symmetric.

Returns:

Type	Description
Float[Array, ' n']	Tuple (eigenvalues, Q) where
Float[Array, 'n n']	self == Q @ diag(eigenvalues) @ Q^T.

Source code in src/gaussx/_operators/_sum_kronecker.py

def eigendecompose(
    self,
) -> tuple[Float[Array, " n"], Float[Array, "n n"]]:
    r"""Eigendecompose via joint eigendecomposition of the second pair.

    Decomposes ``A_2 = Q_C \Lambda_C Q_C^T`` and
    ``B_2 = Q_D \Lambda_D Q_D^T``, then transforms the first pair
    into the eigenbasis and diagonalizes the result.

    !!! note

        This forms a dense ``(n_c n_d) x (n_c n_d)`` matrix
        internally and is O((n_c n_d)^3).  Intended for moderate
        factor sizes (e.g. multi-output GPs where task dimension
        is small).

    Raises:
        ValueError: If the factors of ``kron2`` are not symmetric.

    Returns:
        Tuple ``(eigenvalues, Q)`` where
        ``self == Q @ diag(eigenvalues) @ Q^T``.
    """
    if len(self.operators) != 2:
        count = len(self.operators)
        raise ValueError(
            f"eigendecompose requires exactly two Kronecker products, got {count}."
        )
    A2_op, B2_op = self.kron2.operators
    A1_op, B1_op = self.kron1.operators
    # The final ``eigh(transformed)`` call requires ``transformed`` to
    # be symmetric, which in turn requires *both* kron pairs to have
    # symmetric factors (so that ``A1_tilde`` and ``B1_tilde`` stay
    # symmetric under the ``Q^T A Q`` rotation).
    if not lx.is_symmetric(A2_op) or not lx.is_symmetric(B2_op):
        raise ValueError("eigendecompose requires kron2 factors to be symmetric.")
    if not lx.is_symmetric(A1_op) or not lx.is_symmetric(B1_op):
        raise ValueError("eigendecompose requires kron1 factors to be symmetric.")

    A1, B1 = (op.as_matrix() for op in self.kron1.operators)

    # Per-factor eigendecomposition routed through the structural
    # primitive: Diagonal / BlockDiag / nested Kronecker factors of
    # ``A_2`` or ``B_2`` skip materialization here.
    evals_c, Q_C = eig(A2_op)
    evals_d, Q_D = eig(B2_op)

    # Transform first pair into eigenbasis of second pair
    A1_tilde = Q_C.T @ A1 @ Q_C
    B1_tilde = Q_D.T @ B1 @ Q_D

    # kron(A1_tilde, B1_tilde) + diag(evals_c kron evals_d)
    diag_vals = rearrange(
        evals_c[:, None] * evals_d[None, :],
        "a b -> (a b)",
    )
    transformed = jnp.kron(A1_tilde, B1_tilde) + jnp.diag(diag_vals)
    evals, V = jnp.linalg.eigh(transformed)

    Q = jnp.kron(Q_C, Q_D) @ V
    return evals, Q

SumOfKroneckersSqrt

Bases: SqrtOperator

Lazy Lanczos square-root operator for SumOfKroneckers covariances.

What gaussx.sqrt returns for a SumOfKroneckers. Not to be confused with KroneckerSumSqrt, the closed-form root of a KroneckerSum.

Specialization of SqrtOperator that narrows original to a SumOfKroneckers operator. mv computes sqrt(A) v via matfree's Lanczos matrix-function product without materializing the full square root.

Parameters:

Name	Type	Description	Default
original	SumOfKroneckers	The SumOfKroneckers covariance to take the square root of.	required
lanczos_order	int	Number of Lanczos iterations; clamped to the operator size.	_DEFAULT_LANCZOS_ORDER

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

class SumOfKroneckersSqrt(SqrtOperator):
    """Lazy Lanczos square-root operator for ``SumOfKroneckers`` covariances.

    What `gaussx.sqrt` returns for a `SumOfKroneckers`. Not to be confused
    with `KroneckerSumSqrt`, the closed-form root of a `KroneckerSum`.

    Specialization of `SqrtOperator` that narrows ``original`` to a
    `SumOfKroneckers` operator. `mv` computes ``sqrt(A) v`` via
    matfree's Lanczos matrix-function product without materializing the full
    square root.

    Args:
        original: The ``SumOfKroneckers`` covariance to take the square root of.
        lanczos_order: Number of Lanczos iterations; clamped to the operator
            size.
    """

    original: SumOfKroneckers

    def __init__(
        self,
        original: SumOfKroneckers,
        lanczos_order: int = _DEFAULT_LANCZOS_ORDER,
    ) -> None:
        super().__init__(original, lanczos_order=lanczos_order)

Low-rank updates

\(L + U\,\mathrm{diag}(d)\,V^\top\) with Woodbury-efficient solves and matrix-determinant-lemma logdets. The factories build the common special cases directly from arrays. orthonormal=True is an unchecked claim that \(U\) has orthonormal columns (\(U^\top U = I_k\): truncated SVD, Nyström, ensemble factors). It matters only for a symmetric update on a scaled-identity base, \(cI + U\,\mathrm{diag}(d)\,U^\top\) — build one with low_rank_plus_identity(U, d, scale=c, orthonormal=True) — where solve and logdet take \(O(nk)\) / \(O(k)\) closed forms with no \(k \times k\) factorisation, and inv and sqrt stay orthonormal low-rank updates:

\[ (cI + UDU^\top)^{-1} = c^{-1}\Big(I - U\,\mathrm{diag}\big(\tfrac{d}{c+d}\big)U^\top\Big), \qquad \log\lvert\det(cI + UDU^\top)\rvert = (n-k)\log\lvert c\rvert + \textstyle\sum_i \log\lvert c + d_i\rvert . \]

On any other base (e.g. a general diagonal) the flag does not change the result, and tag inference never depends on it.

SVDLowRankUpdate is deprecated

It remains a LowRankUpdate subclass — so isinstance checks and singledispatch registrations keyed on it still work — but it forces orthonormal=True and emits a DeprecationWarning on construction. New code should use LowRankUpdate(base, U, S, V, orthonormal=True) or svd_low_rank_plus_diag. It will be removed in gaussx 0.7.0.

Structured linear algebra and Gaussian primitives for JAX.

LowRankUpdate

Bases: AbstractLinearOperator

Low-rank update operator L + U diag(d) Vᵀ.

Represents a base operator L plus a rank-k update. When L is cheap to solve (e.g. diagonal), the Woodbury identity gives efficient solves for the full operator.

Parameters:

Name	Type	Description	Default
base	AbstractLinearOperator	The base operator L, with shape (m, n).	required
U	Float[Array, 'n k']	Left factor, shape (m, k).	required
d	Float[Array, ' k'] | None	Diagonal scaling, shape (k,). Defaults to ones.	None
V	Float[Array, 'n k'] | None	Right factor, shape (n, k). Defaults to U for square operators, yielding the symmetric update L + U diag(d) Uᵀ.	None
tags	object | frozenset[object]	Extra lineax tags -- the caller's structural claims.	frozenset()
orthonormal	bool	The caller's claim that U (and V) have orthonormal columns, UᵀU = I_k. Like a tag, it is never checked. It takes effect only for a symmetric update (V omitted or U itself) on a scaled-identity base c·I (a lineax IdentityLinearOperator, possibly scaled, negated or tagged -- low_rank_plus_identity(..., orthonormal=True) builds one): gaussx.solve and gaussx.logdet then take O(nk) / O(k) closed forms with no k x k factorisation, and gaussx.inv / gaussx.sqrt return another orthonormal LowRankUpdate on a scaled-identity base. On any other base (e.g. a general diagonal, where orthonormality does not simplify Woodbury) it has no effect on the result.	False

Tags are inferred from structure only, never from array values, so the same call gives the same tags (and pytree structure) eagerly and under jax.jit:

• symmetric_tag when the base is symmetric and the factors are shared -- V omitted or passed as the same object as U. This is recorded in the static symmetric_factors field. Value-equal but distinct factors (V = U.copy()) are not inferred symmetric; pass tags=lx.symmetric_tag to claim it.
• positive_semidefinite_tag additionally when the base is PSD and d is omitted (all ones). A caller-supplied d can have any sign, so pass tags=lx.positive_semidefinite_tag to claim PSD.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> base = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0, 3.0]))
>>> U = jnp.array([[1.0], [0.0], [1.0]])
>>> K = gaussx.LowRankUpdate(base, U)  # base + U Uᵀ
>>> K.rank
1
>>> dense = K.as_matrix()
>>> b = jnp.ones(3)
>>> x = gaussx.solve(K, b)  # Woodbury: one 1 x 1 solve, no 3 x 3 factorisation
>>> bool(jnp.allclose(dense @ x, b, atol=1e-5))
True

Source code in src/gaussx/_operators/_low_rank_update.py

class LowRankUpdate(lx.AbstractLinearOperator):
    """Low-rank update operator ``L + U diag(d) Vᵀ``.

    Represents a base operator *L* plus a rank-k update. When *L*
    is cheap to solve (e.g. diagonal), the Woodbury identity gives
    efficient solves for the full operator.

    Args:
        base: The base operator *L*, with shape ``(m, n)``.
        U: Left factor, shape ``(m, k)``.
        d: Diagonal scaling, shape ``(k,)``. Defaults to ones.
        V: Right factor, shape ``(n, k)``. Defaults to *U* for
            square operators, yielding the symmetric update
            ``L + U diag(d) Uᵀ``.
        tags: Extra lineax tags -- the caller's structural claims.
        orthonormal: The caller's claim that *U* (and *V*) have orthonormal
            columns, ``UᵀU = I_k``. Like a tag, it is never checked. It
            takes effect only for a symmetric update (*V* omitted or *U*
            itself) on a scaled-identity base ``c·I`` (a lineax
            ``IdentityLinearOperator``, possibly scaled, negated or
            tagged -- ``low_rank_plus_identity(..., orthonormal=True)``
            builds one): `gaussx.solve` and `gaussx.logdet` then take
            ``O(nk)`` / ``O(k)`` closed forms with no ``k x k``
            factorisation, and `gaussx.inv` / `gaussx.sqrt` return another
            orthonormal ``LowRankUpdate`` on a scaled-identity base. On any
            other base (e.g. a general diagonal, where orthonormality does
            not simplify Woodbury) it has no effect on the result.

    Tags are inferred from structure only, never from array values, so the
    same call gives the same tags (and pytree structure) eagerly and under
    ``jax.jit``:

    - ``symmetric_tag`` when the base is symmetric and the factors are
      shared -- *V* omitted or passed as the same object as *U*. This is
      recorded in the static ``symmetric_factors`` field. Value-equal but
      distinct factors (``V = U.copy()``) are not inferred symmetric; pass
      ``tags=lx.symmetric_tag`` to claim it.
    - ``positive_semidefinite_tag`` additionally when the base is PSD and
      *d* is omitted (all ones). A caller-supplied *d* can have any sign,
      so pass ``tags=lx.positive_semidefinite_tag`` to claim PSD.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> base = lx.DiagonalLinearOperator(jnp.array([1.0, 2.0, 3.0]))
        >>> U = jnp.array([[1.0], [0.0], [1.0]])
        >>> K = gaussx.LowRankUpdate(base, U)  # base + U Uᵀ
        >>> K.rank
        1
        >>> dense = K.as_matrix()
        >>> b = jnp.ones(3)
        >>> x = gaussx.solve(K, b)  # Woodbury: one 1 x 1 solve, no 3 x 3 factorisation
        >>> bool(jnp.allclose(dense @ x, b, atol=1e-5))
        True
    """

    base: lx.AbstractLinearOperator
    U: Float[Array, "m k"]
    d: Float[Array, " k"]
    V: Float[Array, "n k"]
    orthonormal: bool = eqx.field(static=True)
    symmetric_factors: bool = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        base: lx.AbstractLinearOperator,
        U: Float[Array, "n k"],
        d: Float[Array, " k"] | None = None,
        V: Float[Array, "n k"] | None = None,
        *,
        tags: object | frozenset[object] = frozenset(),
        orthonormal: bool = False,
    ) -> None:
        m = base.out_size()
        n = base.in_size()
        k = U.shape[1] if U.ndim == 2 else 1
        if U.ndim == 1:
            U = U[:, None]
        unit_weights = d is None
        if d is None:
            d = jnp.ones(k, dtype=U.dtype)
        symmetric_factors = V is None or V is U
        if V is None:
            V = U
        if V.ndim == 1:
            V = V[:, None]
        if U.shape[0] != m or V.shape[0] != n:
            raise ValueError(
                f"U must have {m} rows and V must have {n} rows to match "
                f"base operator, "
                f"got U.shape={U.shape}, V.shape={V.shape}."
            )
        if U.shape[1] != d.shape[0] or V.shape[1] != d.shape[0]:
            raise ValueError(
                f"Rank dimensions must match: U has {U.shape[1]} cols, "
                f"V has {V.shape[1]} cols, d has {d.shape[0]} entries."
            )
        self.base = base
        self.U = U
        self.d = d
        self.V = V
        self.orthonormal = orthonormal
        self.symmetric_factors = symmetric_factors and m == n

        inferred_tags = _infer_tags(
            base,
            symmetric_factors=self.symmetric_factors,
            unit_weights=unit_weights,
        )
        self.tags = _to_frozenset(tags) | inferred_tags | {low_rank_tag}

    @property
    def rank(self) -> int:
        """Rank of the low-rank update."""
        return self.d.shape[0]

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " m"]:
        # (L + U diag(d) V^T) x = L x + U (d * (V^T x))
        base_part = self.base.mv(vector)
        vtx = self.V.T @ vector  # (k,)
        scaled = self.d * vtx  # (k,)
        update_part = self.U @ scaled  # (m,)
        return base_part + update_part

    def as_matrix(self) -> Float[Array, "m n"]:
        L = self.base.as_matrix()
        return L + self.U @ jnp.diag(self.d) @ self.V.T

    def transpose(self) -> LowRankUpdate:
        # Shared factors are passed once so the transpose keeps them shared
        # even under tracing, where ``U`` and ``V`` are distinct tracers.
        return LowRankUpdate(
            self.base.T,
            self.V,
            self.d,
            None if self.symmetric_factors else self.U,
            tags=lx.transpose_tags(self.tags),
            orthonormal=self.orthonormal,
        )

    def _dtype(self):
        # ``mv`` returns the promoted dtype of base, U, d and V; the declared
        # structure must agree, since lineax compares the two in solves.
        return jnp.result_type(self.base.in_structure().dtype, self.U, self.d, self.V)

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct(self.base.in_structure().shape, self._dtype())

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct(self.base.out_structure().shape, self._dtype())

rank: int property

Rank of the low-rank update.

low_rank_plus_diag(diag: Float[Array, ' n'], U: Float[Array, 'n k'], d: Float[Array, ' k'] | None = None, V: Float[Array, 'n k'] | None = None, *, psd: bool = False) -> LowRankUpdate

Construct diag(diag) + U diag(d) Vᵀ.

Common pattern for inducing-point / Nystrom approximations where the base is a diagonal matrix.

Parameters:

Name	Type	Description	Default
diag	Float[Array, ' n']	Diagonal entries, shape (n,).	required
U	Float[Array, 'n k']	Left factor, shape (n, k).	required
d	Float[Array, ' k'] | None	Diagonal scaling, shape (k,). Defaults to ones.	None
V	Float[Array, 'n k'] | None	Right factor, shape (n, k). Defaults to U.	None
psd	bool	Claim that diag and d are non-negative (and the factors shared), so the base and the update are tagged positive semidefinite. The sign of diag is never inspected, so without this the result is only symmetric-tagged.	False

Returns:

Type	Description
LowRankUpdate	A LowRankUpdate with a DiagonalLinearOperator base.

Source code in src/gaussx/_operators/_low_rank_update.py

def low_rank_plus_diag(
    diag: Float[Array, " n"],
    U: Float[Array, "n k"],
    d: Float[Array, " k"] | None = None,
    V: Float[Array, "n k"] | None = None,
    *,
    psd: bool = False,
) -> LowRankUpdate:
    """Construct ``diag(diag) + U diag(d) Vᵀ``.

    Common pattern for inducing-point / Nystrom approximations
    where the base is a diagonal matrix.

    Args:
        diag: Diagonal entries, shape ``(n,)``.
        U: Left factor, shape ``(n, k)``.
        d: Diagonal scaling, shape ``(k,)``. Defaults to ones.
        V: Right factor, shape ``(n, k)``. Defaults to *U*.
        psd: Claim that ``diag`` and ``d`` are non-negative (and the
            factors shared), so the base and the update are tagged
            positive semidefinite. The sign of ``diag`` is never inspected,
            so without this the result is only symmetric-tagged.

    Returns:
        A ``LowRankUpdate`` with a ``DiagonalLinearOperator`` base.
    """
    return _low_rank_update_with_diag_base(diag, U, d, V, psd=psd)

low_rank_plus_identity(U: Float[Array, 'n k'], d: Float[Array, ' k'] | None = None, V: Float[Array, 'n k'] | None = None, *, scale: float = 1.0, psd: bool = False, orthonormal: bool = False) -> LowRankUpdate

Construct scale * I + U diag(d) Vᵀ.

Common pattern for regularised low-rank models (e.g. noise + signal).

Parameters:

Name	Type	Description	Default
U	Float[Array, 'n k']	Left factor, shape (n, k).	required
d	Float[Array, ' k'] | None	Diagonal scaling, shape (k,). Defaults to ones.	None
V	Float[Array, 'n k'] | None	Right factor, shape (n, k). Defaults to U.	None
scale	float	Scalar multiplier on the identity. Default 1.0.	1.0
psd	bool	Claim that scale and d are non-negative (and the factors shared), so the result is tagged positive semidefinite. A Python-number scale >= 0 already makes the base PSD (it is static, not an array value), so the default low_rank_plus_identity(U) is PSD-tagged without it.	False
orthonormal	bool	Claim that U has orthonormal columns (unchecked). The base is then a scaled lineax IdentityLinearOperator rather than a DiagonalLinearOperator, so a symmetric update (V omitted) takes the closed-form orthonormal solve / logdet / inv / sqrt (see LowRankUpdate).	False

Returns:

Type	Description
LowRankUpdate	A LowRankUpdate with a scaled identity base.

Source code in src/gaussx/_operators/_low_rank_update.py

def low_rank_plus_identity(
    U: Float[Array, "n k"],
    d: Float[Array, " k"] | None = None,
    V: Float[Array, "n k"] | None = None,
    *,
    scale: float = 1.0,
    psd: bool = False,
    orthonormal: bool = False,
) -> LowRankUpdate:
    """Construct ``scale * I + U diag(d) Vᵀ``.

    Common pattern for regularised low-rank models (e.g. noise + signal).

    Args:
        U: Left factor, shape ``(n, k)``.
        d: Diagonal scaling, shape ``(k,)``. Defaults to ones.
        V: Right factor, shape ``(n, k)``. Defaults to *U*.
        scale: Scalar multiplier on the identity. Default 1.0.
        psd: Claim that ``scale`` and ``d`` are non-negative (and the
            factors shared), so the result is tagged positive semidefinite.
            A Python-number ``scale >= 0`` already makes the base PSD
            (it is static, not an array value), so the default
            ``low_rank_plus_identity(U)`` is PSD-tagged without it.
        orthonormal: Claim that *U* has orthonormal columns (unchecked).
            The base is then a scaled lineax ``IdentityLinearOperator``
            rather than a ``DiagonalLinearOperator``, so a symmetric update
            (*V* omitted) takes the closed-form orthonormal ``solve`` /
            ``logdet`` / ``inv`` / ``sqrt`` (see `LowRankUpdate`).

    Returns:
        A ``LowRankUpdate`` with a scaled identity base.
    """
    n = U.shape[0]
    static_nonnegative = isinstance(scale, (int, float)) and scale >= 0
    if orthonormal:
        identity = lx.IdentityLinearOperator(jax.ShapeDtypeStruct((n,), U.dtype))
        base: lx.AbstractLinearOperator = jnp.asarray(scale, dtype=U.dtype) * identity
        if psd or static_nonnegative:
            base = lx.TaggedLinearOperator(base, lx.positive_semidefinite_tag)
        tags = (
            frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
            if psd
            else frozenset()
        )
        return LowRankUpdate(base, U, d, V, tags=tags, orthonormal=True)
    diag = jnp.full(n, scale, dtype=U.dtype)
    return _low_rank_update_with_diag_base(
        diag, U, d, V, psd=psd, psd_base=static_nonnegative
    )

svd_low_rank_plus_diag(diag: Float[Array, ' n'], U: Float[Array, 'n k'], S: Float[Array, ' k'], V: Float[Array, 'n k'], *, psd: bool = False) -> LowRankUpdate

Construct diag(diag) + U diag(S) Vᵀ from a truncated SVD.

Parameters:

Name	Type	Description	Default
diag	Float[Array, ' n']	Diagonal entries, shape (n,).	required
U	Float[Array, 'n k']	Left singular vectors, shape (n, k).	required
S	Float[Array, ' k']	Singular values, shape (k,).	required
V	Float[Array, 'n k']	Right singular vectors, shape (n, k). Pass the same array object as U for a symmetric update; value-equal copies are not inferred symmetric.	required
psd	bool	Claim that diag is non-negative and the update is symmetric, so the result is tagged positive semidefinite.	False

Returns:

Type	Description
LowRankUpdate	A LowRankUpdate with a DiagonalLinearOperator base.

Source code in src/gaussx/_operators/_low_rank_update.py

def svd_low_rank_plus_diag(
    diag: Float[Array, " n"],
    U: Float[Array, "n k"],
    S: Float[Array, " k"],
    V: Float[Array, "n k"],
    *,
    psd: bool = False,
) -> LowRankUpdate:
    """Construct ``diag(diag) + U diag(S) Vᵀ`` from a truncated SVD.

    Args:
        diag: Diagonal entries, shape ``(n,)``.
        U: Left singular vectors, shape ``(n, k)``.
        S: Singular values, shape ``(k,)``.
        V: Right singular vectors, shape ``(n, k)``. Pass the same array
            object as *U* for a symmetric update; value-equal copies are
            not inferred symmetric.
        psd: Claim that ``diag`` is non-negative and the update is
            symmetric, so the result is tagged positive semidefinite.

    Returns:
        A ``LowRankUpdate`` with a ``DiagonalLinearOperator`` base.
    """
    return _low_rank_update_with_diag_base(diag, U, S, V, orthonormal=True, psd=psd)

Banded & Toeplitz

Block-tridiagonal operators solve in \(O(N d^3)\) via block-banded Cholesky — the precision structure of Markovian (state-space) GPs. Symmetric Toeplitz operators get \(O(n \log n)\) matvecs and sampling via FFT circulant embedding.

Structured linear algebra and Gaussian primitives for JAX.

BlockTriDiag

Bases: AbstractLinearOperator

Block-tridiagonal operator with a symmetric off-diagonal band.

Represents the structure:

[D_1  A_1^T              ]
[A_1  D_2   A_2^T        ]
[     A_2   D_3   ...    ]
[               A_{N-1} D_N]

where D_k are (d, d) diagonal blocks and A_k are (d, d) sub-diagonal blocks. This is the precision matrix structure arising from discretized SDEs in state-space GP inference.

All primitives (solve, logdet, cholesky, diag, trace) exploit the banded structure for O(Nd³) cost instead of O((Nd)³).

The off-diagonal band is symmetric by construction (A_k below, A_kᵀ above), so the operator is symmetric exactly when every D_k is. With symmetric=True (the default) it carries lineax.symmetric_tag and gaussx.solve / gaussx.logdet use the block Cholesky, which reads only the lower triangle of each D_k; a concrete non-symmetric D_k is rejected with a ValueError (the check is skipped under tracing, where the flag is a claim). Pass symmetric=False for non-symmetric diagonal blocks: the tag is then omitted and the solve and log-determinant fall back to a dense LU.

+, -, unary - and scalar * keep the tags that still hold: symmetry when both operands are symmetric, positive-semidefiniteness for a sum of two PSD operators or a concretely non-negative scale (a negative one, or negation, swaps it with negative-semidefiniteness).

Parameters:

Name	Type	Description	Default
diagonal	Float[Array, 'N d d']	Diagonal blocks, shape (N, d, d).	required
sub_diagonal	Float[Array, 'Nm1 d d']	Sub-diagonal blocks, shape (N-1, d, d).	required
symmetric	bool	Whether every D_k is symmetric (see above).	True
tags	object | frozenset[object]	Additional lineax tags.	frozenset()

Raises:

Type	Description
ValueError	If symmetric=True and a concrete D_k is not symmetric to within sqrt(eps) of its largest entry.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> diagonal = jnp.stack([3.0 * jnp.eye(2)] * 3)  # three 2 x 2 blocks
>>> sub_diagonal = jnp.stack([0.5 * jnp.eye(2)] * 2)
>>> T = gaussx.BlockTriDiag(diagonal, sub_diagonal)
>>> T.as_matrix().shape
(6, 6)
>>> b = jnp.ones(6)
>>> x = gaussx.solve(T, b)  # block-banded, O(N d³)
>>> bool(jnp.allclose(T.as_matrix() @ x, b, atol=1e-5))
True

Source code in src/gaussx/_operators/_block_tridiag.py

class BlockTriDiag(lx.AbstractLinearOperator):
    r"""Block-tridiagonal operator with a symmetric off-diagonal band.

    Represents the structure:

        [D_1  A_1^T              ]
        [A_1  D_2   A_2^T        ]
        [     A_2   D_3   ...    ]
        [               A_{N-1} D_N]

    where ``D_k`` are ``(d, d)`` diagonal blocks and ``A_k`` are
    ``(d, d)`` sub-diagonal blocks. This is the precision matrix
    structure arising from discretized SDEs in state-space GP inference.

    All primitives (solve, logdet, cholesky, diag, trace) exploit the
    banded structure for O(Nd³) cost instead of O((Nd)³).

    The off-diagonal band is symmetric by construction (``A_k`` below,
    ``A_kᵀ`` above), so the operator is symmetric exactly when every ``D_k``
    is. With ``symmetric=True`` (the default) it carries
    ``lineax.symmetric_tag`` and `gaussx.solve` / `gaussx.logdet` use the
    block Cholesky, which reads only the lower triangle of each ``D_k``; a
    concrete non-symmetric ``D_k`` is rejected with a ``ValueError`` (the
    check is skipped under tracing, where the flag is a claim). Pass
    ``symmetric=False`` for non-symmetric diagonal blocks: the tag is then
    omitted and the solve and log-determinant fall back to a dense LU.

    ``+``, ``-``, unary ``-`` and scalar ``*`` keep the tags that still hold:
    symmetry when both operands are symmetric, positive-semidefiniteness for
    a sum of two PSD operators or a concretely non-negative scale (a negative
    one, or negation, swaps it with negative-semidefiniteness).

    Args:
        diagonal: Diagonal blocks, shape ``(N, d, d)``.
        sub_diagonal: Sub-diagonal blocks, shape ``(N-1, d, d)``.
        symmetric: Whether every ``D_k`` is symmetric (see above).
        tags: Additional lineax tags.

    Raises:
        ValueError: If ``symmetric=True`` and a concrete ``D_k`` is not
            symmetric to within ``sqrt(eps)`` of its largest entry.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> diagonal = jnp.stack([3.0 * jnp.eye(2)] * 3)  # three 2 x 2 blocks
        >>> sub_diagonal = jnp.stack([0.5 * jnp.eye(2)] * 2)
        >>> T = gaussx.BlockTriDiag(diagonal, sub_diagonal)
        >>> T.as_matrix().shape
        (6, 6)
        >>> b = jnp.ones(6)
        >>> x = gaussx.solve(T, b)  # block-banded, O(N d³)
        >>> bool(jnp.allclose(T.as_matrix() @ x, b, atol=1e-5))
        True
    """

    diagonal: Float[Array, "N d d"]
    sub_diagonal: Float[Array, "Nm1 d d"]
    symmetric: bool = eqx.field(static=True)
    _num_blocks: int = eqx.field(static=True)
    _block_size: int = eqx.field(static=True)
    _size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        diagonal: Float[Array, "N d d"],
        sub_diagonal: Float[Array, "Nm1 d d"],
        *,
        symmetric: bool = True,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        if diagonal.ndim != 3:
            raise ValueError(
                f"diagonal must have 3 dimensions (N, d, d), got {diagonal.ndim}."
            )
        if sub_diagonal.ndim != 3:
            raise ValueError(
                f"sub_diagonal must have 3 dimensions (N-1, d, d), "
                f"got {sub_diagonal.ndim}."
            )
        N, d, d2 = diagonal.shape
        if d != d2:
            raise ValueError(f"Diagonal blocks must be square, got ({d}, {d2}).")
        if sub_diagonal.shape[0] != N - 1:
            raise ValueError(
                f"sub_diagonal must have {N - 1} blocks, got {sub_diagonal.shape[0]}."
            )
        if sub_diagonal.shape[1] != d or sub_diagonal.shape[2] != d:
            raise ValueError(
                f"Sub-diagonal blocks must have shape ({d}, {d}), "
                f"got ({sub_diagonal.shape[1]}, {sub_diagonal.shape[2]})."
            )
        if symmetric and not _is_concretely_symmetric_or_traced(diagonal):
            raise ValueError(
                "BlockTriDiag diagonal blocks are not symmetric; pass "
                "symmetric=False for a non-symmetric operator (gh-344)."
            )
        self.diagonal = diagonal
        self.sub_diagonal = sub_diagonal
        self.symmetric = symmetric
        self._num_blocks = N
        self._block_size = d
        self._size = N * d
        self._dtype = str(diagonal.dtype)
        structural: set[object] = {block_tridiagonal_tag}
        if symmetric:
            structural.add(lx.symmetric_tag)
        self.tags = _to_frozenset(tags) | structural

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        N = self._num_blocks
        d = self._block_size
        x = rearrange(vector, "(N d) -> N d", N=N, d=d)
        # Dₖ xₖ for all k
        result = einsum(self.diagonal, x, "N d1 d2, N d2 -> N d1")
        if N == 1:
            # einx cannot contract the empty (0, d, d) sub-diagonal (gh-304).
            return rearrange(result, "N d -> (N d)")
        # Aₖ xₖ₋₁ for k = 1, ..., N-1 (sub-diagonal)
        sub_contrib = einsum(self.sub_diagonal, x[:-1], "N d1 d2, N d2 -> N d1")
        result = result.at[1:].add(sub_contrib)
        # Aₖᵀ xₖ₊₁ for k = 0, ..., N-2 (super-diagonal)
        super_contrib = einsum(self.sub_diagonal, x[1:], "N d1 d2, N d1 -> N d2")
        result = result.at[:-1].add(super_contrib)
        return rearrange(result, "N d -> (N d)")

    def as_matrix(self) -> Float[Array, "n n"]:
        N = self._num_blocks
        d = self._block_size
        n = self._size
        mat = jnp.zeros((n, n), dtype=jnp.dtype(self._dtype))
        for k in range(N):
            r = k * d
            mat = mat.at[r : r + d, r : r + d].set(self.diagonal[k])
        for k in range(N - 1):
            r = (k + 1) * d
            c = k * d
            mat = mat.at[r : r + d, c : c + d].set(self.sub_diagonal[k])
            mat = mat.at[c : c + d, r : r + d].set(self.sub_diagonal[k].T)
        return mat

    def transpose(self) -> BlockTriDiag:
        # as_matrix puts sub[k] at (k+1,k) and sub[k].T at (k,k+1).
        # Transposing: new (k+1,k) = old (k,k+1).T = sub[k].
        # So new sub_diagonal = self.sub_diagonal (unchanged).
        return BlockTriDiag(
            rearrange(self.diagonal, "N i j -> N j i"),
            self.sub_diagonal,
            symmetric=self.symmetric,
            tags=lx.transpose_tags(self._extra_tags()),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def _extra_tags(self) -> frozenset[object]:
        """Tags other than the structural ones the constructor adds."""
        return self.tags - {block_tridiagonal_tag, lx.symmetric_tag}

    def _definiteness(self) -> frozenset[object]:
        return self.tags & frozenset(_DEFINITENESS_SWAP)

    def add(self, other: BlockTriDiag) -> BlockTriDiag:
        """Add two block-tridiagonal operators (e.g. prior + likelihood sites).

        The sum is symmetric if both operands are, and keeps a definiteness
        tag (PSD or NSD) that both operands carry.
        """
        return BlockTriDiag(
            self.diagonal + other.diagonal,
            self.sub_diagonal + other.sub_diagonal,
            symmetric=self.symmetric and other.symmetric,
            tags=self._definiteness() & other._definiteness(),
        )

    def __add__(self, other):
        if not isinstance(other, BlockTriDiag):
            # A foreign operand (e.g. ``s * Identity``) becomes a lineax sum
            # instead of crashing on ``other.diagonal`` (gh-337).
            return super().__add__(other)
        return self.add(other)

    def __radd__(self, other: object) -> BlockTriDiag:
        if isinstance(other, BlockTriDiag):
            return other.add(self)
        if other == 0:
            return self
        return NotImplemented

    def __sub__(self, other):
        if not isinstance(other, BlockTriDiag):
            return super().__sub__(other)
        return BlockTriDiag(
            self.diagonal - other.diagonal,
            self.sub_diagonal - other.sub_diagonal,
            symmetric=self.symmetric and other.symmetric,
        )

    def __neg__(self) -> BlockTriDiag:
        return BlockTriDiag(
            -self.diagonal,
            -self.sub_diagonal,
            symmetric=self.symmetric,
            tags=frozenset(_DEFINITENESS_SWAP[tag] for tag in self._definiteness()),
        )

    def __mul__(self, other: object) -> BlockTriDiag:
        scalar = jnp.asarray(other)
        if scalar.ndim != 0:
            msg = "BlockTriDiag can only be multiplied by a scalar"
            raise TypeError(msg)
        # A definiteness tag survives a scale of known sign only.
        tags: frozenset[object] = frozenset()
        if not isinstance(scalar, Tracer):
            value = float(np.real(np.asarray(scalar)))
            if value >= 0:
                tags = self._definiteness()
            else:
                tags = frozenset(_DEFINITENESS_SWAP[t] for t in self._definiteness())
        return BlockTriDiag(
            scalar * self.diagonal,
            scalar * self.sub_diagonal,
            symmetric=self.symmetric,
            tags=tags,
        )

    def __rmul__(self, other: object) -> BlockTriDiag:
        return self.__mul__(other)

add(other: BlockTriDiag) -> BlockTriDiag

Add two block-tridiagonal operators (e.g. prior + likelihood sites).

The sum is symmetric if both operands are, and keeps a definiteness tag (PSD or NSD) that both operands carry.

Source code in src/gaussx/_operators/_block_tridiag.py

def add(self, other: BlockTriDiag) -> BlockTriDiag:
    """Add two block-tridiagonal operators (e.g. prior + likelihood sites).

    The sum is symmetric if both operands are, and keeps a definiteness
    tag (PSD or NSD) that both operands carry.
    """
    return BlockTriDiag(
        self.diagonal + other.diagonal,
        self.sub_diagonal + other.sub_diagonal,
        symmetric=self.symmetric and other.symmetric,
        tags=self._definiteness() & other._definiteness(),
    )

LowerBlockTriDiag

Bases: AbstractLinearOperator

Lower triangular block-bidiagonal Cholesky factor.

Represents:

[L_1              ]
[B_1  L_2          ]
[     B_2  L_3     ]
[          ...  L_N]

where L_k are (d, d) lower-triangular blocks and B_k are (d, d) sub-diagonal blocks.

Parameters:

Name	Type	Description	Default
diagonal	Float[Array, 'N d d']	Lower-triangular diagonal blocks, shape (N, d, d).	required
sub_diagonal	Float[Array, 'Nm1 d d']	Sub-diagonal blocks, shape (N-1, d, d).	required

Source code in src/gaussx/_operators/_block_tridiag.py

class LowerBlockTriDiag(lx.AbstractLinearOperator):
    """Lower triangular block-bidiagonal Cholesky factor.

    Represents:

        [L_1              ]
        [B_1  L_2          ]
        [     B_2  L_3     ]
        [          ...  L_N]

    where ``L_k`` are ``(d, d)`` lower-triangular blocks and ``B_k`` are
    ``(d, d)`` sub-diagonal blocks.

    Args:
        diagonal: Lower-triangular diagonal blocks, shape ``(N, d, d)``.
        sub_diagonal: Sub-diagonal blocks, shape ``(N-1, d, d)``.
    """

    diagonal: Float[Array, "N d d"]
    sub_diagonal: Float[Array, "Nm1 d d"]
    _num_blocks: int = eqx.field(static=True)
    _block_size: int = eqx.field(static=True)
    _size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        diagonal: Float[Array, "N d d"],
        sub_diagonal: Float[Array, "Nm1 d d"],
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        N, d, _ = diagonal.shape
        self.diagonal = diagonal
        self.sub_diagonal = sub_diagonal
        self._num_blocks = N
        self._block_size = d
        self._size = N * d
        self._dtype = str(diagonal.dtype)
        self.tags = _to_frozenset(tags) | {lx.lower_triangular_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        N = self._num_blocks
        d = self._block_size
        x = rearrange(vector, "(N d) -> N d", N=N, d=d)
        # Lₖ xₖ
        result = einsum(self.diagonal, x, "N d1 d2, N d2 -> N d1")
        if N == 1:
            return rearrange(result, "N d -> (N d)")
        # Bₖ xₖ₋₁
        sub_contrib = einsum(self.sub_diagonal, x[:-1], "N d1 d2, N d2 -> N d1")
        result = result.at[1:].add(sub_contrib)
        return rearrange(result, "N d -> (N d)")

    def as_matrix(self) -> Float[Array, "n n"]:
        N = self._num_blocks
        d = self._block_size
        n = self._size
        mat = jnp.zeros((n, n), dtype=jnp.dtype(self._dtype))
        for k in range(N):
            r = k * d
            mat = mat.at[r : r + d, r : r + d].set(self.diagonal[k])
        for k in range(N - 1):
            r = (k + 1) * d
            c = k * d
            mat = mat.at[r : r + d, c : c + d].set(self.sub_diagonal[k])
        return mat

    def transpose(self) -> UpperBlockTriDiag:
        """Transpose gives upper block-bidiagonal."""
        return UpperBlockTriDiag(
            rearrange(self.diagonal, "N i j -> N j i"),
            _transpose_blocks(self.sub_diagonal),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

transpose() -> UpperBlockTriDiag

Transpose gives upper block-bidiagonal.

Source code in src/gaussx/_operators/_block_tridiag.py

def transpose(self) -> UpperBlockTriDiag:
    """Transpose gives upper block-bidiagonal."""
    return UpperBlockTriDiag(
        rearrange(self.diagonal, "N i j -> N j i"),
        _transpose_blocks(self.sub_diagonal),
    )

UpperBlockTriDiag

Bases: AbstractLinearOperator

Upper triangular block-bidiagonal (transpose of LowerBlockTriDiag).

Represents:

[U_1  C_1            ]
[     U_2  C_2        ]
[          ...   C_{N-1}]
[               U_N  ]

where U_k are upper-triangular diagonal blocks and C_k are super-diagonal blocks.

Source code in src/gaussx/_operators/_block_tridiag.py

class UpperBlockTriDiag(lx.AbstractLinearOperator):
    """Upper triangular block-bidiagonal (transpose of LowerBlockTriDiag).

    Represents:

        [U_1  C_1            ]
        [     U_2  C_2        ]
        [          ...   C_{N-1}]
        [               U_N  ]

    where ``U_k`` are upper-triangular diagonal blocks and ``C_k`` are
    super-diagonal blocks.
    """

    diagonal: Float[Array, "N d d"]
    super_diagonal: Float[Array, "Nm1 d d"]
    _num_blocks: int = eqx.field(static=True)
    _block_size: int = eqx.field(static=True)
    _size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        diagonal: Float[Array, "N d d"],
        super_diagonal: Float[Array, "Nm1 d d"],
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        N, d, _ = diagonal.shape
        self.diagonal = diagonal
        self.super_diagonal = super_diagonal
        self._num_blocks = N
        self._block_size = d
        self._size = N * d
        self._dtype = str(diagonal.dtype)
        self.tags = _to_frozenset(tags) | {lx.upper_triangular_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        N = self._num_blocks
        d = self._block_size
        x = rearrange(vector, "(N d) -> N d", N=N, d=d)
        # Uₖ xₖ
        result = einsum(self.diagonal, x, "N d1 d2, N d2 -> N d1")
        if N == 1:
            return rearrange(result, "N d -> (N d)")
        # Cₖ xₖ₊₁
        super_contrib = einsum(self.super_diagonal, x[1:], "N d1 d2, N d2 -> N d1")
        result = result.at[:-1].add(super_contrib)
        return rearrange(result, "N d -> (N d)")

    def as_matrix(self) -> Float[Array, "n n"]:
        N = self._num_blocks
        d = self._block_size
        n = self._size
        mat = jnp.zeros((n, n), dtype=jnp.dtype(self._dtype))
        for k in range(N):
            r = k * d
            mat = mat.at[r : r + d, r : r + d].set(self.diagonal[k])
        for k in range(N - 1):
            r = k * d
            c = (k + 1) * d
            mat = mat.at[r : r + d, c : c + d].set(self.super_diagonal[k])
        return mat

    def transpose(self) -> LowerBlockTriDiag:
        return LowerBlockTriDiag(
            rearrange(self.diagonal, "N i j -> N j i"),
            _transpose_blocks(self.super_diagonal),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

Toeplitz

Bases: AbstractLinearOperator

Symmetric Toeplitz matrix from its first column.

K_{ij} = c_{|i-j|}. Stored as O(n) with O(n log n) matvec via circulant embedding and FFT.

For stationary kernels on regular 1-D grids the full kernel matrix is Toeplitz, so this gives an asymptotic win over dense storage.

Parameters:

Name	Type	Description	Default
column	Float[Array, ' n']	First column of the Toeplitz matrix, shape (n,). Real; integer columns are promoted to floating point.	required

Raises:

Type	Description
TypeError	If column is complex.
ValueError	If column is not rank 1 or is empty.

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> import gaussx
>>> T = gaussx.Toeplitz(jnp.array([2.0, 0.5, 0.25]))  # first column
>>> T.as_matrix().tolist()
[[2.0, 0.5, 0.25], [0.5, 2.0, 0.5], [0.25, 0.5, 2.0]]
>>> [float(v) for v in T.mv(jnp.ones(3))]  # FFT matvec
[2.75, 3.0, 2.75]

Source code in src/gaussx/_operators/_toeplitz.py

class Toeplitz(lx.AbstractLinearOperator):
    r"""Symmetric Toeplitz matrix from its first column.

    ``K_{ij} = c_{|i-j|}``.  Stored as O(n) with O(n log n) matvec
    via circulant embedding and FFT.

    For stationary kernels on regular 1-D grids the full kernel matrix
    is Toeplitz, so this gives an asymptotic win over dense storage.

    Args:
        column: First column of the Toeplitz matrix, shape ``(n,)``. Real;
            integer columns are promoted to floating point.

    Raises:
        TypeError: If ``column`` is complex.
        ValueError: If ``column`` is not rank 1 or is empty.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> import gaussx
        >>> T = gaussx.Toeplitz(jnp.array([2.0, 0.5, 0.25]))  # first column
        >>> T.as_matrix().tolist()
        [[2.0, 0.5, 0.25], [0.5, 2.0, 0.5], [0.25, 0.5, 2.0]]
        >>> [float(v) for v in T.mv(jnp.ones(3))]  # FFT matvec
        [2.75, 3.0, 2.75]
    """

    column: Float[Array, " n"]
    _size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        column: Float[Array, " n"],
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        self.column = _as_floating_column(column)
        self._size = self.column.shape[0]
        self._dtype = str(self.column.dtype)
        self.tags = _to_frozenset(tags) | {lx.symmetric_tag}

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        return _toeplitz_mv(self.column, vector)

    def as_matrix(self) -> Float[Array, "n n"]:
        n = self._size
        indices = jnp.abs(jnp.arange(n)[:, None] - jnp.arange(n)[None, :])
        return self.column[indices]

    def transpose(self) -> Toeplitz:
        # Symmetric Toeplitz is self-transpose.
        return self

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

ToeplitzCholesky

Bases: AbstractLinearOperator

Circulant-embedding sample factor for a symmetric positive Toeplitz matrix.

The operator has shape (n, embedding_factor * n) and satisfies L @ L.T == Toeplitz(column) — it is a rectangular sample factor, not a traditional lower-triangular Cholesky factor. Applying it to standard normal white noise gives samples from 𝒩(0, Toeplitz(column)) when the Wood--Chan condition holds.

The Wood--Chan non-negativity check is implemented with eqx.error_if so it is JIT-friendly: the error fires at evaluation time rather than tracing time. If the embedding's spectrum has a materially negative eigenvalue, bump embedding_factor (typically 4 or 8 suffices for well-behaved covariances).

Parameters:

Name	Type	Description	Default
column	Float[Array, ' n']	First column of the Toeplitz matrix, shape (n,).	required
embedding_factor	int	Circulant embedding size as a multiple of n.	2

Source code in src/gaussx/_operators/_toeplitz.py

class ToeplitzCholesky(lx.AbstractLinearOperator):
    """Circulant-embedding sample factor for a symmetric positive Toeplitz matrix.

    The operator has shape ``(n, embedding_factor * n)`` and satisfies
    ``L @ L.T == Toeplitz(column)`` — it is a rectangular sample factor,
    *not* a traditional lower-triangular Cholesky factor. Applying it to
    standard normal white noise gives samples from
    ``𝒩(0, Toeplitz(column))`` when the Wood--Chan condition holds.

    The Wood--Chan non-negativity check is implemented with
    ``eqx.error_if`` so it is JIT-friendly: the error fires at evaluation
    time rather than tracing time. If the embedding's spectrum has a
    materially negative eigenvalue, bump ``embedding_factor`` (typically
    ``4`` or ``8`` suffices for well-behaved covariances).

    Args:
        column: First column of the Toeplitz matrix, shape ``(n,)``.
        embedding_factor: Circulant embedding size as a multiple of ``n``.
    """

    sqrt_spectrum: Float[Array, " m_rfft"]
    _size: int = eqx.field(static=True)
    _embedding_size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        column: Float[Array, " n"],
        *,
        embedding_factor: int = 2,
    ) -> None:
        column = _as_floating_column(column)
        self.sqrt_spectrum = _circulant_sqrt_spectrum(
            column,
            embedding_factor=embedding_factor,
        )
        self._size = column.shape[0]
        self._embedding_size = embedding_factor * self._size
        self._dtype = str(column.dtype)
        self.tags = frozenset()

    def mv(self, vector: Float[Array, " m"]) -> Float[Array, " n"]:
        return _circulant_sqrt_mv(self.sqrt_spectrum, vector, self._size)

    def as_matrix(self) -> Float[Array, "n m"]:
        eye = jnp.eye(self._embedding_size, dtype=jnp.dtype(self._dtype))
        return jax.vmap(self.mv, in_axes=1, out_axes=1)(eye)

    def transpose(self) -> _ToeplitzCholeskyTranspose:
        return _ToeplitzCholeskyTranspose(self)

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._embedding_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

Fast-diagonalisable operators

DiagonalizedOperator is an operator diagonal in a fast transform basis, \(A = V^{-1}\Lambda V\), given by a forward/inverse transform pair (FFT, orthonormal DCT/DST, spherical harmonics, or a dense eigenvector matrix) and the eigenvalue array \(\Lambda\). solve, logdet, inv, sqrt and trace are elementwise in \(\Lambda\), and shifts/scalings such as \(A - \lambda I\) stay diagonalised, so a spectral Helmholtz solve is never materialised. circulant / circulant_from_symbol are the FFT special case (periodic stencils, stationary covariances on periodic grids), and a KroneckerSum whose factors are all diagonalised solves through the composed per-axis transforms. Use DiagonalizedOperator.from_eigen_factorization for dense non-symmetric diagonalisable factors such as Chebyshev collocation blocks.

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

n = 128
k = 2 * jnp.pi * jnp.fft.fftfreq(n)
symbol = (2 * jnp.cos(k) - 2)[:, None] + (2 * jnp.cos(k) - 2)[None, :]
laplacian = gaussx.circulant_from_symbol(symbol)  # periodic 5-point ∇²
helmholtz = laplacian - 1.0 * lx.IdentityLinearOperator(laplacian.in_structure())
psi = gaussx.solve(helmholtz, f)  # two FFTs, no matrix

Structured linear algebra and Gaussian primitives for JAX.

DiagonalizedOperator

Bases: AbstractLinearOperator

Linear operator A x = V⁻¹(Λ ⊙ V x) given by a transform pair.

forward applies V (field of shape in_shape → coefficients) and inverse applies V⁻¹ (coefficients → field); eigenvalues holds Λ in the coefficient layout. The operator acts on flat vectors of length prod(in_shape), as every lineax operator does; mv reshapes internally.

Examples of (forward, inverse, Λ):

• Periodic grid: fftn / ifftn and the symbol of the stencil — see circulant and circulant_from_symbol.
• Dirichlet / Neumann grid: an orthonormal DST / DCT pair and the corresponding finite-difference or spectral eigenvalues.
• A dense diagonalisable matrix M = V Λ V⁻¹ (e.g. a Chebyshev collocation block): DiagonalizedOperator.from_eigen_factorization.

Algebra stays closed. A + c·I, A − c·I, c·A, A / c, −A and A ± B (for B with the same transform pair) return another DiagonalizedOperator with updated eigenvalues, so a shifted Helmholtz operator A − λI is never materialised.

Structured primitives. gaussx.solve divides by Λ (eigenvalues that are exactly zero get a zero coefficient, i.e. the pseudo-inverse for a normal operator); gaussx.logdet is Σ log|λ| (the slogdet convention, so −inf for a singular operator); gaussx.inv, gaussx.sqrt and gaussx.trace are elementwise in Λ. A gaussx.KroneckerSum whose factors are all DiagonalizedOperator s is solved through the composed per-axis transforms.

Transpose. With normal=True (V⁻¹ ∝ Vᴴ, e.g. FFT or orthonormal DCT/DST), Aᵀ for a real operator (real_output=True) is the same pair with conjugated eigenvalues; for a complex one it is the same eigenvalues with the conjugated pair conj(V), conj(V)⁻¹, since Aᵀ x = conj(Aᴴ conj(x)). Otherwise pass transpose_pair=(forward_T, inverse_T) applying V⁻ᵀ and Vᵀ in the same roles; without either, transpose raises.

Parameters:

Name	Type	Description	Default
eigenvalues	Inexact[ArrayLike, ' ...']	Λ in the coefficient layout; eigenvalues.size must equal prod(in_shape). May be complex.	required
forward	Transform	V: array of shape in_shape → coefficient array.	required
inverse	Transform	V⁻¹: coefficient array → array of shape in_shape.	required
in_shape	tuple[int, ...]	Field shape the transforms act on.	required
real_output	bool	Return the real part of mv (default True), as needed when V is complex but A is real (FFT). The primitives work on Λ directly, so this is only consistent when V⁻¹ Λ V is itself real (for the FFT: a conjugate-even Λ); otherwise mv applies Re(A) while solve and logdet describe A. The generic class cannot check an arbitrary transform pair; circulant_from_symbol does.	True
normal	bool	Whether V⁻¹ is proportional to Vᴴ (see Transpose).	False
transpose_pair	tuple[Transform, Transform] | None	Optional (forward_T, inverse_T) for Aᵀ.	None
tags	object | frozenset[object]	lineax tags. symmetric_tag is added automatically when normal and real_output are set and the eigenvalues have a real dtype. (Real eigenvalues in a unitary basis make A Hermitian; it is also symmetric only when it is real.)	frozenset()

The transform callables are stored as static fields, so they must be hashable (module-level functions, functools.partial of them, or closures); keep them stable between calls to avoid jit recompiles.

Source code in src/gaussx/_operators/_diagonalised.py

class DiagonalizedOperator(lx.AbstractLinearOperator):
    r"""Linear operator ``A x = V⁻¹(Λ ⊙ V x)`` given by a transform pair.

    ``forward`` applies ``V`` (field of shape ``in_shape`` → coefficients) and
    ``inverse`` applies ``V⁻¹`` (coefficients → field); ``eigenvalues`` holds
    ``Λ`` in the coefficient layout. The operator acts on **flat** vectors of
    length ``prod(in_shape)``, as every lineax operator does; ``mv`` reshapes
    internally.

    Examples of ``(forward, inverse, Λ)``:

    - Periodic grid: ``fftn`` / ``ifftn`` and the symbol of the stencil — see
      `circulant` and `circulant_from_symbol`.
    - Dirichlet / Neumann grid: an orthonormal DST / DCT pair and the
      corresponding finite-difference or spectral eigenvalues.
    - A dense diagonalisable matrix ``M = V Λ V⁻¹`` (e.g. a Chebyshev
      collocation block): `DiagonalizedOperator.from_eigen_factorization`.

    **Algebra stays closed.** ``A + c·I``, ``A − c·I``, ``c·A``, ``A / c``,
    ``−A`` and ``A ± B`` (for ``B`` with the same transform pair) return
    another `DiagonalizedOperator` with updated eigenvalues, so a shifted
    Helmholtz operator ``A − λI`` is never materialised.

    **Structured primitives.** `gaussx.solve` divides by ``Λ`` (eigenvalues
    that are exactly zero get a zero coefficient, i.e. the pseudo-inverse for
    a normal operator); `gaussx.logdet` is ``Σ log|λ|`` (the ``slogdet``
    convention, so ``−inf`` for a singular operator); `gaussx.inv`,
    `gaussx.sqrt` and `gaussx.trace` are elementwise in ``Λ``. A
    `gaussx.KroneckerSum` whose factors are all `DiagonalizedOperator` s is
    solved through the composed per-axis transforms.

    **Transpose.** With ``normal=True`` (``V⁻¹ ∝ Vᴴ``, e.g. FFT or orthonormal
    DCT/DST), ``Aᵀ`` for a real operator (``real_output=True``) is the same
    pair with conjugated eigenvalues; for a complex one it is the same
    eigenvalues with the conjugated pair ``conj(V)``, ``conj(V)⁻¹``, since
    ``Aᵀ x = conj(Aᴴ conj(x))``. Otherwise pass
    ``transpose_pair=(forward_T, inverse_T)`` applying ``V⁻ᵀ`` and ``Vᵀ`` in
    the same roles; without either, `transpose` raises.

    Args:
        eigenvalues: ``Λ`` in the coefficient layout; ``eigenvalues.size``
            must equal ``prod(in_shape)``. May be complex.
        forward: ``V``: array of shape ``in_shape`` → coefficient array.
        inverse: ``V⁻¹``: coefficient array → array of shape ``in_shape``.
        in_shape: Field shape the transforms act on.
        real_output: Return the real part of ``mv`` (default ``True``), as
            needed when ``V`` is complex but ``A`` is real (FFT). The
            primitives work on ``Λ`` directly, so this is only consistent
            when ``V⁻¹ Λ V`` is itself real (for the FFT: a conjugate-even
            ``Λ``); otherwise ``mv`` applies ``Re(A)`` while ``solve`` and
            ``logdet`` describe ``A``. The generic class cannot check an
            arbitrary transform pair; `circulant_from_symbol` does.
        normal: Whether ``V⁻¹`` is proportional to ``Vᴴ`` (see Transpose).
        transpose_pair: Optional ``(forward_T, inverse_T)`` for ``Aᵀ``.
        tags: lineax tags. ``symmetric_tag`` is added automatically when
            ``normal`` and ``real_output`` are set and the eigenvalues have a
            real dtype. (Real eigenvalues in a unitary basis make ``A``
            Hermitian; it is also symmetric only when it is real.)

    The transform callables are stored as static fields, so they must be
    hashable (module-level functions, ``functools.partial`` of them, or
    closures); keep them stable between calls to avoid ``jit`` recompiles.
    """

    eigenvalues: Inexact[Array, " ..."]
    forward: Transform = eqx.field(static=True)
    inverse: Transform = eqx.field(static=True)
    in_shape: tuple[int, ...] = eqx.field(static=True)
    real_output: bool = eqx.field(static=True)
    normal: bool = eqx.field(static=True)
    transpose_pair: tuple[Transform, Transform] | None = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        eigenvalues: Inexact[ArrayLike, " ..."],
        forward: Transform,
        inverse: Transform,
        in_shape: tuple[int, ...],
        *,
        real_output: bool = True,
        normal: bool = False,
        transpose_pair: tuple[Transform, Transform] | None = None,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        eigenvalues = jnp.asarray(eigenvalues)
        in_shape = tuple(int(d) for d in in_shape)
        size = 1
        for d in in_shape:
            size *= d
        if eigenvalues.size != size:
            raise ValueError(
                f"eigenvalues has {eigenvalues.size} entries but in_shape "
                f"{in_shape} has {size}."
            )
        self.eigenvalues = eigenvalues
        self.forward = forward
        self.inverse = inverse
        self.in_shape = in_shape
        self.real_output = real_output
        self.normal = normal
        self.transpose_pair = transpose_pair
        tags = _to_frozenset(tags)
        if normal and real_output and not jnp.iscomplexobj(eigenvalues):
            tags = tags | {lx.symmetric_tag}
        self.tags = tags

    # ------------------------------------------------------------------
    # Constructors
    # ------------------------------------------------------------------

    @classmethod
    def from_eigen_factorization(
        cls,
        factorization,
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> DiagonalizedOperator:
        """Operator ``M = V Λ V⁻¹`` from a `gaussx.EigenDecomposition`.

        Uses the dense eigenvector matrices as the transform pair
        (``forward = V⁻¹ ·``, ``inverse = V ·``), so non-symmetric but
        diagonalisable factors — e.g. Chebyshev collocation second-derivative
        blocks — get an exact structured solve and transpose.
        """
        V = factorization.eigenvectors
        V_inv = factorization.eigenvectors_inv
        n = V.shape[0]
        return cls(
            factorization.eigenvalues,
            forward=_MatrixApply(V_inv),
            inverse=_MatrixApply(V),
            in_shape=(n,),
            transpose_pair=(_MatrixApply(V.T), _MatrixApply(V_inv.T)),
            tags=tags,
        )

    # ------------------------------------------------------------------
    # Flat-vector transforms (used by solves and Kronecker-sum composition)
    # ------------------------------------------------------------------

    @property
    def size(self) -> int:
        """Number of degrees of freedom, ``prod(in_shape)``."""
        return int(self.eigenvalues.size)

    def forward_flat(self, vector: Shaped[Array, " n"]) -> Shaped[Array, " n"]:
        """``V`` on a flat vector, returning flat coefficients."""
        return _flatten(self.forward(_unflatten(vector, self.in_shape)))

    def inverse_flat(self, coeffs: Shaped[Array, " n"]) -> Shaped[Array, " n"]:
        """``V⁻¹`` on flat coefficients, returning a flat vector."""
        return _flatten(
            self.inverse(_unflatten(coeffs, self.eigenvalues.shape)),
        )

    def eigenvalues_flat(self) -> Inexact[Array, " n"]:
        """``Λ`` flattened to match `forward_flat`."""
        return _flatten(self.eigenvalues)

    def with_eigenvalues(
        self, eigenvalues: Inexact[ArrayLike, " ..."]
    ) -> DiagonalizedOperator:
        """Same transform pair with new eigenvalues.

        Structural tags are recomputed: the automatic ``symmetric_tag`` is
        re-derived, and definiteness tags (``positive_semidefinite_tag``,
        ``negative_semidefinite_tag``) are dropped because a shift or sign
        change can invalidate them. Re-tag the result if you know it holds.
        """
        user_tags = self.tags - {
            lx.symmetric_tag,
            lx.positive_semidefinite_tag,
            lx.negative_semidefinite_tag,
        }
        return DiagonalizedOperator(
            jnp.asarray(eigenvalues),
            self.forward,
            self.inverse,
            self.in_shape,
            real_output=self.real_output,
            normal=self.normal,
            transpose_pair=self.transpose_pair,
            tags=user_tags,
        )

    # ------------------------------------------------------------------
    # lineax interface
    # ------------------------------------------------------------------

    def mv(self, vector: Shaped[Array, " n"]) -> Inexact[Array, " n"]:
        coeffs = self.forward(_unflatten(vector, self.in_shape))
        out = _flatten(self.inverse(self.eigenvalues * coeffs))
        return jnp.real(out) if self.real_output else out

    def as_matrix(self) -> Inexact[Array, "n n"]:
        eye = jnp.eye(self.size, dtype=self.out_structure().dtype)
        return jax.vmap(self.mv, out_axes=1)(eye)

    def transpose(self) -> DiagonalizedOperator:
        tags = lx.transpose_tags(self.tags)
        if self.normal and not self.real_output:
            # Conjugated eigenvalues would give Aᴴ, not Aᵀ (gh-330).
            return DiagonalizedOperator(
                self.eigenvalues,
                _conjugated(self.forward),
                _conjugated(self.inverse),
                self.in_shape,
                real_output=False,
                normal=True,
                tags=tags,
            )
        if self.normal:
            return DiagonalizedOperator(
                jnp.conj(self.eigenvalues),
                self.forward,
                self.inverse,
                self.in_shape,
                real_output=self.real_output,
                normal=True,
                transpose_pair=self.transpose_pair,
                tags=tags,
            )
        if self.transpose_pair is None:
            raise NotImplementedError(
                "DiagonalizedOperator.transpose needs normal=True or a transpose_pair."
            )
        forward_t, inverse_t = self.transpose_pair
        return DiagonalizedOperator(
            self.eigenvalues,
            forward_t,
            inverse_t,
            self.in_shape,
            real_output=self.real_output,
            transpose_pair=(self.forward, self.inverse),
            tags=tags,
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self.size,), self._dtype())

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self.size,), self._dtype())

    def _dtype(self):
        dtype = self.eigenvalues.dtype
        if self.real_output:
            return jnp.real(jnp.zeros((), dtype)).dtype
        return jnp.result_type(dtype, jnp.complex64)

    # ------------------------------------------------------------------
    # Closed algebra: shifts, scalings, sums in the same basis
    # ------------------------------------------------------------------

    def _same_basis(self, other: object) -> bool:
        return (
            isinstance(other, DiagonalizedOperator)
            and other.forward is self.forward
            and other.inverse is self.inverse
            and other.in_shape == self.in_shape
            and other.eigenvalues.shape == self.eigenvalues.shape
        )

    def __add__(self, other):
        shift = _identity_multiple(other)
        if shift is not None:
            return self.with_eigenvalues(self.eigenvalues + shift)
        if self._same_basis(other):
            return self.with_eigenvalues(self.eigenvalues + other.eigenvalues)
        return super().__add__(other)

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        shift = _identity_multiple(other)
        if shift is not None:
            return self.with_eigenvalues(self.eigenvalues - shift)
        if self._same_basis(other):
            return self.with_eigenvalues(self.eigenvalues - other.eigenvalues)
        return super().__sub__(other)

    def __mul__(self, other):
        if _is_scalar(other):
            return self.with_eigenvalues(self.eigenvalues * other)
        return super().__mul__(other)

    def __rmul__(self, other):
        if _is_scalar(other):
            return self.with_eigenvalues(self.eigenvalues * other)
        return super().__rmul__(other)

    def __truediv__(self, other):
        if _is_scalar(other):
            return self.with_eigenvalues(self.eigenvalues / other)
        return super().__truediv__(other)

    def __neg__(self):
        return self.with_eigenvalues(-self.eigenvalues)

size: int property

Number of degrees of freedom, prod(in_shape).

from_eigen_factorization(factorization, *, tags: object | frozenset[object] = frozenset()) -> DiagonalizedOperator classmethod

Operator M = V Λ V⁻¹ from a gaussx.EigenDecomposition.

Uses the dense eigenvector matrices as the transform pair (forward = V⁻¹ ·, inverse = V ·), so non-symmetric but diagonalisable factors — e.g. Chebyshev collocation second-derivative blocks — get an exact structured solve and transpose.

Source code in src/gaussx/_operators/_diagonalised.py

@classmethod
def from_eigen_factorization(
    cls,
    factorization,
    *,
    tags: object | frozenset[object] = frozenset(),
) -> DiagonalizedOperator:
    """Operator ``M = V Λ V⁻¹`` from a `gaussx.EigenDecomposition`.

    Uses the dense eigenvector matrices as the transform pair
    (``forward = V⁻¹ ·``, ``inverse = V ·``), so non-symmetric but
    diagonalisable factors — e.g. Chebyshev collocation second-derivative
    blocks — get an exact structured solve and transpose.
    """
    V = factorization.eigenvectors
    V_inv = factorization.eigenvectors_inv
    n = V.shape[0]
    return cls(
        factorization.eigenvalues,
        forward=_MatrixApply(V_inv),
        inverse=_MatrixApply(V),
        in_shape=(n,),
        transpose_pair=(_MatrixApply(V.T), _MatrixApply(V_inv.T)),
        tags=tags,
    )

forward_flat(vector: Shaped[Array, ' n']) -> Shaped[Array, ' n']

V on a flat vector, returning flat coefficients.

Source code in src/gaussx/_operators/_diagonalised.py

def forward_flat(self, vector: Shaped[Array, " n"]) -> Shaped[Array, " n"]:
    """``V`` on a flat vector, returning flat coefficients."""
    return _flatten(self.forward(_unflatten(vector, self.in_shape)))

inverse_flat(coeffs: Shaped[Array, ' n']) -> Shaped[Array, ' n']

V⁻¹ on flat coefficients, returning a flat vector.

Source code in src/gaussx/_operators/_diagonalised.py

def inverse_flat(self, coeffs: Shaped[Array, " n"]) -> Shaped[Array, " n"]:
    """``V⁻¹`` on flat coefficients, returning a flat vector."""
    return _flatten(
        self.inverse(_unflatten(coeffs, self.eigenvalues.shape)),
    )

eigenvalues_flat() -> Inexact[Array, ' n']

Λ flattened to match forward_flat.

Source code in src/gaussx/_operators/_diagonalised.py

def eigenvalues_flat(self) -> Inexact[Array, " n"]:
    """``Λ`` flattened to match `forward_flat`."""
    return _flatten(self.eigenvalues)

with_eigenvalues(eigenvalues: Inexact[ArrayLike, ' ...']) -> DiagonalizedOperator

Same transform pair with new eigenvalues.

Structural tags are recomputed: the automatic symmetric_tag is re-derived, and definiteness tags (positive_semidefinite_tag, negative_semidefinite_tag) are dropped because a shift or sign change can invalidate them. Re-tag the result if you know it holds.

Source code in src/gaussx/_operators/_diagonalised.py

def with_eigenvalues(
    self, eigenvalues: Inexact[ArrayLike, " ..."]
) -> DiagonalizedOperator:
    """Same transform pair with new eigenvalues.

    Structural tags are recomputed: the automatic ``symmetric_tag`` is
    re-derived, and definiteness tags (``positive_semidefinite_tag``,
    ``negative_semidefinite_tag``) are dropped because a shift or sign
    change can invalidate them. Re-tag the result if you know it holds.
    """
    user_tags = self.tags - {
        lx.symmetric_tag,
        lx.positive_semidefinite_tag,
        lx.negative_semidefinite_tag,
    }
    return DiagonalizedOperator(
        jnp.asarray(eigenvalues),
        self.forward,
        self.inverse,
        self.in_shape,
        real_output=self.real_output,
        normal=self.normal,
        transpose_pair=self.transpose_pair,
        tags=user_tags,
    )

circulant(first_column: Inexact[ArrayLike, ' ...'], *, symmetric: bool | None = None, tags: object | frozenset[object] = frozenset()) -> DiagonalizedOperator

(Block-)circulant operator from its first column.

A circulant matrix C[i, j] = c[(i − j) mod n] (and its n-D block-circulant analogue, with first_column of the grid's shape) is diagonalised by the DFT: C = F⁻¹ diag(F c) F. Periodic finite-difference stencils and stationary covariances on periodic grids are circulant.

Compare gaussx.Toeplitz, which represents a (non-periodic) symmetric Toeplitz matrix with an FFT matvec by circulant embedding but has no closed-form solve; a circulant matrix has exact elementwise solve / logdet / sqrt.

A real even kernel, c[k] = c[−k mod n], has a real symbol and gives a symmetric operator. Say so with symmetric=True: the decision then depends only on the flag, so construction under jax.jit gives the same eigenvalue dtype and tags as eager construction (gh-440).

Parameters:

Name	Type	Description	Default
first_column	Inexact[ArrayLike, ' ...']	c, shape = grid shape ((n,) in 1D).	required
symmetric	bool | None	True for a real even kernel: the symbol is real and the operator is tagged symmetric (a concrete column that is not even raises). False keeps the complex symbol. None (the deprecated default) infers it from the column's values, which only works eagerly — a traced column gives a complex symbol and no tag — and warns when it infers True. The default becomes False in gaussx 0.7.0.	None
tags	object | frozenset[object]	Additional lineax tags.	frozenset()

Returns:

Type	Description
DiagonalizedOperator	A DiagonalizedOperator with eigenvalues fftn(c).

Raises:

Type	Description
ValueError	If symmetric=True and the column is complex, or concrete and not even.

Source code in src/gaussx/_operators/_diagonalised.py

def circulant(
    first_column: Inexact[ArrayLike, " ..."],
    *,
    symmetric: bool | None = None,
    tags: object | frozenset[object] = frozenset(),
) -> DiagonalizedOperator:
    r"""(Block-)circulant operator from its first column.

    A circulant matrix ``C[i, j] = c[(i − j) mod n]`` (and its n-D
    block-circulant analogue, with ``first_column`` of the grid's shape) is
    diagonalised by the DFT: ``C = F⁻¹ diag(F c) F``. Periodic
    finite-difference stencils and stationary covariances on periodic grids
    are circulant.

    Compare `gaussx.Toeplitz`, which represents a (non-periodic) symmetric
    Toeplitz matrix with an FFT matvec by circulant embedding but has no
    closed-form solve; a circulant matrix has exact elementwise
    solve / logdet / sqrt.

    A real even kernel, ``c[k] = c[−k mod n]``, has a real symbol and gives a
    symmetric operator. Say so with ``symmetric=True``: the decision then
    depends only on the flag, so construction under `jax.jit` gives the same
    eigenvalue dtype and tags as eager construction (gh-440).

    Args:
        first_column: ``c``, shape = grid shape (``(n,)`` in 1D).
        symmetric: ``True`` for a real even kernel: the symbol is real and
            the operator is tagged symmetric (a concrete column that is not
            even raises). ``False`` keeps the complex symbol. ``None`` (the
            deprecated default) infers it from the column's *values*, which
            only works eagerly — a traced column gives a complex symbol and
            no tag — and warns when it infers ``True``. The default
            becomes ``False`` in gaussx 0.7.0.
        tags: Additional lineax tags.

    Returns:
        A `DiagonalizedOperator` with eigenvalues ``fftn(c)``.

    Raises:
        ValueError: If ``symmetric=True`` and the column is complex, or
            concrete and not even.
    """
    column = jnp.asarray(first_column)
    symbol = jnp.fft.fftn(column)
    if symmetric is None:
        symmetric = _is_even_real_kernel(column)
        if symmetric:
            warn_deprecated(
                "circulant infers symmetry from the column's values, so a traced "
                "column gives a different operator (gh-440). Pass symmetric=True "
                "for a real even kernel; the default becomes False in gaussx 0.7.0."
            )
    elif symmetric:
        _check_even_real_kernel(column)
    if symmetric:
        # c[k] = c[−k] (mod n) and real ⇒ real symbol ⇒ symmetric operator.
        symbol = jnp.real(symbol)
    return circulant_from_symbol(
        symbol,
        real_output=not jnp.iscomplexobj(column),
        tags=tags,
    )

circulant_from_symbol(symbol: Inexact[ArrayLike, ' ...'], *, real_output: bool = True, tags: object | frozenset[object] = frozenset()) -> DiagonalizedOperator

(Block-)circulant operator from its DFT symbol.

A x = ifftn(symbol ⊙ fftn(x)) on a periodic grid of shape symbol.shape. Use this when the eigenvalues are known in closed form, e.g. the five-point Laplacian Σ_d (2 cos(2π k_d / n_d) − 2) / h_d² or the spectral −|k|². A real-dtype symbol gives a symmetric operator (tagged automatically); add lx.positive_semidefinite_tag / negative_semidefinite_tag yourself when it applies.

With real_output=True the symbol must be conjugate-even, symbol[k] = conj(symbol[−k mod n]) (for a real symbol: even), since only then is ifftn(symbol ⊙ fftn(x)) real for real x. Any other symbol would make mv apply the real part of the operator while solve / logdet / inv / sqrt describe the complex one.

Parameters:

Name	Type	Description	Default
symbol	Inexact[ArrayLike, ' ...']	Eigenvalues in numpy.fft frequency order, shape = grid shape.	required
real_output	bool	Take the real part of mv (default True).	True
tags	object | frozenset[object]	Additional lineax tags.	frozenset()

Returns:

Type	Description
DiagonalizedOperator	A DiagonalizedOperator with the FFT transform pair.

Raises:

Type	Description
ValueError	If real_output and a concrete symbol is not conjugate-even. A traced symbol is checked at run time and raises an EquinoxRuntimeError instead.

Source code in src/gaussx/_operators/_diagonalised.py

def circulant_from_symbol(
    symbol: Inexact[ArrayLike, " ..."],
    *,
    real_output: bool = True,
    tags: object | frozenset[object] = frozenset(),
) -> DiagonalizedOperator:
    r"""(Block-)circulant operator from its DFT symbol.

    ``A x = ifftn(symbol ⊙ fftn(x))`` on a periodic grid of shape
    ``symbol.shape``. Use this when the eigenvalues are known in closed form,
    e.g. the five-point Laplacian ``Σ_d (2 cos(2π k_d / n_d) − 2) / h_d²`` or
    the spectral ``−|k|²``. A real-dtype symbol gives a symmetric operator
    (tagged automatically); add ``lx.positive_semidefinite_tag`` /
    ``negative_semidefinite_tag`` yourself when it applies.

    With ``real_output=True`` the symbol must be conjugate-even,
    ``symbol[k] = conj(symbol[−k mod n])`` (for a real symbol: even), since
    only then is ``ifftn(symbol ⊙ fftn(x))`` real for real ``x``. Any other
    symbol would make ``mv`` apply the real part of the operator while
    ``solve`` / ``logdet`` / ``inv`` / ``sqrt`` describe the complex one.

    Args:
        symbol: Eigenvalues in `numpy.fft` frequency order, shape = grid shape.
        real_output: Take the real part of ``mv`` (default ``True``).
        tags: Additional lineax tags.

    Returns:
        A `DiagonalizedOperator` with the FFT transform pair.

    Raises:
        ValueError: If ``real_output`` and a concrete ``symbol`` is not
            conjugate-even. A traced symbol is checked at run time and
            raises an ``EquinoxRuntimeError`` instead.
    """
    symbol = jnp.asarray(symbol)
    if real_output:
        symbol = _check_conjugate_even(symbol)
    return DiagonalizedOperator(
        symbol,
        _fftn,
        _ifftn,
        symbol.shape,
        real_output=real_output,
        normal=True,
        tags=tags,
    )

as_diagonalized(operator: lx.AbstractLinearOperator) -> DiagonalizedOperator | None

Express operator as one DiagonalizedOperator, if possible.

Returns the operator itself for a DiagonalizedOperator, the composed operator for a gaussx.KroneckerSum whose factors are (recursively) diagonalisable, and None otherwise.

For A ⊕ B with A = V_A⁻¹ Λ_A V_A and B = V_B⁻¹ Λ_B V_B, A ⊕ B = (V_A ⊗ V_B)⁻¹ (Λ_A ⊕ Λ_B) (V_A ⊗ V_B): the forward transform applies V_A along the leading axis of the (n_A, n_B) field and V_B along the trailing one, and the eigenvalues are all pairwise sums λ^A_i + λ^B_j.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	The operator to express in a transform basis.	required

Returns:

Type	Description
DiagonalizedOperator | None	The equivalent DiagonalizedOperator, or None when operator
DiagonalizedOperator | None	is neither a DiagonalizedOperator nor a gaussx.KroneckerSum of
DiagonalizedOperator | None	diagonalisable factors.

Source code in src/gaussx/_operators/_diagonalised.py

def as_diagonalized(operator: lx.AbstractLinearOperator) -> DiagonalizedOperator | None:
    """Express ``operator`` as one `DiagonalizedOperator`, if possible.

    Returns the operator itself for a `DiagonalizedOperator`, the composed
    operator for a `gaussx.KroneckerSum` whose factors are (recursively)
    diagonalisable, and ``None`` otherwise.

    For ``A ⊕ B`` with ``A = V_A⁻¹ Λ_A V_A`` and ``B = V_B⁻¹ Λ_B V_B``,
    ``A ⊕ B = (V_A ⊗ V_B)⁻¹ (Λ_A ⊕ Λ_B) (V_A ⊗ V_B)``: the forward transform
    applies ``V_A`` along the leading axis of the ``(n_A, n_B)`` field and
    ``V_B`` along the trailing one, and the eigenvalues are all pairwise sums
    ``λ^A_i + λ^B_j``.

    Args:
        operator: The operator to express in a transform basis.

    Returns:
        The equivalent `DiagonalizedOperator`, or ``None`` when ``operator``
        is neither a `DiagonalizedOperator` nor a `gaussx.KroneckerSum` of
        diagonalisable factors.
    """

    if isinstance(operator, DiagonalizedOperator):
        return operator
    if not isinstance(operator, KroneckerSum):
        return None
    a = as_diagonalized(operator.A)
    b = as_diagonalized(operator.B)
    if a is None or b is None:
        return None
    lam = jnp.add.outer(a.eigenvalues_flat(), b.eigenvalues_flat())
    real_output = a.real_output and b.real_output
    return DiagonalizedOperator(
        lam,
        ft.partial(_kron_forward, a, b),
        ft.partial(_kron_inverse, a, b),
        (a.size, b.size),
        real_output=real_output,
        normal=a.normal and b.normal,
    )

SpectralFunction is \(f(A_1\oplus\cdots\oplus A_d)\) for symmetric factors, stored as the factor eigendecompositions plus the elementwise function \(f\): mv, solve, logdet, diag_inv (also with pinv=True) and sqrt_matmul act one axis at a time through the shared eigenvectors. It is the operator behind the SPDE (Matérn) precision on a regular grid, spde_precision_grid, and as a factor of a shifted Kronecker product \(A\otimes f(B) + cI\) it keeps its eigenbasis.

Structured linear algebra and Gaussian primitives for JAX.

SpectralFunction

Bases: AbstractLinearOperator

f(B) for a symmetric B with a factored eigenbasis.

B is typically a Kronecker sum A₁ ⊕ … ⊕ A_d of symmetric factors (a grid Laplacian L_H ⊕ L_W). With A_m = V_m Λ_m V_mᵀ,

\[ f(A_1 \oplus \cdots \oplus A_d) = (V_1 \otimes \cdots \otimes V_d)\, f(\Lambda_1 \oplus \cdots \oplus \Lambda_d)\, (V_1 \otimes \cdots \otimes V_d)^\top, \]

so mv, solve, logdet, diag_inv and sqrt_matmul all act on the (n_1, …, n_d) grid one axis at a time, O(N Σ_m n_m), and the joint eigenbasis is never formed. The factor eigendecompositions are computed once, at construction.

gaussx.solve, gaussx.logdet, gaussx.diag_inv (including pinv=True) and gaussx.diag dispatch here, and a SpectralFunction factor of a shifted Kronecker product A ⊗ f(B) + c·I keeps its eigenbasis there too.

Parameters:

Name	Type	Description	Default
base	AbstractLinearOperator	The symmetric operator B (a gaussx.KroneckerSum, a gaussx.DiagonalizedOperator, a diagonal or a symmetric dense operator). Its factors are eigendecomposed with eigh; build it from concrete values (outside jit), or use from_eigen_factorizations to supply the decompositions.	required
fn	Callable[[Array], Array]	The elementwise function f, applied to the eigenvalue tensor of B. Must be real-valued on the spectrum. Pass an equinox.Module (or a jax.tree_util.Partial) when it carries parameters that are traced or differentiated, so they stay pytree leaves.	required
tags	object | frozenset[object]	lineax tags. symmetric_tag is always added.	frozenset()

Raises:

Type	Description
ValueError	If base has no factored eigenbasis (it is not symmetric).

Examples:

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

def path_laplacian(n):
    off = -jnp.ones(n - 1)
    L = jnp.diag(jnp.r_[1.0, 2.0 * jnp.ones(n - 2), 1.0])
    return lx.MatrixLinearOperator(
        L + jnp.diag(off, 1) + jnp.diag(off, -1), lx.symmetric_tag
    )

# (κ²I + L_H ⊕ L_W)², a Matérn nu = 1 precision on a 4 x 5 grid
base = gaussx.KroneckerSum(path_laplacian(4), path_laplacian(5))
Q = gaussx.SpectralFunction(base, lambda lam: (0.5 + lam) ** 2)
dense = jnp.linalg.matrix_power(0.5 * jnp.eye(20) + base.as_matrix(), 2)
assert jnp.allclose(Q.as_matrix(), dense, atol=1e-4)
assert jnp.allclose(gaussx.diag_inv(Q), jnp.diag(jnp.linalg.inv(dense)))

Source code in src/gaussx/_operators/_spectral_function.py

class SpectralFunction(lx.AbstractLinearOperator):
    r"""``f(B)`` for a symmetric ``B`` with a factored eigenbasis.

    ``B`` is typically a Kronecker sum ``A₁ ⊕ … ⊕ A_d`` of symmetric factors
    (a grid Laplacian ``L_H ⊕ L_W``). With ``A_m = V_m Λ_m V_mᵀ``,

    $$
    f(A_1 \oplus \cdots \oplus A_d)
    = (V_1 \otimes \cdots \otimes V_d)\,
      f(\Lambda_1 \oplus \cdots \oplus \Lambda_d)\,
      (V_1 \otimes \cdots \otimes V_d)^\top,
    $$

    so `mv`, `solve`, `logdet`, `diag_inv` and `sqrt_matmul` all act on the
    ``(n_1, …, n_d)`` grid one axis at a time, ``O(N Σ_m n_m)``, and the
    joint eigenbasis is never formed. The factor eigendecompositions are
    computed once, at construction.

    `gaussx.solve`, `gaussx.logdet`, `gaussx.diag_inv` (including
    ``pinv=True``) and `gaussx.diag` dispatch here, and a `SpectralFunction`
    factor of a shifted Kronecker product ``A ⊗ f(B) + c·I`` keeps its
    eigenbasis there too.

    Args:
        base: The symmetric operator ``B`` (a `gaussx.KroneckerSum`, a
            `gaussx.DiagonalizedOperator`, a diagonal or a symmetric dense
            operator). Its factors are eigendecomposed with ``eigh``; build
            it from concrete values (outside ``jit``), or use
            `from_eigen_factorizations` to supply the decompositions.
        fn: The elementwise function ``f``, applied to the eigenvalue tensor
            of ``B``. Must be real-valued on the spectrum. Pass an
            `equinox.Module` (or a `jax.tree_util.Partial`) when it carries
            parameters that are traced or differentiated, so they stay
            pytree leaves.
        tags: lineax tags. ``symmetric_tag`` is always added.

    Raises:
        ValueError: If ``base`` has no factored eigenbasis (it is not
            symmetric).

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

        def path_laplacian(n):
            off = -jnp.ones(n - 1)
            L = jnp.diag(jnp.r_[1.0, 2.0 * jnp.ones(n - 2), 1.0])
            return lx.MatrixLinearOperator(
                L + jnp.diag(off, 1) + jnp.diag(off, -1), lx.symmetric_tag
            )

        # (κ²I + L_H ⊕ L_W)², a Matérn nu = 1 precision on a 4 x 5 grid
        base = gaussx.KroneckerSum(path_laplacian(4), path_laplacian(5))
        Q = gaussx.SpectralFunction(base, lambda lam: (0.5 + lam) ** 2)
        dense = jnp.linalg.matrix_power(0.5 * jnp.eye(20) + base.as_matrix(), 2)
        assert jnp.allclose(Q.as_matrix(), dense, atol=1e-4)
        assert jnp.allclose(gaussx.diag_inv(Q), jnp.diag(jnp.linalg.inv(dense)))
        ```
    """

    base: lx.AbstractLinearOperator
    fn: Callable[[Array], Array]
    eigen: FactoredEigen
    _size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        base: lx.AbstractLinearOperator,
        fn: Callable[[Array], Array],
        *,
        tags: object | frozenset[object] = frozenset(),
        _eigen: FactoredEigen | None = None,
    ) -> None:
        if _eigen is None:
            _eigen = factored_eigen(base)
            if _eigen is None:
                raise ValueError(
                    "SpectralFunction needs a symmetric base with a factored "
                    f"eigenbasis, got {type(base).__name__}."
                )
        self.base = base
        self.fn = fn
        self.eigen = _eigen
        self._size = int(_eigen.eigenvalues.size)
        spectrum = jax.eval_shape(fn, _eigen.eigenvalues)
        self._dtype = str(jnp.result_type(spectrum.dtype, _eigen.eigenvalues.dtype))
        self.tags = _to_frozenset(tags) | {lx.symmetric_tag}

    @classmethod
    def from_eigen_factorizations(
        cls,
        factors: Sequence[EigenDecomposition],
        fn: Callable[[Array], Array],
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> SpectralFunction:
        """``f(A₁ ⊕ … ⊕ A_d)`` from precomputed symmetric factorizations.

        Skips the ``eigh`` of `__init__`, so it is cheap inside ``jit``; the
        base is the nested `gaussx.KroneckerSum` of the reassembled factors.

        Args:
            factors: One `gaussx.EigenDecomposition` per axis, in Kronecker
                (row-major) order, each of a symmetric matrix (orthonormal
                eigenvectors, as from ``from_matrix(..., symmetric=True)``).
            fn: The elementwise function ``f`` (see the class docstring).
            tags: lineax tags.

        Returns:
            The operator.
        """
        if not factors:
            raise ValueError("factors must contain at least one factorization.")
        bases = tuple(f.eigenvectors for f in factors)
        eigenvalues = None
        for factor in factors:
            if eigenvalues is None:
                eigenvalues = factor.eigenvalues
                continue
            names = " ".join(f"a{m}" for m in range(eigenvalues.ndim))
            eigenvalues = einx.add(
                f"{names}, b -> {names} b", eigenvalues, factor.eigenvalues
            )
        assert eigenvalues is not None
        operators = [
            lx.MatrixLinearOperator(f.as_matrix(), lx.symmetric_tag) for f in factors
        ]
        base = operators[-1]
        for operator in reversed(operators[:-1]):
            base = KroneckerSum(operator, base)
        return cls(base, fn, tags=tags, _eigen=FactoredEigen(bases, eigenvalues))

    def spectrum(self) -> Float[Array, " ..."]:
        """``f(λ)`` on the eigenvalue tensor, shape ``(n_1, …, n_d)``."""
        return jnp.asarray(self.fn(self.eigen.eigenvalues), dtype=self._dtype)

    def factored_eigen(self) -> FactoredEigen:
        """The factored eigendecomposition of ``f(B)`` (same basis as ``B``)."""
        return FactoredEigen(self.eigen.bases, self.spectrum())

    def _apply(self, vector: Array, values: Array) -> Array:
        coeffs = self.eigen.to_eigenbasis(vector)
        return self.eigen._output(self.eigen.from_eigenbasis(coeffs * values), vector)

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        return self._apply(vector, self.spectrum())

    def solve(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        """``f(B)⁻¹ b``."""
        return self._apply(vector, 1.0 / self.spectrum())

    def logdet(self) -> Float[Array, ""]:
        """``log|det f(B)| = Σ log|f(λ)|``."""
        return jnp.sum(jnp.log(jnp.abs(self.spectrum())))

    def diag_inv(self, *, pinv: bool = False) -> Float[Array, " n"]:
        """``diag(f(B)⁻¹)`` (or of the pseudo-inverse), from the factor bases."""
        return self.factored_eigen().diag_inv(pinv=pinv)

    def diagonal(self) -> Float[Array, " n"]:
        """``diag(f(B))``, from the factor bases."""
        return FactoredEigen(self.eigen.bases, 1.0 / self.spectrum()).diag_inv()

    def sqrt_matmul(
        self, vector: Float[Array, " n"], *, inverse: bool = False
    ) -> Float[Array, " n"]:
        """``f(B)^{1/2} x``, or ``f(B)^{-1/2} x`` with ``inverse=True``.

        The symmetric square root: with ``inverse=True`` and ``z ~ N(0, I)``,
        ``f(B)^{-1/2} z`` is an exact sample with precision ``f(B)``.
        """
        root = jnp.sqrt(self.spectrum())
        return self._apply(vector, 1.0 / root if inverse else root)

    def as_matrix(self) -> Float[Array, "n n"]:
        eye = jnp.eye(self._size, dtype=self._dtype)
        # Rows are f(B) e_i, i.e. columns of the symmetric matrix.
        return rearrange(jax.vmap(self.mv)(eye), "i j -> j i")

    def transpose(self) -> SpectralFunction:
        return self

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._size,), jnp.dtype(self._dtype))

from_eigen_factorizations(factors: Sequence[EigenDecomposition], fn: Callable[[Array], Array], *, tags: object | frozenset[object] = frozenset()) -> SpectralFunction classmethod

f(A₁ ⊕ … ⊕ A_d) from precomputed symmetric factorizations.

Skips the eigh of __init__, so it is cheap inside jit; the base is the nested gaussx.KroneckerSum of the reassembled factors.

Parameters:

Name	Type	Description	Default
factors	Sequence[EigenDecomposition]	One gaussx.EigenDecomposition per axis, in Kronecker (row-major) order, each of a symmetric matrix (orthonormal eigenvectors, as from from_matrix(..., symmetric=True)).	required
fn	Callable[[Array], Array]	The elementwise function f (see the class docstring).	required
tags	object | frozenset[object]	lineax tags.	frozenset()

Returns:

Type	Description
SpectralFunction	The operator.

Source code in src/gaussx/_operators/_spectral_function.py

@classmethod
def from_eigen_factorizations(
    cls,
    factors: Sequence[EigenDecomposition],
    fn: Callable[[Array], Array],
    *,
    tags: object | frozenset[object] = frozenset(),
) -> SpectralFunction:
    """``f(A₁ ⊕ … ⊕ A_d)`` from precomputed symmetric factorizations.

    Skips the ``eigh`` of `__init__`, so it is cheap inside ``jit``; the
    base is the nested `gaussx.KroneckerSum` of the reassembled factors.

    Args:
        factors: One `gaussx.EigenDecomposition` per axis, in Kronecker
            (row-major) order, each of a symmetric matrix (orthonormal
            eigenvectors, as from ``from_matrix(..., symmetric=True)``).
        fn: The elementwise function ``f`` (see the class docstring).
        tags: lineax tags.

    Returns:
        The operator.
    """
    if not factors:
        raise ValueError("factors must contain at least one factorization.")
    bases = tuple(f.eigenvectors for f in factors)
    eigenvalues = None
    for factor in factors:
        if eigenvalues is None:
            eigenvalues = factor.eigenvalues
            continue
        names = " ".join(f"a{m}" for m in range(eigenvalues.ndim))
        eigenvalues = einx.add(
            f"{names}, b -> {names} b", eigenvalues, factor.eigenvalues
        )
    assert eigenvalues is not None
    operators = [
        lx.MatrixLinearOperator(f.as_matrix(), lx.symmetric_tag) for f in factors
    ]
    base = operators[-1]
    for operator in reversed(operators[:-1]):
        base = KroneckerSum(operator, base)
    return cls(base, fn, tags=tags, _eigen=FactoredEigen(bases, eigenvalues))

spectrum() -> Float[Array, ' ...']

f(λ) on the eigenvalue tensor, shape (n_1, …, n_d).

Source code in src/gaussx/_operators/_spectral_function.py

def spectrum(self) -> Float[Array, " ..."]:
    """``f(λ)`` on the eigenvalue tensor, shape ``(n_1, …, n_d)``."""
    return jnp.asarray(self.fn(self.eigen.eigenvalues), dtype=self._dtype)

factored_eigen() -> FactoredEigen

The factored eigendecomposition of f(B) (same basis as B).

Source code in src/gaussx/_operators/_spectral_function.py

def factored_eigen(self) -> FactoredEigen:
    """The factored eigendecomposition of ``f(B)`` (same basis as ``B``)."""
    return FactoredEigen(self.eigen.bases, self.spectrum())

solve(vector: Float[Array, ' n']) -> Float[Array, ' n']

f(B)⁻¹ b.

Source code in src/gaussx/_operators/_spectral_function.py

def solve(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
    """``f(B)⁻¹ b``."""
    return self._apply(vector, 1.0 / self.spectrum())

logdet() -> Float[Array, '']

log|det f(B)| = Σ log|f(λ)|.

Source code in src/gaussx/_operators/_spectral_function.py

def logdet(self) -> Float[Array, ""]:
    """``log|det f(B)| = Σ log|f(λ)|``."""
    return jnp.sum(jnp.log(jnp.abs(self.spectrum())))

diag_inv(*, pinv: bool = False) -> Float[Array, ' n']

diag(f(B)⁻¹) (or of the pseudo-inverse), from the factor bases.

Source code in src/gaussx/_operators/_spectral_function.py

def diag_inv(self, *, pinv: bool = False) -> Float[Array, " n"]:
    """``diag(f(B)⁻¹)`` (or of the pseudo-inverse), from the factor bases."""
    return self.factored_eigen().diag_inv(pinv=pinv)

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

diag(f(B)), from the factor bases.

Source code in src/gaussx/_operators/_spectral_function.py

def diagonal(self) -> Float[Array, " n"]:
    """``diag(f(B))``, from the factor bases."""
    return FactoredEigen(self.eigen.bases, 1.0 / self.spectrum()).diag_inv()

sqrt_matmul(vector: Float[Array, ' n'], *, inverse: bool = False) -> Float[Array, ' n']

f(B)^{1/2} x, or f(B)^{-1/2} x with inverse=True.

The symmetric square root: with inverse=True and z ~ N(0, I), f(B)^{-1/2} z is an exact sample with precision f(B).

Source code in src/gaussx/_operators/_spectral_function.py

def sqrt_matmul(
    self, vector: Float[Array, " n"], *, inverse: bool = False
) -> Float[Array, " n"]:
    """``f(B)^{1/2} x``, or ``f(B)^{-1/2} x`` with ``inverse=True``.

    The symmetric square root: with ``inverse=True`` and ``z ~ N(0, I)``,
    ``f(B)^{-1/2} z`` is an exact sample with precision ``f(B)``.
    """
    root = jnp.sqrt(self.spectrum())
    return self._apply(vector, 1.0 / root if inverse else root)

Interpolated & masked operators

Grid-interpolated (KISS-GP style) and masked operators, plus the grid and cubic-interpolation helpers that build them. Kernel operators themselves (KernelOperator, ImplicitKernelOperator, ImplicitCrossKernelOperator) moved to kernellib in gaussx 0.2.0; they are lineax operators, so every gaussx primitive and solver strategy works on them unchanged.

Structured linear algebra and Gaussian primitives for JAX.

InterpolatedOperator

Bases: AbstractLinearOperator

Structured Kernel Interpolation: K \approx W K_{uu} W^T.

W is a sparse interpolation matrix with p nonzeros per row (e.g. cubic interpolation weights). The base operator K_{uu} acts on the inducing grid (typically Toeplitz for stationary kernels).

Total matvec cost: O(n p + m log m) when the base is Toeplitz, essentially linear in n.

Parameters:

Name	Type	Description	Default
base_operator	AbstractLinearOperator	The inducing-point kernel K_{uu}, shape (m, m).	required
interp_indices	Int[Array, 'n p']	Integer indices into the inducing grid, shape (n, p) where p is the interpolation order.	required
interp_values	Float[Array, 'n p']	Interpolation weights, shape (n, p).	required
tags	object | frozenset[object]	Extra lineax tags. Symmetry and positive/negative semi-definiteness are inherited from base_operator structurally (W A W^T preserves them for any real W), so they need not be passed; caller tags are added on top.	frozenset()

Source code in src/gaussx/_operators/_interpolated.py

class InterpolatedOperator(lx.AbstractLinearOperator):
    r"""Structured Kernel Interpolation: ``K \approx W K_{uu} W^T``.

    ``W`` is a sparse interpolation matrix with ``p`` nonzeros per row
    (e.g. cubic interpolation weights).  The base operator ``K_{uu}``
    acts on the inducing grid (typically Toeplitz for stationary
    kernels).

    Total matvec cost: ``O(n p + m log m)`` when the base is Toeplitz,
    essentially linear in ``n``.

    Args:
        base_operator: The inducing-point kernel ``K_{uu}``, shape ``(m, m)``.
        interp_indices: Integer indices into the inducing grid,
            shape ``(n, p)`` where ``p`` is the interpolation order.
        interp_values: Interpolation weights, shape ``(n, p)``.
        tags: Extra lineax tags. Symmetry and positive/negative
            semi-definiteness are inherited from ``base_operator``
            structurally (``W A W^T`` preserves them for any real ``W``),
            so they need not be passed; caller tags are added on top.
    """

    base_operator: lx.AbstractLinearOperator
    interp_indices: Int[Array, "n p"]
    interp_values: Float[Array, "n p"]
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _m: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        base_operator: lx.AbstractLinearOperator,
        interp_indices: Int[Array, "n p"],
        interp_values: Float[Array, "n p"],
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:
        if interp_indices.shape != interp_values.shape:
            raise ValueError(
                f"interp_indices and interp_values must have the same shape, "
                f"got {interp_indices.shape} and {interp_values.shape}."
            )
        m = base_operator.in_size()
        out_m = base_operator.out_size()
        if m != out_m:
            raise ValueError(f"base_operator must be square, got shape ({out_m}, {m}).")
        n = interp_indices.shape[0]
        self.base_operator = base_operator
        self.interp_indices = jnp.asarray(interp_indices)
        self.interp_values = jnp.asarray(interp_values)
        self._in_size = n
        self._out_size = n
        self._m = m
        self._dtype = str(
            jnp.result_type(
                jnp.dtype(_resolve_dtype(base_operator)),
                self.interp_values.dtype,
            )
        )
        self.tags = _to_frozenset(tags) | _inherited_tags(base_operator)

    def mv(self, vector: Float[Array, " n"]) -> Float[Array, " n"]:
        # W^T v: scatter into inducing space
        wt_v = jnp.zeros(self._m, dtype=vector.dtype)
        wt_v = wt_v.at[self.interp_indices].add(self.interp_values * vector[:, None])
        # K_uu (W^T v)
        k_wt_v = self.base_operator.mv(wt_v)
        # W (K_uu W^T v): gather back
        return jnp.sum(self.interp_values * k_wt_v[self.interp_indices], axis=-1)

    def as_matrix(self) -> Float[Array, "n n"]:
        W = self._build_W()
        K_uu = self.base_operator.as_matrix()
        return W @ K_uu @ W.T

    def transpose(self) -> InterpolatedOperator:
        # W K_uu W^T is symmetric when K_uu is symmetric, so self-transpose.
        # In general, (W K_uu W^T)^T = W K_uu^T W^T.
        return InterpolatedOperator(
            self.base_operator.T,
            self.interp_indices,
            self.interp_values,
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

    def _build_W(self) -> Float[Array, "n m"]:
        """Build the dense interpolation matrix W (n x m)."""
        n = self._in_size
        W = jnp.zeros((n, self._m), dtype=jnp.dtype(self._dtype))
        rows = jnp.arange(n)[:, None]  # (n, 1) broadcast with (n, p)
        W = W.at[rows, self.interp_indices].add(self.interp_values)
        return W

MaskedOperator

Bases: AbstractLinearOperator

Row/column-masked view of a base operator.

Given a base operator A of shape (N, N) and boolean masks, produces the sub-matrix A[row_mask][:, col_mask].

Matvec is computed without materializing the sub-matrix: zero-pad input to full size, apply base matvec, then extract masked rows.

Fast masked solves (capacitance method). For a square mask (row_mask == col_mask == m) whose base B has a fast solve (a gaussx.DiagonalizedOperator, gaussx.circulant, or a Kronecker sum of them), pass coupling_indices: the flat indices C of the masked-out degrees of freedom that B couples to masked-in rows (B[m, j] ≠ 0). Then gaussx.solve(op, f) solves B[m][:, m] x = f exactly by the capacitance method: it finds y on the full grid with y[C] = 0 and (B y)[i] = f[i] at every i ∉ C (point sources on C absorb the constraints), which on m is exactly the masked system because the remaining masked-out entries do not couple to m. The |C| × |C| capacitance factorisation is built once, here in the constructor, so each solve costs two base solves.

For a local stencil (e.g. a finite-difference Laplacian on a grid) C is the one-cell ring outside the mask — use grid_coupling_indices, with periodic=True on axes where B wraps. Passing all masked-out indices is always exact (for any B) but costs O(|C|) base solves to build.

If B is singular (e.g. a periodic Poisson operator, whose solve is the pseudo-inverse), its null vector must enter the capacitance system (see gaussx.CapacitanceSolver). For a diagonalised B with a normal transform and exactly one zero eigenvalue this is derived automatically; otherwise pass null_vector (and left_null_vector if B is not symmetric).

Parameters:

Name	Type	Description	Default
base	AbstractLinearOperator	The underlying (N, N) linear operator.	required
row_mask	Bool[Array, ' N']	Boolean mask of length N selecting output rows.	required
col_mask	Bool[Array, ' N']	Boolean mask of length N selecting input columns.	required
coupling_indices	Int[Array, ' Nc'] | None	Optional flat indices C (see above); enables the capacitance solve. Requires row_mask == col_mask (concrete).	None
null_vector	Float[Array, ' N'] | None	Right null vector of a singular base; derived automatically for diagonalised bases when possible.	None
left_null_vector	Float[Array, ' N'] | None	Left null vector of a singular non-symmetric base; defaults to null_vector.	None

Attributes:

Name	Type	Description
capacitance	CapacitanceSolver | None	The precomputed gaussx.CapacitanceSolver on the full index space, or None when no coupling_indices were given (solve then falls back to the dense path). It stores no base_solve: solve builds one from base at call time, so base stays a differentiable pytree leaf and solve is differentiated implicitly (not through the cached factorisation).

The capacitance is a cache of base. Gradients with respect to base are exact, but replacing base in an existing operator (e.g. eqx.apply_updates after an optimiser step, or eqx.tree_at) leaves the factorisation describing the old base, and solve is then wrong. Rebuild with with_base, or construct the operator inside the loss from the current parameters. That also works under jit, except that a traced base's null vector cannot be derived: for a singular base built under jit pass null_vector (and left_null_vector) explicitly. A traced diagonalised base with a zero eigenvalue and no null_vector raises at run time rather than solving the wrong system.

Source code in src/gaussx/_operators/_masked.py

class MaskedOperator(lx.AbstractLinearOperator):
    """Row/column-masked view of a base operator.

    Given a base operator A of shape ``(N, N)`` and boolean masks,
    produces the sub-matrix ``A[row_mask][:, col_mask]``.

    Matvec is computed without materializing the sub-matrix:
    zero-pad input to full size, apply base matvec, then extract
    masked rows.

    **Fast masked solves (capacitance method).** For a square mask
    (``row_mask == col_mask == m``) whose base ``B`` has a fast solve (a
    `gaussx.DiagonalizedOperator`, `gaussx.circulant`, or a Kronecker sum of
    them), pass ``coupling_indices``: the flat indices ``C`` of the
    masked-*out* degrees of freedom that ``B`` couples to masked-in rows
    (``B[m, j] ≠ 0``). Then ``gaussx.solve(op, f)`` solves
    ``B[m][:, m] x = f`` exactly by the capacitance method: it finds ``y`` on
    the full grid with ``y[C] = 0`` and ``(B y)[i] = f[i]`` at every ``i ∉ C``
    (point sources on ``C`` absorb the constraints), which on ``m`` is exactly
    the masked system because the remaining masked-out entries do not couple
    to ``m``. The ``|C| × |C|`` capacitance factorisation is built once, here
    in the constructor, so each solve costs two base solves.

    For a local stencil (e.g. a finite-difference Laplacian on a grid) ``C``
    is the one-cell ring outside the mask — use `grid_coupling_indices`, with
    ``periodic=True`` on axes where ``B`` wraps. Passing *all* masked-out
    indices is always exact (for any ``B``) but costs ``O(|C|)`` base solves
    to build.

    If ``B`` is singular (e.g. a periodic Poisson operator, whose solve is the
    pseudo-inverse), its null vector must enter the capacitance system (see
    `gaussx.CapacitanceSolver`). For a diagonalised ``B`` with a normal
    transform and exactly one zero eigenvalue this is derived automatically;
    otherwise pass ``null_vector`` (and ``left_null_vector`` if ``B`` is not
    symmetric).

    Args:
        base: The underlying ``(N, N)`` linear operator.
        row_mask: Boolean mask of length N selecting output rows.
        col_mask: Boolean mask of length N selecting input columns.
        coupling_indices: Optional flat indices ``C`` (see above); enables the
            capacitance solve. Requires ``row_mask == col_mask`` (concrete).
        null_vector: Right null vector of a singular ``base``; derived
            automatically for diagonalised bases when possible.
        left_null_vector: Left null vector of a singular non-symmetric
            ``base``; defaults to ``null_vector``.

    Attributes:
        capacitance: The precomputed `gaussx.CapacitanceSolver` on the full
            index space, or ``None`` when no ``coupling_indices`` were given
            (``solve`` then falls back to the dense path). It stores no
            ``base_solve``: ``solve`` builds one from ``base`` at call time,
            so ``base`` stays a differentiable pytree leaf and ``solve`` is
            differentiated implicitly (not through the cached factorisation).

    **The capacitance is a cache of ``base``.** Gradients with respect to
    ``base`` are exact, but replacing ``base`` in an existing operator (e.g.
    ``eqx.apply_updates`` after an optimiser step, or ``eqx.tree_at``) leaves
    the factorisation describing the *old* base, and ``solve`` is then
    wrong. Rebuild with `with_base`, or construct the operator inside the
    loss from the current parameters. That also works under ``jit``, except
    that a traced base's null vector cannot be derived: for a singular base
    built under ``jit`` pass ``null_vector`` (and ``left_null_vector``)
    explicitly. A traced diagonalised base with a zero eigenvalue and no
    ``null_vector`` raises at run time rather than solving the wrong system.
    """

    base: lx.AbstractLinearOperator
    row_mask: Bool[Array, " N"]
    col_mask: Bool[Array, " N"]
    capacitance: CapacitanceSolver | None
    _in_size: int = eqx.field(static=True)
    _out_size: int = eqx.field(static=True)
    _dtype: str = eqx.field(static=True)
    tags: frozenset[object] = eqx.field(static=True)

    def __init__(
        self,
        base: lx.AbstractLinearOperator,
        row_mask: Bool[Array, " N"],
        col_mask: Bool[Array, " N"],
        *,
        tags: object | frozenset[object] = frozenset(),
        coupling_indices: Int[Array, " Nc"] | None = None,
        null_vector: Float[Array, " N"] | None = None,
        left_null_vector: Float[Array, " N"] | None = None,
    ) -> None:
        if base.in_size() != base.out_size():
            raise ValueError(
                f"Base operator must be square, got in_size={base.in_size()}, "
                f"out_size={base.out_size()}."
            )
        n = base.in_size()
        if row_mask.shape != (n,) or col_mask.shape != (n,):
            raise ValueError(
                f"Masks must have shape ({n},), got row_mask={row_mask.shape}, "
                f"col_mask={col_mask.shape}."
            )
        self.base = base
        self.row_mask = jnp.asarray(row_mask, dtype=bool)
        self.col_mask = jnp.asarray(col_mask, dtype=bool)
        # Counted with NumPy so constant masks also work inside ``jit``
        # (``jnp.sum`` would be staged); traced masks cannot give a static
        # size and raise.
        self._in_size = int(np.sum(np.asarray(col_mask)))
        self._out_size = int(np.sum(np.asarray(row_mask)))
        struct = base.out_structure()
        leaves = jax.tree.leaves(struct)
        self._dtype = str(leaves[0].dtype)
        self.tags = _to_frozenset(tags)
        self.capacitance = None
        if coupling_indices is not None:
            self.capacitance = _build_capacitance(
                base,
                self.row_mask,
                self.col_mask,
                jnp.asarray(coupling_indices),
                null_vector,
                left_null_vector,
            )

    def mv(self, vector: Float[Array, " m"]) -> Float[Array, " k"]:
        # Scatter input into full-size vector at col_mask positions
        n = self.base.in_size()
        col_indices = jnp.where(self.col_mask, size=self._in_size)[0]
        full_v = jnp.zeros(n, dtype=vector.dtype).at[col_indices].set(vector)
        # Apply base operator
        full_out = self.base.mv(full_v)
        # Gather output at row_mask positions
        row_indices = jnp.where(self.row_mask, size=self._out_size)[0]
        return full_out[row_indices]

    def as_matrix(self) -> Float[Array, "k m"]:
        full = self.base.as_matrix()
        row_indices = jnp.where(self.row_mask, size=self._out_size)[0]
        col_indices = jnp.where(self.col_mask, size=self._in_size)[0]
        return full[jnp.ix_(row_indices, col_indices)]

    def with_base(
        self,
        base: lx.AbstractLinearOperator,
        *,
        null_vector: Float[Array, " N"] | None = None,
        left_null_vector: Float[Array, " N"] | None = None,
    ) -> MaskedOperator:
        """The same masked view of a new base, with its capacitance rebuilt.

        Use this after updating ``base`` (e.g. an optimiser step): the
        capacitance factorisation is a cache of the base and does not follow
        leaf updates. Costs the ``|C|`` base solves of construction.

        Args:
            base: The new ``(N, N)`` base operator.
            null_vector: Right null vector of a singular ``base``. As in the
                constructor, it is derived automatically for a singular
                diagonalised base; the old one is *not* reused, since the
                update may change (or remove) the null space.
            left_null_vector: Left null vector of a singular non-symmetric
                ``base``; defaults to ``null_vector``.

        Returns:
            A new `MaskedOperator` over ``base`` with the same masks, tags and
            coupling indices.
        """
        if self.capacitance is None:
            return MaskedOperator(base, self.row_mask, self.col_mask, tags=self.tags)
        return MaskedOperator(
            base,
            self.row_mask,
            self.col_mask,
            tags=self.tags,
            coupling_indices=self.capacitance.boundary_indices,
            null_vector=null_vector,
            left_null_vector=left_null_vector,
        )

    def transpose(self) -> MaskedOperator:
        # Built field by field rather than through ``__init__``, which needs
        # concrete masks (so it can run under a trace, e.g. in lineax's solve
        # VJP). ``(B[m][:, m])ᵀ = (Bᵀ)[m][:, m]``, and the capacitance of
        # ``Bᵀ`` is the transposed factorisation of ``B``'s.
        return _replace_fields(
            self,
            base=self.base.T,
            row_mask=self.col_mask,
            col_mask=self.row_mask,
            capacitance=(
                None if self.capacitance is None else self.capacitance.transpose()
            ),
            _in_size=self._out_size,
            _out_size=self._in_size,
            tags=lx.transpose_tags(self.tags),
        )

    def in_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._in_size,), jnp.dtype(self._dtype))

    def out_structure(self) -> jax.ShapeDtypeStruct:
        return jax.ShapeDtypeStruct((self._out_size,), jnp.dtype(self._dtype))

with_base(base: lx.AbstractLinearOperator, *, null_vector: Float[Array, ' N'] | None = None, left_null_vector: Float[Array, ' N'] | None = None) -> MaskedOperator

The same masked view of a new base, with its capacitance rebuilt.

Use this after updating base (e.g. an optimiser step): the capacitance factorisation is a cache of the base and does not follow leaf updates. Costs the |C| base solves of construction.

Parameters:

Name	Type	Description	Default
base	AbstractLinearOperator	The new (N, N) base operator.	required
null_vector	Float[Array, ' N'] | None	Right null vector of a singular base. As in the constructor, it is derived automatically for a singular diagonalised base; the old one is not reused, since the update may change (or remove) the null space.	None
left_null_vector	Float[Array, ' N'] | None	Left null vector of a singular non-symmetric base; defaults to null_vector.	None

Returns:

Type	Description
MaskedOperator	A new MaskedOperator over base with the same masks, tags and
MaskedOperator	coupling indices.

Source code in src/gaussx/_operators/_masked.py

def with_base(
    self,
    base: lx.AbstractLinearOperator,
    *,
    null_vector: Float[Array, " N"] | None = None,
    left_null_vector: Float[Array, " N"] | None = None,
) -> MaskedOperator:
    """The same masked view of a new base, with its capacitance rebuilt.

    Use this after updating ``base`` (e.g. an optimiser step): the
    capacitance factorisation is a cache of the base and does not follow
    leaf updates. Costs the ``|C|`` base solves of construction.

    Args:
        base: The new ``(N, N)`` base operator.
        null_vector: Right null vector of a singular ``base``. As in the
            constructor, it is derived automatically for a singular
            diagonalised base; the old one is *not* reused, since the
            update may change (or remove) the null space.
        left_null_vector: Left null vector of a singular non-symmetric
            ``base``; defaults to ``null_vector``.

    Returns:
        A new `MaskedOperator` over ``base`` with the same masks, tags and
        coupling indices.
    """
    if self.capacitance is None:
        return MaskedOperator(base, self.row_mask, self.col_mask, tags=self.tags)
    return MaskedOperator(
        base,
        self.row_mask,
        self.col_mask,
        tags=self.tags,
        coupling_indices=self.capacitance.boundary_indices,
        null_vector=null_vector,
        left_null_vector=left_null_vector,
    )

grid_coupling_indices(mask: Bool[Array, ' *grid'], *, periodic: bool | tuple[bool, ...] = False, connectivity: int = 1) -> Int[Array, ' Nc']

Flat indices of masked-out grid cells adjacent to the masked-in region.

This is the coupling set C that MaskedOperator needs for a local stencil operator on a regular grid: every cell outside mask that the stencil of some cell inside mask reaches. With connectivity=1 neighbours differ by ±1 along one axis (the 5-point / 7-point stencil); connectivity=mask.ndim includes diagonals (9-point / 27-point).

Parameters:

Name	Type	Description	Default
mask	Bool[Array, ' *grid']	Boolean grid, True for the unknowns (masked-in cells).	required
periodic	bool | tuple[bool, ...]	Whether each axis wraps around (a bool applies to all axes). Must match the base operator: a periodic operator couples cells across the edge, so an edge-touching mask needs the wrapped neighbours in C.	False
connectivity	int	Maximum number of axes a neighbour offset may change.	1

Returns:

Type	Description
Int[Array, ' Nc']	Flat (row-major) indices of the coupled masked-out cells.

Source code in src/gaussx/_operators/_masked.py

def grid_coupling_indices(
    mask: Bool[Array, " *grid"],
    *,
    periodic: bool | tuple[bool, ...] = False,
    connectivity: int = 1,
) -> Int[Array, " Nc"]:
    """Flat indices of masked-out grid cells adjacent to the masked-in region.

    This is the coupling set ``C`` that `MaskedOperator` needs for a local
    stencil operator on a regular grid: every cell outside ``mask`` that the
    stencil of some cell inside ``mask`` reaches. With ``connectivity=1``
    neighbours differ by ±1 along one axis (the 5-point / 7-point stencil);
    ``connectivity=mask.ndim`` includes diagonals (9-point / 27-point).

    Args:
        mask: Boolean grid, ``True`` for the unknowns (masked-in cells).
        periodic: Whether each axis wraps around (a bool applies to all
            axes). Must match the base operator: a periodic operator couples
            cells across the edge, so an edge-touching mask needs the wrapped
            neighbours in ``C``.
        connectivity: Maximum number of axes a neighbour offset may change.

    Returns:
        Flat (row-major) indices of the coupled masked-out cells.
    """
    m = np.asarray(mask, dtype=bool)
    ndim = m.ndim
    wrap = (periodic,) * ndim if isinstance(periodic, bool) else tuple(periodic)
    if len(wrap) != ndim:
        raise ValueError(f"periodic has {len(wrap)} entries for a {ndim}-D mask.")
    if not 1 <= connectivity <= ndim:
        raise ValueError(f"connectivity must be in [1, {ndim}], got {connectivity}.")
    reach = np.zeros_like(m)
    for offset in itertools.product((-1, 0, 1), repeat=ndim):
        if 0 < np.count_nonzero(offset) <= connectivity:
            reach |= _shift(m, offset, wrap)
    return jnp.asarray(np.flatnonzero((reach & ~m).ravel()))

create_grid(grid_sizes: list[int] | tuple[int, ...], grid_bounds: list[tuple[float, float]] | tuple[tuple[float, float], ...]) -> list[Float[Array, ' n']]

Create a regular grid from per-dimension sizes and bounds.

Parameters:

Name	Type	Description	Default
grid_sizes	list[int] | tuple[int, ...]	Number of points per dimension, length D.	required
grid_bounds	list[tuple[float, float]] | tuple[tuple[float, float], ...]	(lo, hi) bounds per dimension, length D.	required

Returns:

Type	Description
list[Float[Array, ' n']]	List of 1-D arrays, one per dimension.

Source code in src/gaussx/_operators/_grid.py

def create_grid(
    grid_sizes: list[int] | tuple[int, ...],
    grid_bounds: list[tuple[float, float]] | tuple[tuple[float, float], ...],
) -> list[Float[Array, " n"]]:
    """Create a regular grid from per-dimension sizes and bounds.

    Args:
        grid_sizes: Number of points per dimension, length D.
        grid_bounds: (lo, hi) bounds per dimension, length D.

    Returns:
        List of 1-D arrays, one per dimension.
    """
    return [
        jnp.linspace(lo, hi, n)
        for n, (lo, hi) in zip(grid_sizes, grid_bounds, strict=True)
    ]

grid_data(grid: list[Float[Array, ' n']]) -> Float[Array, 'G D']

Expand a grid to the full Cartesian product of grid points.

Parameters:

Name	Type	Description	Default
grid	list[Float[Array, ' n']]	List of 1-D arrays, one per dimension.	required

Returns:

Type	Description
Float[Array, 'G D']	Shape (prod(sizes), D) array of all grid points.

Source code in src/gaussx/_operators/_grid.py

def grid_data(grid: list[Float[Array, " n"]]) -> Float[Array, "G D"]:
    """Expand a grid to the full Cartesian product of grid points.

    Args:
        grid: List of 1-D arrays, one per dimension.

    Returns:
        Shape ``(prod(sizes), D)`` array of all grid points.
    """
    meshes = jnp.meshgrid(*grid, indexing="ij")  # D arrays of shape sizes
    stacked = jnp.stack(meshes, axis=0)  # (D, *sizes)
    return rearrange(stacked, "D ... -> (...) D")

cubic_interpolation_weights(x_target: Float[Array, 'B D'], grid: list[Float[Array, ' n']]) -> tuple[Int[Array, 'B K'], Float[Array, 'B K']]

Compute cubic interpolation indices and weights for SKI.

For each target point, finds the 4ᴰ nearest grid neighbors and computes cubic convolution interpolation weights.

Parameters:

Name	Type	Description	Default
x_target	Float[Array, 'B D']	Target points, shape (B, D).	required
grid	list[Float[Array, ' n']]	List of 1-D arrays, one per dimension (length D). Each dimension must have ≥ 4 points.	required

Returns:

Type	Description
Int[Array, 'B K']	Tuple (indices, weights) both of shape (B, 4ᴰ):
Float[Array, 'B K']	indices: Flat indices into the grid (product of sizes).
tuple[Int[Array, 'B K'], Float[Array, 'B K']]	weights: Interpolation weights summing to ≈ 1.

Raises:

Type	Description
ValueError	If any grid dimension has fewer than 4 points.

Source code in src/gaussx/_operators/_grid.py

def cubic_interpolation_weights(
    x_target: Float[Array, "B D"],
    grid: list[Float[Array, " n"]],
) -> tuple[Int[Array, "B K"], Float[Array, "B K"]]:
    """Compute cubic interpolation indices and weights for SKI.

    For each target point, finds the 4ᴰ nearest grid neighbors and
    computes cubic convolution interpolation weights.

    Args:
        x_target: Target points, shape ``(B, D)``.
        grid: List of 1-D arrays, one per dimension (length D).
            Each dimension must have ≥ 4 points.

    Returns:
        Tuple ``(indices, weights)`` both of shape ``(B, 4ᴰ)``:

        - ``indices``: Flat indices into the grid (product of sizes).
        - ``weights``: Interpolation weights summing to ≈ 1.

    Raises:
        ValueError: If any grid dimension has fewer than 4 points.
    """
    D = len(grid)
    grid_sizes = [g.shape[0] for g in grid]

    for d, n in enumerate(grid_sizes):
        if n < 4:
            msg = (
                f"Cubic interpolation requires at least 4 grid points per "
                f"dimension, but dimension {d} has {n}."
            )
            raise ValueError(msg)

    # Strides for manual flat indexing: strides[d] = ∏ sizes[d+1:]
    strides = []
    s = 1
    for d in range(D - 1, -1, -1):
        strides.append(s)
        s *= grid_sizes[d]
    strides = strides[::-1]

    # Per-dimension: 4 indices and weights
    dim_indices = []  # each (B, 4)
    dim_weights = []  # each (B, 4)

    for d in range(D):
        g = grid[d]
        n = g.shape[0]
        x_d = x_target[:, d]  # (B,)

        h = (g[-1] - g[0]) / (n - 1)  # cell spacing
        cont_idx = (x_d - g[0]) / h  # continuous index, (B,)

        # Integer cell index, clamped for 4 neighbors
        cell = jnp.floor(cont_idx).astype(jnp.int32)
        cell = jnp.clip(cell, 1, n - 3)  # (B,)

        t = cont_idx - cell.astype(x_d.dtype)  # fractional part, (B,)

        # 4 neighbor indices: cell−1, cell, cell+1, cell+2
        idx = jnp.stack([cell - 1, cell, cell + 1, cell + 2], axis=-1)  # (B, 4)
        idx = jnp.clip(idx, 0, n - 1)

        dim_indices.append(idx)
        dim_weights.append(_cubic_weights_1d(t))

    # Cross-dimensional outer product of indices/weights
    flat_indices = dim_indices[0] * strides[0]  # (B, 4)
    weights = dim_weights[0]  # (B, 4)

    for d in range(1, D):
        # Outer product: (B, K) × (B, 4) → (B, K·4)
        fi_exp = flat_indices[:, :, None]  # (B, K, 1)
        di_exp = dim_indices[d][:, None, :] * strides[d]  # (B, 1, 4)
        flat_indices = rearrange(fi_exp + di_exp, "B K four -> B (K four)")

        w_exp = weights[:, :, None]  # (B, K, 1)
        dw_exp = dim_weights[d][:, None, :]  # (B, 1, 4)
        weights = rearrange(w_exp * dw_exp, "B K four -> B (K four)")

    return flat_indices, weights

Lazy algebra & sampling

Sum / scale / compose operators without materializing, sample \(\varepsilon \sim \mathcal{N}(0, A)\) for the structured families, and solve bordered systems through the capacitance (Schur-complement) form.

Structured linear algebra and Gaussian primitives for JAX.

CapacitanceSolver

Bases: Module

Solve a base system subject to homogeneous point constraints.

Given a fast base solver B^{-1} and a set of N_b constrained indices b_1, …, b_{N_b}, this returns x with x[b_k] = 0 and B x = f at every unconstrained index, via the capacitance-matrix correction:

1. Base solve: u = B^{-1} f
2. Sample boundary: u_b = u[boundary]
3. Correction: α = C^{-1} u_b
4. Subtract: x = u − B^{-1}(Σ_k α_k e_{b_k})

where C[k, l] = (B^{-1} e_{b_l})[b_k] is the N_b × N_b capacitance matrix, LU-factorized once at construction.

Cost and storage: only the factorized C is stored — O(N_b²) memory. The Green's functions B^{-1} e_{b_k} are never kept; step 4 is a second base solve with the point sources scattered into a zero field. Each call therefore costs two base solves plus an O(N_b²) back-substitution. Construction costs N_b base solves (run one column at a time under jit, so no dense (N_b, n) array is formed).

Singular base operators. If B is singular with right null vector r (e.g. the constant for a periodic or Neumann Poisson operator) and base_solve is the pseudo-inverse that projects out that mode, the plain method leaves the null-mode component undetermined and the interior residual is a nonzero constant. Pass null_vector=r (and, for a non-symmetric B, left_null_vector=w with wᵀ B = 0; it defaults to r). The solution is then x = u − B^{+}(Σ α_k e_{b_k}) + β r with (α, β) from the bordered system

[ C −r_b ] [α] [ u_b ] [ w_bᵀ 0 ] [β] = [ wᵀ f ]

whose second row makes the effective source f − Σ α_k e_{b_k} orthogonal to the left null space, so that B x = f holds exactly away from the constrained indices. Without null_vector a singular base cannot be detected here, so callers must pass it.

The solver operates on flat vectors. Any reshaping between fields and flat vectors, and any masking of the exterior, is the caller's responsibility -- keeping grid/mask concepts out of this class.

Parameters:

Name	Type	Description	Default
base_solve	Callable[[Float[Array, ' n']], Float[Array, ' n']]	Callable applying the base inverse B^{-1} (or the pseudo-inverse B^{+} for a singular B) to a flat right-hand side of length n. Stored as a static field, so it must be hashable (functions and closures are).	required
boundary_indices	Int[Array, ' Nb']	Flat indices of the constrained degrees of freedom, shape (N_b,).	required
n	int	Length of the flat solution vector.	required
null_vector	Float[Array, ' n'] | None	Right null vector r of a singular base operator, shape (n,); None (default) for a non-singular base.	None
left_null_vector	Float[Array, ' n'] | None	Left null vector w of a singular base operator; defaults to null_vector (correct for symmetric B). Ignored when null_vector is None.	None
keep_base_solve	bool	When False, base_solve is only used to build the capacitance matrix and is not stored; pass it to __call__ instead. A static closure over an operator hides the operator's arrays from autodiff and makes every instance a new jit cache key, so gaussx.MaskedOperator keeps its base as a pytree leaf and uses this.	True

Attributes:

Name	Type	Description
base_solve	Callable[[Float[Array, ' n']], Float[Array, ' n']] | None	The base inverse callable (static), or None when built with keep_base_solve=False.
boundary_indices	Int[Array, ' Nb']	The constrained indices.
capacitance_lu	tuple[Float[Array, 'm m'], Int[Array, ' m']]	LU factorization (lu, pivots) of C, or of the bordered (N_b + 1) × (N_b + 1) system when a null vector is given.
null_vector	Float[Array, ' n'] | None	r, or None.
left_null_vector	Float[Array, ' n'] | None	w, or None.
transposed	bool	Whether this solver is the transpose of a built one (it then solves with Bᵀ using the transposed factorisation).

Source code in src/gaussx/_operators/_capacitance.py

class CapacitanceSolver(eqx.Module):
    r"""Solve a base system subject to homogeneous point constraints.

    Given a fast base solver ``B^{-1}`` and a set of ``N_b`` constrained indices
    ``b_1, …, b_{N_b}``, this returns ``x`` with ``x[b_k] = 0`` and
    ``B x = f`` at every unconstrained index, via the capacitance-matrix
    correction:

    1. Base solve:        ``u = B^{-1} f``
    2. Sample boundary:   ``u_b = u[boundary]``
    3. Correction:        ``α = C^{-1} u_b``
    4. Subtract:          ``x = u − B^{-1}(Σ_k α_k e_{b_k})``

    where ``C[k, l] = (B^{-1} e_{b_l})[b_k]`` is the ``N_b × N_b`` capacitance
    matrix, LU-factorized once at construction.

    Cost and storage: only the factorized ``C`` is stored — ``O(N_b²)``
    memory. The Green's functions ``B^{-1} e_{b_k}`` are never kept; step 4 is
    a second base solve with the point sources scattered into a zero field.
    Each call therefore costs two base solves plus an ``O(N_b²)``
    back-substitution. Construction costs ``N_b`` base solves (run one column
    at a time under ``jit``, so no dense ``(N_b, n)`` array is formed).

    **Singular base operators.** If ``B`` is singular with right null vector
    ``r`` (e.g. the constant for a periodic or Neumann Poisson operator) and
    ``base_solve`` is the pseudo-inverse that projects out that mode, the plain
    method leaves the null-mode component undetermined and the interior
    residual is a nonzero constant. Pass ``null_vector=r`` (and, for a
    non-symmetric ``B``, ``left_null_vector=w`` with ``wᵀ B = 0``; it defaults
    to ``r``). The solution is then ``x = u − B^{+}(Σ α_k e_{b_k}) + β r`` with
    ``(α, β)`` from the bordered system

    ``[ C      −r_b ] [α]   [ u_b  ]``
    ``[ w_bᵀ    0   ] [β] = [ wᵀ f ]``

    whose second row makes the effective source ``f − Σ α_k e_{b_k}``
    orthogonal to the left null space, so that ``B x = f`` holds exactly away
    from the constrained indices. Without ``null_vector`` a singular base
    cannot be detected here, so callers must pass it.

    The solver operates on **flat** vectors. Any reshaping between fields and
    flat vectors, and any masking of the exterior, is the caller's
    responsibility -- keeping grid/mask concepts out of this class.

    Args:
        base_solve: Callable applying the base inverse ``B^{-1}`` (or the
            pseudo-inverse ``B^{+}`` for a singular ``B``) to a flat
            right-hand side of length ``n``. Stored as a static field, so it
            must be hashable (functions and closures are).
        boundary_indices: Flat indices of the constrained degrees of freedom,
            shape ``(N_b,)``.
        n: Length of the flat solution vector.
        null_vector: Right null vector ``r`` of a singular base operator,
            shape ``(n,)``; ``None`` (default) for a non-singular base.
        left_null_vector: Left null vector ``w`` of a singular base operator;
            defaults to ``null_vector`` (correct for symmetric ``B``). Ignored
            when ``null_vector`` is ``None``.
        keep_base_solve: When ``False``, ``base_solve`` is only used to build
            the capacitance matrix and is not stored; pass it to `__call__`
            instead. A static closure over an operator hides the operator's
            arrays from autodiff and makes every instance a new ``jit``
            cache key, so `gaussx.MaskedOperator` keeps its base as a pytree
            leaf and uses this.

    Attributes:
        base_solve: The base inverse callable (static), or ``None`` when
            built with ``keep_base_solve=False``.
        boundary_indices: The constrained indices.
        capacitance_lu: LU factorization ``(lu, pivots)`` of ``C``, or of the
            bordered ``(N_b + 1) × (N_b + 1)`` system when a null vector is
            given.
        null_vector: ``r``, or ``None``.
        left_null_vector: ``w``, or ``None``.
        transposed: Whether this solver is the `transpose` of a built one
            (it then solves with ``Bᵀ`` using the transposed factorisation).
    """

    base_solve: Callable[[Float[Array, " n"]], Float[Array, " n"]] | None = eqx.field(
        static=True
    )
    boundary_indices: Int[Array, " Nb"]
    capacitance_lu: tuple[Float[Array, "m m"], Int[Array, " m"]]
    null_vector: Float[Array, " n"] | None
    left_null_vector: Float[Array, " n"] | None
    transposed: bool = eqx.field(static=True)

    def __init__(
        self,
        base_solve: Callable[[Float[Array, " n"]], Float[Array, " n"]],
        boundary_indices: Int[Array, " Nb"],
        n: int,
        *,
        null_vector: Float[Array, " n"] | None = None,
        left_null_vector: Float[Array, " n"] | None = None,
        keep_base_solve: bool = True,
    ):
        indices = jnp.asarray(boundary_indices)
        capacitance = _capacitance_matrix(base_solve, indices, n)

        if null_vector is None:
            system = capacitance
            left_null_vector = None
        else:
            null_vector = jnp.asarray(null_vector)
            left_null_vector = (
                null_vector
                if left_null_vector is None
                else jnp.asarray(left_null_vector)
            )
            r_b = null_vector[indices]
            l_b = left_null_vector[indices]
            system = jnp.block(
                [
                    [capacitance, -r_b[:, None]],
                    [l_b[None, :], jnp.zeros((1, 1), dtype=capacitance.dtype)],
                ]
            )

        self.base_solve = base_solve if keep_base_solve else None
        self.boundary_indices = indices
        self.capacitance_lu = jsl.lu_factor(system)
        self.null_vector = null_vector
        self.left_null_vector = left_null_vector
        self.transposed = False

    def __call__(
        self,
        rhs: Float[Array, " n"],
        *,
        base_solve: Callable[[Float[Array, " n"]], Float[Array, " n"]] | None = None,
    ) -> Float[Array, " n"]:
        """Solve the constrained system for a flat right-hand side ``rhs``.

        Args:
            rhs: Flat right-hand side, shape ``(n,)``.
            base_solve: The base inverse to use; overrides the stored one
                and is required when the solver was built with
                ``keep_base_solve=False``. For a solver returned by
                `transpose` it must apply ``B^{-T}``.

        Returns:
            The constrained solution, shape ``(n,)``.
        """
        base_solve = self.base_solve if base_solve is None else base_solve
        if base_solve is None:
            raise ValueError(
                "This CapacitanceSolver stores no base_solve; pass base_solve=."
            )
        # A transposed solver solves with the transposed factorisation. The
        # bordered system for Bᵀ is D Sᵀ D with D = diag(1, …, 1, -1).
        trans = 1 if self.transposed else 0
        u = base_solve(rhs)
        u_b = u[self.boundary_indices]
        if self.null_vector is None:
            alpha = jsl.lu_solve(self.capacitance_lu, u_b, trans=trans)
            beta = None
        else:
            assert self.left_null_vector is not None
            load = jnp.dot(self.left_null_vector, rhs)
            sign = -1.0 if self.transposed else 1.0
            sol = jsl.lu_solve(
                self.capacitance_lu, jnp.append(u_b, sign * load), trans=trans
            )
            alpha, beta = sol[:-1], sign * sol[-1]
        # Constrained indices are distinct; saying so keeps the scatter
        # transposable (reverse mode through a solve).
        sources = (
            jnp.zeros_like(u).at[self.boundary_indices].set(alpha, unique_indices=True)
        )
        x = u - base_solve(sources)
        if beta is not None:
            assert self.null_vector is not None
            x = x + beta * self.null_vector
        return x

    def transpose(self) -> CapacitanceSolver:
        """The capacitance solver for ``Bᵀ``, reusing this factorisation.

        ``Cᵀ`` is the capacitance matrix of ``Bᵀ`` on the same indices, and
        ``Bᵀ``'s right/left null vectors are ``B``'s left/right ones, so no
        base solve is needed. The result stores no ``base_solve``: pass
        ``Bᵀ``'s inverse to `__call__`.
        """
        return _replace_fields(
            self,
            base_solve=None,
            null_vector=self.left_null_vector,
            left_null_vector=self.null_vector,
            transposed=not self.transposed,
        )

transpose() -> CapacitanceSolver

The capacitance solver for Bᵀ, reusing this factorisation.

Cᵀ is the capacitance matrix of Bᵀ on the same indices, and Bᵀ's right/left null vectors are B's left/right ones, so no base solve is needed. The result stores no base_solve: pass Bᵀ's inverse to __call__.

Source code in src/gaussx/_operators/_capacitance.py

def transpose(self) -> CapacitanceSolver:
    """The capacitance solver for ``Bᵀ``, reusing this factorisation.

    ``Cᵀ`` is the capacitance matrix of ``Bᵀ`` on the same indices, and
    ``Bᵀ``'s right/left null vectors are ``B``'s left/right ones, so no
    base solve is needed. The result stores no ``base_solve``: pass
    ``Bᵀ``'s inverse to `__call__`.
    """
    return _replace_fields(
        self,
        base_solve=None,
        null_vector=self.left_null_vector,
        left_null_vector=self.null_vector,
        transposed=not self.transposed,
    )

sum_operator(*operators: lx.AbstractLinearOperator, tags: object | frozenset[object] = frozenset()) -> lx.AbstractLinearOperator

Lazy sum (A + B + …) v = A v + B v + ….

Defers materialization so that structured sub-operators keep their efficient matvec. All operators must have the same input and output sizes. Returns a (possibly tagged) chain of lineax AddLinearOperator compositions.

Parameters:

Name	Type	Description	Default
*operators	AbstractLinearOperator	Two or more lineax.AbstractLinearOperator instances with matching shapes.	()
tags	object | frozenset[object]	Optional explicit lineax tags for the combined operator.	frozenset()

Returns:

Type	Description
AbstractLinearOperator	The lazy sum as a lineax operator.

Source code in src/gaussx/_operators/_lazy_algebra.py

def sum_operator(
    *operators: lx.AbstractLinearOperator,
    tags: object | frozenset[object] = frozenset(),
) -> lx.AbstractLinearOperator:
    """Lazy sum ``(A + B + …) v = A v + B v + …``.

    Defers materialization so that structured sub-operators keep their
    efficient matvec. All operators must have the same input and output
    sizes. Returns a (possibly tagged) chain of lineax
    ``AddLinearOperator`` compositions.

    Args:
        *operators: Two or more ``lineax.AbstractLinearOperator`` instances
            with matching shapes.
        tags: Optional explicit lineax tags for the combined operator.

    Returns:
        The lazy sum as a lineax operator.
    """
    if len(operators) < 2:
        raise ValueError("sum_operator requires at least two operators.")
    in0 = operators[0].in_size()
    out0 = operators[0].out_size()
    for i, op in enumerate(operators[1:], 1):
        if op.in_size() != in0 or op.out_size() != out0:
            raise ValueError(
                f"Shape mismatch: operator 0 has shape ({out0}, {in0}) "
                f"but operator {i} has shape ({op.out_size()}, {op.in_size()})."
            )
    return _maybe_tag(ft.reduce(_op.add, operators), tags)

scaled_operator(operator: lx.AbstractLinearOperator, scalar: float | Float[Array, ''], *, tags: object | frozenset[object] = frozenset()) -> lx.AbstractLinearOperator

Lazy scalar multiply (c A) v = c (A v).

Returns a (possibly tagged) lineax MulLinearOperator.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A lineax.AbstractLinearOperator.	required
scalar	float | Float[Array, '']	A scalar multiplier.	required
tags	object | frozenset[object]	Optional explicit lineax tags for the scaled operator.	frozenset()

Returns:

Type	Description
AbstractLinearOperator	The lazy scaled operator.

Source code in src/gaussx/_operators/_lazy_algebra.py

def scaled_operator(
    operator: lx.AbstractLinearOperator,
    scalar: float | Float[Array, ""],
    *,
    tags: object | frozenset[object] = frozenset(),
) -> lx.AbstractLinearOperator:
    """Lazy scalar multiply ``(c A) v = c (A v)``.

    Returns a (possibly tagged) lineax ``MulLinearOperator``.

    Args:
        operator: A ``lineax.AbstractLinearOperator``.
        scalar: A scalar multiplier.
        tags: Optional explicit lineax tags for the scaled operator.

    Returns:
        The lazy scaled operator.
    """
    scalar_array = jnp.asarray(scalar)
    if scalar_array.ndim != 0:
        msg = (
            "scaled_operator scalar must be a rank-0 scalar, got "
            f"shape {scalar_array.shape}."
        )
        raise ValueError(msg)
    return _maybe_tag(scalar_array * operator, tags)

product_operator(left: lx.AbstractLinearOperator, right: lx.AbstractLinearOperator, *, tags: object | frozenset[object] = frozenset()) -> lx.AbstractLinearOperator

Lazy matmul (A B) v = A (B v).

The inner dimension must match: left.in_size() == right.out_size(). Returns a (possibly tagged) lineax ComposedLinearOperator.

Parameters:

Name	Type	Description	Default
left	AbstractLinearOperator	The left operator A.	required
right	AbstractLinearOperator	The right operator B.	required
tags	object | frozenset[object]	Optional explicit lineax tags for the composed operator.	frozenset()

Returns:

Type	Description
AbstractLinearOperator	The lazy product as a lineax operator.

Source code in src/gaussx/_operators/_lazy_algebra.py

def product_operator(
    left: lx.AbstractLinearOperator,
    right: lx.AbstractLinearOperator,
    *,
    tags: object | frozenset[object] = frozenset(),
) -> lx.AbstractLinearOperator:
    """Lazy matmul ``(A B) v = A (B v)``.

    The inner dimension must match: ``left.in_size() == right.out_size()``.
    Returns a (possibly tagged) lineax ``ComposedLinearOperator``.

    Args:
        left: The left operator A.
        right: The right operator B.
        tags: Optional explicit lineax tags for the composed operator.

    Returns:
        The lazy product as a lineax operator.
    """
    if left.in_size() != right.out_size():
        raise ValueError(
            f"Inner dimension mismatch: left.in_size()={left.in_size()} "
            f"!= right.out_size()={right.out_size()}."
        )
    return _maybe_tag(left @ right, tags)

kronecker_sum_sample(A_op: lx.AbstractLinearOperator, B_op: lx.AbstractLinearOperator, *, key: jax.Array, num_samples: int = 1) -> Float[Array, 'num_samples n_a n_b']

Sample from 𝒩(0, A ⊕ B) using per-factor eigendecompositions.

Draws zero-mean samples with covariance A ⊕ B by applying the matrix-free KroneckerSumSqrt to standard normal noise, avoiding materialization of the full (n_a · n_b, n_a · n_b) covariance.

Parameters:

Name	Type	Description	Default
A_op	AbstractLinearOperator	Symmetric PSD factor, shape (n_a, n_a).	required
B_op	AbstractLinearOperator	Symmetric PSD factor, shape (n_b, n_b).	required
key	Array	PRNG key for the standard normal draws.	required
num_samples	int	Number of samples to draw.	1

Returns:

Type	Description
Float[Array, 'num_samples n_a n_b']	Samples of shape (num_samples, n_a, n_b).

Raises:

Type	Description
ValueError	If num_samples is less than 1.

Source code in src/gaussx/_operators/_kronecker_sum.py

def kronecker_sum_sample(
    A_op: lx.AbstractLinearOperator,
    B_op: lx.AbstractLinearOperator,
    *,
    key: jax.Array,
    num_samples: int = 1,
) -> Float[Array, "num_samples n_a n_b"]:
    """Sample from ``𝒩(0, A ⊕ B)`` using per-factor eigendecompositions.

    Draws zero-mean samples with covariance ``A ⊕ B`` by applying the
    matrix-free `KroneckerSumSqrt` to standard normal noise, avoiding
    materialization of the full ``(n_a · n_b, n_a · n_b)`` covariance.

    Args:
        A_op: Symmetric PSD factor, shape ``(n_a, n_a)``.
        B_op: Symmetric PSD factor, shape ``(n_b, n_b)``.
        key: PRNG key for the standard normal draws.
        num_samples: Number of samples to draw.

    Returns:
        Samples of shape ``(num_samples, n_a, n_b)``.

    Raises:
        ValueError: If ``num_samples`` is less than 1.
    """
    if num_samples <= 0:
        raise ValueError(f"num_samples must be at least 1, got {num_samples}.")

    sqrt_op = KroneckerSumSqrt(A_op, B_op)
    eps = jax.random.normal(
        key,
        (num_samples, sqrt_op.in_size()),
        dtype=jnp.dtype(sqrt_op._dtype),
    )
    samples = jax.vmap(sqrt_op.mv)(eps)
    return rearrange(samples, "s (a b) -> s a b", a=sqrt_op._n_a, b=sqrt_op._n_b)

sum_of_kroneckers_sample(op: SumOfKroneckers, *, key: jax.Array, num_samples: int = 1, lanczos_order: int = 50) -> Float[Array, 'num_samples n']

Sample from 𝒩(0, op) with matrix-free Lanczos square roots.

The square-root action is evaluated by matfree Lanczos against op.mv. This avoids materialising the dense (n_A n_B) × (n_A n_B) covariance and costs lanczos_order operator matvecs per sample.

Parameters:

Name	Type	Description	Default
op	SumOfKroneckers	Positive-semidefinite SumOfKroneckers covariance operator.	required
key	Array	JAX PRNG key.	required
num_samples	int	Number of independent samples to draw.	1
lanczos_order	int	Lanczos truncation order.	50

Returns:

Type	Description
Float[Array, 'num_samples n']	Samples with shape (num_samples, op.in_size()).

Source code in src/gaussx/_operators/_sum_kronecker.py

def sum_of_kroneckers_sample(
    op: SumOfKroneckers,
    *,
    key: jax.Array,
    num_samples: int = 1,
    lanczos_order: int = 50,
) -> Float[Array, "num_samples n"]:
    r"""Sample from ``𝒩(0, op)`` with matrix-free Lanczos square roots.

    The square-root action is evaluated by ``matfree`` Lanczos against
    ``op.mv``. This avoids materialising the dense ``(n_A n_B) ×
    (n_A n_B)`` covariance and costs ``lanczos_order`` operator
    matvecs per sample.

    Args:
        op: Positive-semidefinite `SumOfKroneckers` covariance operator.
        key: JAX PRNG key.
        num_samples: Number of independent samples to draw.
        lanczos_order: Lanczos truncation order.

    Returns:
        Samples with shape ``(num_samples, op.in_size())``.
    """
    # lazy import, cycle: _primitives._sqrt -> _operators._sum_kronecker
    from gaussx._primitives._sqrt import sqrt

    if op.in_size() != op.out_size():
        raise ValueError(
            "sum_of_kroneckers_sample requires a square SumOfKroneckers, got "
            f"in_size={op.in_size()} and out_size={op.out_size()}."
        )
    if num_samples < 1:
        raise ValueError(f"num_samples must be at least 1, got {num_samples}.")

    sqrt_op = sqrt(op, lanczos_order=lanczos_order)
    eps = jax.random.normal(
        key, (num_samples, op.in_size()), dtype=op.in_structure().dtype
    )
    return jax.vmap(sqrt_op.mv)(eps)

toeplitz_sample(column: Float[Array, ' n'], *, key: jax.Array, num_samples: int = 1, embedding_factor: int = 2) -> Float[Array, 'num_samples n']

Sample from 𝒩(0, Toeplitz(column)) via FFT circulant embedding.

The Wood--Chan non-negativity check is JIT-friendly (via eqx.error_if). If the embedding fails for the given embedding_factor, the error fires at evaluation time — bump embedding_factor (typically 4 or 8) for difficult covariances rather than relying on a runtime fallback.

Parameters:

Name	Type	Description	Default
column	Float[Array, ' n']	First column of the covariance matrix.	required
key	Array	JAX PRNG key used to draw white noise.	required
num_samples	int	Number of independent samples to draw.	1
embedding_factor	int	Circulant embedding size as a multiple of n.	2

Returns:

Type	Description
Float[Array, 'num_samples n']	Samples with shape (num_samples, n).

Source code in src/gaussx/_operators/_toeplitz.py

def toeplitz_sample(
    column: Float[Array, " n"],
    *,
    key: jax.Array,
    num_samples: int = 1,
    embedding_factor: int = 2,
) -> Float[Array, "num_samples n"]:
    """Sample from ``𝒩(0, Toeplitz(column))`` via FFT circulant embedding.

    The Wood--Chan non-negativity check is JIT-friendly (via
    ``eqx.error_if``). If the embedding fails for the given
    ``embedding_factor``, the error fires at evaluation time — bump
    ``embedding_factor`` (typically ``4`` or ``8``) for difficult
    covariances rather than relying on a runtime fallback.

    Args:
        column: First column of the covariance matrix.
        key: JAX PRNG key used to draw white noise.
        num_samples: Number of independent samples to draw.
        embedding_factor: Circulant embedding size as a multiple of ``n``.

    Returns:
        Samples with shape ``(num_samples, n)``.
    """
    if num_samples < 1:
        raise ValueError("num_samples must be at least 1.")

    column = _as_floating_column(column)
    factor = ToeplitzCholesky(column, embedding_factor=embedding_factor)
    noise = jr.normal(key, (num_samples, factor.in_size()), dtype=column.dtype)
    return jax.vmap(factor.mv)(noise)

Structural tags & predicates

Tags mark structure and properties on operators; the is_* predicates are what the primitives consult when choosing an algorithm. The property tags (positive_semidefinite_tag, symmetric_tag, the triangular tags, …) are re-exported from lineax so user code only needs one import. is_eigen_reducible reports whether a SumOfKroneckers (or the equivalent sum_operator chain) takes the exact two-term reduction.

Structured linear algebra and Gaussian primitives for JAX.

is_diagonal = lx.is_diagonal module-attribute

is_symmetric = lx.is_symmetric module-attribute

is_positive_semidefinite = lx.is_positive_semidefinite module-attribute

is_negative_semidefinite = lx.is_negative_semidefinite module-attribute

is_lower_triangular = lx.is_lower_triangular module-attribute

is_upper_triangular = lx.is_upper_triangular module-attribute

kronecker_tag = _Tag('kronecker_tag') module-attribute

Operator is a Kronecker product.

kronecker_sum_tag = _Tag('kronecker_sum_tag') module-attribute

Operator is a Kronecker sum A (+) B = A (x) I_b + I_a (x) B.

block_diagonal_tag = _Tag('block_diagonal_tag') module-attribute

Operator is block diagonal.

block_tridiagonal_tag = _Tag('block_tridiagonal_tag') module-attribute

Operator is block tridiagonal.

low_rank_tag = _Tag('low_rank_tag') module-attribute

Operator has low-rank structure (e.g. L + U D V^T).

diagonal_tag = lx.diagonal_tag module-attribute

symmetric_tag = lx.symmetric_tag module-attribute

positive_semidefinite_tag = lx.positive_semidefinite_tag module-attribute

negative_semidefinite_tag = lx.negative_semidefinite_tag module-attribute

lower_triangular_tag = lx.lower_triangular_tag module-attribute

upper_triangular_tag = lx.upper_triangular_tag module-attribute

tridiagonal_tag = lx.tridiagonal_tag module-attribute

unit_diagonal_tag = lx.unit_diagonal_tag module-attribute

is_kronecker(operator: lx.AbstractLinearOperator) -> bool

Check whether operator carries the Kronecker tag.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator to inspect.	required

Returns:

Type	Description
bool	True if the operator carries the Kronecker tag, else False.

Source code in src/gaussx/_tags.py

@ft.singledispatch
def is_kronecker(operator: lx.AbstractLinearOperator) -> bool:
    """Check whether *operator* carries the Kronecker tag.

    Args:
        operator: Linear operator to inspect.

    Returns:
        ``True`` if the operator carries the Kronecker tag, else ``False``.
    """
    return False

is_kronecker_sum(operator: lx.AbstractLinearOperator) -> bool

Check whether operator carries the Kronecker sum tag.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator to inspect.	required

Returns:

Type	Description
bool	True if the operator carries the Kronecker sum tag, else False.

Source code in src/gaussx/_tags.py

@ft.singledispatch
def is_kronecker_sum(operator: lx.AbstractLinearOperator) -> bool:
    """Check whether *operator* carries the Kronecker sum tag.

    Args:
        operator: Linear operator to inspect.

    Returns:
        ``True`` if the operator carries the Kronecker sum tag, else ``False``.
    """
    return False

is_block_diagonal(operator: lx.AbstractLinearOperator) -> bool

Check whether operator carries the block-diagonal tag.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator to inspect.	required

Returns:

Type	Description
bool	True if the operator carries the block-diagonal tag, else False.

Source code in src/gaussx/_tags.py

@ft.singledispatch
def is_block_diagonal(operator: lx.AbstractLinearOperator) -> bool:
    """Check whether *operator* carries the block-diagonal tag.

    Args:
        operator: Linear operator to inspect.

    Returns:
        ``True`` if the operator carries the block-diagonal tag, else ``False``.
    """
    return False

is_block_tridiagonal(operator: lx.AbstractLinearOperator) -> bool

Check whether operator carries the block-tridiagonal tag.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator to inspect.	required

Returns:

Type	Description
bool	True if the operator carries the block-tridiagonal tag, else
bool	False.

Source code in src/gaussx/_tags.py

@ft.singledispatch
def is_block_tridiagonal(operator: lx.AbstractLinearOperator) -> bool:
    """Check whether *operator* carries the block-tridiagonal tag.

    Args:
        operator: Linear operator to inspect.

    Returns:
        ``True`` if the operator carries the block-tridiagonal tag, else
        ``False``.
    """
    return False

is_low_rank(operator: lx.AbstractLinearOperator) -> bool

Check whether operator carries the low-rank tag.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Linear operator to inspect.	required

Returns:

Type	Description
bool	True if the operator carries the low-rank tag, else False.

Source code in src/gaussx/_tags.py

@ft.singledispatch
def is_low_rank(operator: lx.AbstractLinearOperator) -> bool:
    """Check whether *operator* carries the low-rank tag.

    Args:
        operator: Linear operator to inspect.

    Returns:
        ``True`` if the operator carries the low-rank tag, else ``False``.
    """
    return False

is_eigen_reducible(operator: lx.AbstractLinearOperator) -> bool

Whether a sum of Kronecker products has the exact two-term reduction.

"Eigen-reducible" means operator is A₁ ⊗ B₁ + A₂ ⊗ B₂ — a two-term SumOfKroneckers, or the same two terms built with sum_operator — whose main factors are symmetric and whose other ("anchor") term is positive definite. Then whitening by the anchor leaves one eigendecomposition per factor, so gaussx.solve, gaussx.logdet and SumOfKroneckers.eigendecompose cost O(n_A³ + n_B³) instead of the O((n_A n_B)³) dense fallback. Three or more terms are never eigen-reducible.

A purely structural query — operator types and lineax tags only, no array work — so it is safe to call from solver-strategy selection and under jax.jit. The exact path can still decline an operator this accepts, when a diagonal anchor factor is concretely non-positive.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	Any lineax operator.	required

Returns:

Type	Description
bool	True if the exact two-term reduction applies.

Source code in src/gaussx/_operators/_sum_kronecker.py

def is_eigen_reducible(operator: lx.AbstractLinearOperator) -> bool:
    r"""Whether a sum of Kronecker products has the exact two-term reduction.

    "Eigen-reducible" means ``operator`` is ``A₁ ⊗ B₁ + A₂ ⊗ B₂`` — a
    two-term `SumOfKroneckers`, or the same two terms built with
    `sum_operator` — whose main factors are symmetric and whose other
    ("anchor") term is positive definite. Then whitening by the anchor
    leaves one eigendecomposition per factor, so `gaussx.solve`,
    `gaussx.logdet` and `SumOfKroneckers.eigendecompose` cost
    ``O(n_A³ + n_B³)`` instead of the ``O((n_A n_B)³)`` dense fallback.
    Three or more terms are never eigen-reducible.

    A purely structural query — operator types and lineax tags only, no
    array work — so it is safe to call from solver-strategy selection and
    under ``jax.jit``. The exact path can still decline an operator this
    accepts, when a diagonal anchor factor is concretely non-positive.

    Args:
        operator: Any lineax operator.

    Returns:
        ``True`` if the exact two-term reduction applies.
    """
    terms = _kronecker_terms(operator)
    return terms is not None and _select_anchor(terms) is not None

Deprecated aliases

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

Structured linear algebra and Gaussian primitives for JAX.

SumKronecker

Bases: SumOfKroneckers

Deprecated alias for SumOfKroneckers.

The old name was one word-order away from KroneckerSum, a different operator with a different eigendecomposition — picking the wrong one produced silently wrong solves. Renamed in gh-136.

Subclasses SumOfKroneckers so isinstance checks and singledispatch registrations keyed on this class keep working, and emits a DeprecationWarning on construction. Structural operations that rebuild the operator — transpose / .T — preserve this type, so they warn again. Note that instances built via SumOfKroneckers are not instances of this subclass; internal dispatch keys on the parent. Will be removed in gaussx 0.7.0.

Source code in src/gaussx/_operators/_sum_kronecker.py

class SumKronecker(SumOfKroneckers):
    """Deprecated alias for `SumOfKroneckers`.

    The old name was one word-order away from `KroneckerSum`, a
    different operator with a different eigendecomposition — picking the wrong
    one produced silently wrong solves. Renamed in gh-136.

    Subclasses `SumOfKroneckers` so ``isinstance`` checks and
    ``singledispatch`` registrations keyed on this class keep working, and
    emits a `DeprecationWarning` on construction. Structural operations that
    rebuild the operator — ``transpose`` / ``.T`` — preserve this type, so
    they warn again. Note that instances built via `SumOfKroneckers` are
    *not* instances of this subclass; internal dispatch keys on the parent.
    Will be removed in gaussx 0.7.0.
    """

    def __init__(
        self,
        kron1: Kronecker,
        kron2: Kronecker,
        *krons: Kronecker,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:

        warn_deprecated(
            "SumKronecker is deprecated and will be removed in gaussx 0.7.0; use "
            "SumOfKroneckers "
            "(KroneckerSum remains a different operator)."
        )
        super().__init__(kron1, kron2, *krons, tags=tags)

SumKroneckerSqrt

Bases: SumOfKroneckersSqrt

Deprecated alias for SumOfKroneckersSqrt (gh-297).

Subclasses SumOfKroneckersSqrt so isinstance checks keep working, and emits a DeprecationWarning on construction. gaussx.sqrt returns the new class, which is not an instance of this alias. Will be removed in gaussx 0.7.0.

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

class SumKroneckerSqrt(SumOfKroneckersSqrt):
    """Deprecated alias for `SumOfKroneckersSqrt` (gh-297).

    Subclasses `SumOfKroneckersSqrt` so ``isinstance`` checks keep working,
    and emits a `DeprecationWarning` on construction. `gaussx.sqrt` returns
    the new class, which is *not* an instance of this alias. Will be removed
    in gaussx 0.7.0.
    """

    def __init__(
        self,
        original: SumOfKroneckers,
        lanczos_order: int = _DEFAULT_LANCZOS_ORDER,
    ) -> None:
        warn_deprecated(
            "SumKroneckerSqrt is deprecated and will be removed in gaussx 0.7.0; "
            "use SumOfKroneckersSqrt."
        )
        super().__init__(original, lanczos_order=lanczos_order)

SVDLowRankUpdate

Bases: LowRankUpdate

Deprecated subclass of LowRankUpdate with orthonormal=True.

Preserves the pre-consolidation public API until gaussx 0.7.0:

• Same constructor signature as the old class — S defaults to ones (via the parent LowRankUpdate) if omitted, and V defaults to U so calls like SVDLowRankUpdate(base, U, S) and SVDLowRankUpdate(base, U) continue to work.
• Inherits from LowRankUpdate so isinstance / issubclass checks and singledispatch registrations keyed on this class keep working.
• Forces orthonormal=True and emits a DeprecationWarning on construction.

New code should construct LowRankUpdate(base, U, S, V, orthonormal=True) (or use svd_low_rank_plus_diag) directly. Will be removed in gaussx 0.7.0.

Source code in src/gaussx/_operators/_low_rank_update.py

class SVDLowRankUpdate(LowRankUpdate):
    """Deprecated subclass of `LowRankUpdate` with ``orthonormal=True``.

    Preserves the pre-consolidation public API until gaussx 0.7.0:

    - Same constructor signature as the old class — ``S`` defaults to
      ones (via the parent ``LowRankUpdate``) if omitted, and ``V``
      defaults to ``U`` so calls like ``SVDLowRankUpdate(base, U, S)``
      and ``SVDLowRankUpdate(base, U)`` continue to work.
    - Inherits from `LowRankUpdate` so ``isinstance`` /
      ``issubclass`` checks and ``singledispatch`` registrations keyed
      on this class keep working.
    - Forces ``orthonormal=True`` and emits a
      `DeprecationWarning` on construction.

    New code should construct ``LowRankUpdate(base, U, S, V,
    orthonormal=True)`` (or use `svd_low_rank_plus_diag`)
    directly. Will be removed in gaussx 0.7.0.
    """

    def __init__(
        self,
        base: lx.AbstractLinearOperator,
        U: Float[Array, "n k"],
        S: Float[Array, " k"] | None = None,
        V: Float[Array, "n k"] | None = None,
        *,
        tags: object | frozenset[object] = frozenset(),
    ) -> None:

        warn_deprecated(
            "SVDLowRankUpdate is deprecated and will be removed in gaussx 0.7.0; "
            "use LowRankUpdate(..., orthonormal=True) or svd_low_rank_plus_diag()."
        )
        super().__init__(base, U, S, V, tags=tags, orthonormal=True)

sumkronecker_sample(op: SumOfKroneckers, *, key: jax.Array, num_samples: int = 1, lanczos_order: int = 50) -> Float[Array, 'num_samples n']

Deprecated alias for sum_of_kroneckers_sample (gh-297).

Will be removed in gaussx 0.7.0.

Source code in src/gaussx/_operators/_sum_kronecker.py

def sumkronecker_sample(
    op: SumOfKroneckers,
    *,
    key: jax.Array,
    num_samples: int = 1,
    lanczos_order: int = 50,
) -> Float[Array, "num_samples n"]:
    """Deprecated alias for `sum_of_kroneckers_sample` (gh-297).

    Will be removed in gaussx 0.7.0.
    """
    warn_deprecated(
        "sumkronecker_sample is deprecated and will be removed in gaussx 0.7.0; "
        "use sum_of_kroneckers_sample."
    )
    return sum_of_kroneckers_sample(
        op, key=key, num_samples=num_samples, lanczos_order=lanczos_order
    )