Solvers &amp; Preconditioners &para;

Dense solver strategy using gaussx structural dispatch.

Delegates to gaussx.solve and gaussx.logdet which automatically select the best algorithm based on operator structure (Diagonal, BlockDiag, Kronecker, LowRankUpdate, or dense fallback via lineax).

Source code in src/gaussx/_strategies/_dense.py

class DenseSolver(AbstractSolverStrategy):
    """Dense solver strategy using gaussx structural dispatch.

    Delegates to ``gaussx.solve`` and ``gaussx.logdet`` which
    automatically select the best algorithm based on operator
    structure (Diagonal, BlockDiag, Kronecker, LowRankUpdate,
    or dense fallback via lineax).
    """

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``A x = b`` via structural dispatch (``gaussx.solve``).

        Args:
            operator: Linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        return _solve(operator, vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Compute ``log|det(A)|`` via structural dispatch (``gaussx.logdet``).

        Args:
            operator: Linear operator ``A``.
            key: Unused; present for protocol compatibility with stochastic
                estimators.

        Returns:
            Scalar log-determinant.
        """
        return _logdet(operator)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b via structural dispatch (gaussx.solve).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_dense.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``A x = b`` via structural dispatch (``gaussx.solve``).

    Args:
        operator: Linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    return _solve(operator, vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Compute log|det(A)| via structural dispatch (gaussx.logdet).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`key`	`Array \| None`	Unused; present for protocol compatibility with stochastic estimators.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_dense.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Compute ``log|det(A)|`` via structural dispatch (``gaussx.logdet``).

    Args:
        operator: Linear operator ``A``.
        key: Unused; present for protocol compatibility with stochastic
            estimators.

    Returns:
        Scalar log-determinant.
    """
    return _logdet(operator)

`AutoSolver` ¶

Automatic solver selection based on operator type and size.

Selection logic:

Structured, i.e. any operator with an exact structural solve and log-determinant (Diagonal, DiagonalizedOperator, BlockDiag, Kronecker, LowRankUpdate, KroneckerSum, SpectralFunction, BlockTriDiag and its bidiagonal factors, eigen-reducible sums of Kronecker products; also inside TaggedLinearOperator, c * A, A / c and -A): DenseSolver, whose structural dispatch is exact and cheap
Small dense (N <= size_threshold): DenseSolver
Large PSD: CGSolver. Its logdet is a stochastic, fixed-seed SLQ estimate (the same probes on every call unless a key is passed, e.g. through gaussx.KeyedSolver); for an exact one use ComposedSolver(CGSolver(), DenseLogdet())
Large general: DenseSolver (fallback)

Attributes:

Name	Type	Description
`size_threshold`	`int`	Matrix dimension above which iterative solvers are preferred. Default: 1000.
`throw`	`bool`	Forwarded to the `CGSolver` built for large PSD operators: raise when CG does not converge (the default), or return the last iterate unchecked.
`preconditioner`	`AbstractPreconditioner \| None`	Forwarded to the `CGSolver` built for large PSD operators; ignored when a direct solve is chosen (gh-390).

Source code in src/gaussx/_strategies/_auto.py

class AutoSolver(AbstractSolverStrategy):
    """Automatic solver selection based on operator type and size.

    Selection logic:

    - Structured, i.e. any operator with an exact structural solve and
      log-determinant (Diagonal,
      DiagonalizedOperator, BlockDiag, Kronecker, LowRankUpdate,
      KroneckerSum, SpectralFunction, BlockTriDiag and its bidiagonal
      factors, eigen-reducible sums of Kronecker products; also inside
      ``TaggedLinearOperator``, ``c * A``, ``A / c`` and ``-A``):
      DenseSolver, whose structural dispatch is exact and cheap
    - Small dense (N <= size_threshold): DenseSolver
    - Large PSD: CGSolver. Its ``logdet`` is a **stochastic**, fixed-seed
      SLQ estimate (the same probes on every call unless a ``key`` is
      passed, e.g. through `gaussx.KeyedSolver`); for an exact one use
      ``ComposedSolver(CGSolver(), DenseLogdet())``
    - Large general: DenseSolver (fallback)

    Attributes:
        size_threshold: Matrix dimension above which iterative
            solvers are preferred. Default: 1000.
        throw: Forwarded to the `CGSolver` built for large PSD operators:
            raise when CG does not converge (the default), or return the
            last iterate unchecked.
        preconditioner: Forwarded to the `CGSolver` built for large PSD
            operators; ignored when a direct solve is chosen (gh-390).
    """

    size_threshold: int = eqx.field(static=True, default=1000)
    throw: bool = eqx.field(static=True, default=True)
    preconditioner: AbstractPreconditioner | None = None

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve A x = b with automatically selected algorithm.

        Args:
            operator: The linear operator A.
            vector: The right-hand side b.

        Returns:
            The solution x.
        """
        return self._get_strategy(operator).solve(operator, vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Compute log |det(A)| with automatically selected algorithm.

        Args:
            operator: The linear operator A.
            key: Optional PRNG key (forwarded to stochastic strategies).

        Returns:
            Scalar log |det(A)|.
        """
        return self._get_strategy(operator).logdet(operator, key=key)

    def _get_strategy(
        self, operator: lx.AbstractLinearOperator
    ) -> AbstractSolverStrategy:
        """Select the best solver strategy for the given operator."""

        # Exact structural solve + logdet (gh-321). A sum of Kronecker
        # products only counts when the exact two-term reduction applies;
        # without it ``gaussx.solve`` would materialize, so it falls through
        # to the size/tag rules below and can still pick up CG.
        if _has_structural_logdet(operator):
            return DenseSolver()

        n = operator.in_size()
        if n <= self.size_threshold:
            return DenseSolver()

        # Large operators: use CG for PSD, DenseSolver otherwise
        if lx.is_positive_semidefinite(operator):
            return CGSolver(throw=self.throw, preconditioner=self.preconditioner)

        return DenseSolver()

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b with automatically selected algorithm.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	The linear operator A.	required
`vector`	`Float[Array, ' n']`	The right-hand side b.	required

Returns:

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

Source code in src/gaussx/_strategies/_auto.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve A x = b with automatically selected algorithm.

    Args:
        operator: The linear operator A.
        vector: The right-hand side b.

    Returns:
        The solution x.
    """
    return self._get_strategy(operator).solve(operator, vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Compute log |det(A)| with automatically selected algorithm.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	The linear operator A.	required
`key`	`Array \| None`	Optional PRNG key (forwarded to stochastic strategies).	`None`

Returns:

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

Source code in src/gaussx/_strategies/_auto.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Compute log |det(A)| with automatically selected algorithm.

    Args:
        operator: The linear operator A.
        key: Optional PRNG key (forwarded to stochastic strategies).

    Returns:
        Scalar log |det(A)|.
    """
    return self._get_strategy(operator).logdet(operator, key=key)

`CGSolver` ¶

Iterative CG solver with stochastic log-determinant.

Uses lineax CG for the linear solve and matfree's stochastic Lanczos quadrature (SLQ) for the log-determinant. Suitable for large PSD operators where dense factorization is too expensive.

The relative tolerance defaults to None, which resolves from the operator's dtype at solve time: 1e-5 in float64 and 1e-3 in float32, where a relative residual of 1e-5 is out of reach once the condition number passes about 1e3 (gh-327). Set them explicitly to override.

Attributes:

Name	Type	Description
`rtol`	`float \| None`	Relative tolerance for CG. `None`: `1e-5` in float64, `1e-3` in float32.
`atol`	`float \| None`	Absolute tolerance for CG. `None`: `1e-5` in float64; in float32, `sqrt(eps) * max\|b_i\|`, relative to the right-hand side (the solve runs on `b / max\|b_i\|`; gh-639).
`max_steps`	`int`	Maximum CG iterations.
`num_probes`	`int`	Number of probe vectors for stochastic logdet.
`lanczos_order`	`int`	Order of the Lanczos decomposition for SLQ.
`seed`	`int`	Seed of the SLQ probes when `logdet` gets no `key`. With no key every call reuses the same probes (common random numbers); see `gaussx.KeyedSolver`.
`preconditioner`	`AbstractPreconditioner \| None`	Optional preconditioner. When set, its approximate inverse is passed to lineax CG to accelerate convergence.
`throw`	`bool`	Raise when CG does not converge within `max_steps` (the default). With `False` the last iterate is returned unchecked; an unconverged CG iterate can be far worse than zero, so only use it where the caller checks the result.

Source code in src/gaussx/_strategies/_cg.py

class CGSolver(AbstractSolverStrategy):
    """Iterative CG solver with stochastic log-determinant.

    Uses lineax CG for the linear solve and matfree's stochastic
    Lanczos quadrature (SLQ) for the log-determinant. Suitable
    for large PSD operators where dense factorization is too
    expensive.

    The relative tolerance defaults to ``None``, which resolves from the
    operator's dtype at solve time: ``1e-5`` in float64 and ``1e-3`` in
    float32, where
    a relative residual of ``1e-5`` is out of reach once the condition number
    passes about ``1e3`` (gh-327). Set them explicitly to override.

    Attributes:
        rtol: Relative tolerance for CG. ``None``: ``1e-5`` in float64,
            ``1e-3`` in float32.
        atol: Absolute tolerance for CG. ``None``: ``1e-5`` in float64; in
            float32, ``sqrt(eps) * max|b_i|``, relative to the right-hand
            side (the solve runs on ``b / max|b_i|``; gh-639).
        max_steps: Maximum CG iterations.
        num_probes: Number of probe vectors for stochastic logdet.
        lanczos_order: Order of the Lanczos decomposition for SLQ.
        seed: Seed of the SLQ probes when `logdet` gets no ``key``. With no
            key every call reuses the same probes (common random numbers);
            see `gaussx.KeyedSolver`.
        preconditioner: Optional preconditioner. When set, its approximate
            inverse is passed to lineax CG to accelerate convergence.
        throw: Raise when CG does not converge within ``max_steps`` (the
            default). With ``False`` the last iterate is returned unchecked;
            an unconverged CG iterate can be far worse than zero, so only
            use it where the caller checks the result.
    """

    rtol: float | None = eqx.field(static=True, default=None)
    atol: float | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=1000)
    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    preconditioner: AbstractPreconditioner | None = None
    throw: bool = eqx.field(static=True, default=True)

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``A x = b`` with conjugate gradients.

        Args:
            operator: A PSD linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        dtype = operator_dtype(operator, vector)
        atol, scale = rhs_scaling(self.atol, dtype, 1e-5, vector)
        rhs = vector if scale is None else vector / scale
        solver = lx.CG(
            rtol=resolve_tolerance(self.rtol, dtype, 1e-5),
            atol=atol,
            max_steps=self.max_steps,
        )
        options: dict[str, lx.AbstractLinearOperator] = {}
        if self.preconditioner is not None:
            precond_op = self.preconditioner.as_operator(operator)
            if precond_op is not None:
                # lineax treats `options` as non-differentiable and raises if a
                # tangent reaches it, but a preconditioner built from the
                # (traced) operator carries one. The CG solution does not
                # depend on M^{-1}, only the iteration count does, so stopping
                # its gradient is exact (as `_inv_quad_logdet` does).
                dynamic, static = eqx.partition(precond_op, eqx.is_array)
                options["preconditioner"] = eqx.combine(
                    jax.lax.stop_gradient(dynamic), static
                )
        x = lx.linear_solve(
            operator, rhs, solver, options=options, throw=self.throw
        ).value
        return x if scale is None else x * scale

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic log-determinant via Lanczos quadrature.

        Args:
            operator: A PSD linear operator.
            key: PRNG key for probe vector sampling. If None,
                uses ``jax.random.PRNGKey(seed)``: the same probes on
                every call.

        Returns:
            Scalar estimate of log |det(A)|.
        """
        return SLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            seed=self.seed,
        ).logdet(operator, key=key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b with conjugate gradients.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_cg.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``A x = b`` with conjugate gradients.

    Args:
        operator: A PSD linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    dtype = operator_dtype(operator, vector)
    atol, scale = rhs_scaling(self.atol, dtype, 1e-5, vector)
    rhs = vector if scale is None else vector / scale
    solver = lx.CG(
        rtol=resolve_tolerance(self.rtol, dtype, 1e-5),
        atol=atol,
        max_steps=self.max_steps,
    )
    options: dict[str, lx.AbstractLinearOperator] = {}
    if self.preconditioner is not None:
        precond_op = self.preconditioner.as_operator(operator)
        if precond_op is not None:
            # lineax treats `options` as non-differentiable and raises if a
            # tangent reaches it, but a preconditioner built from the
            # (traced) operator carries one. The CG solution does not
            # depend on M^{-1}, only the iteration count does, so stopping
            # its gradient is exact (as `_inv_quad_logdet` does).
            dynamic, static = eqx.partition(precond_op, eqx.is_array)
            options["preconditioner"] = eqx.combine(
                jax.lax.stop_gradient(dynamic), static
            )
    x = lx.linear_solve(
        operator, rhs, solver, options=options, throw=self.throw
    ).value
    return x if scale is None else x * scale

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log-determinant via Lanczos quadrature.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If None, uses `jax.random.PRNGKey(seed)`: the same probes on every call.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_cg.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic log-determinant via Lanczos quadrature.

    Args:
        operator: A PSD linear operator.
        key: PRNG key for probe vector sampling. If None,
            uses ``jax.random.PRNGKey(seed)``: the same probes on
            every call.

    Returns:
        Scalar estimate of log |det(A)|.
    """
    return SLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        seed=self.seed,
    ).logdet(operator, key=key)

`PreconditionedCGSolver` ¶

CG solver with pivoted partial Cholesky preconditioner.

Uses gaussx.PartialCholeskyPreconditioner's guarded pivoted partial Cholesky to build a rank-k preconditioner, then solves (σ² I + F Fᵀ)⁻¹ v via the Woodbury identity inside lineax CG.

For operators of the form K + σ² I, preconditioning dramatically reduces the number of CG iterations.

Pass a preconditioner built once with gaussx.PartialCholeskyPreconditioner.from_operator or gaussx.NystromPreconditioner.from_operator (on K, with shift=σ²) to reuse it across solves. Otherwise a rank preconditioner_rank factor of A − shift · I is rebuilt from the system operator A at every solve, so the noise is never counted twice (#345).

Attributes:

Name	Type	Description
`preconditioner_rank`	`int`	Rank of the partial Cholesky built per solve. Set to 0 to disable preconditioning (falls back to plain CG). Ignored when `preconditioner` is given.
`shift`	`float`	The noise variance `σ²` in the system `A = K + σ² I`, for the preconditioner built per solve. Must not exceed the noise actually in `A`. Ignored when `preconditioner` is given.
`rtol`	`float \| None`	Relative tolerance for CG. `None`: `1e-5` in float64, `1e-3` in float32 (see `gaussx.CGSolver`).
`atol`	`float \| None`	Absolute tolerance for CG. `None`: `1e-5` in float64; in float32, `sqrt(eps) * max\|b_i\|`, relative to the right-hand side (the solve runs on `b / max\|b_i\|`; gh-639).
`max_steps`	`int`	Maximum CG iterations.
`num_probes`	`int`	Number of probe vectors for stochastic logdet.
`lanczos_order`	`int`	Lanczos iterations for SLQ logdet.
`seed`	`int`	Seed for probe vector generation.
`preconditioner`	`AbstractPreconditioner \| None`	A prebuilt preconditioner, used for every solve.
`throw`	`bool`	Raise when CG does not converge within `max_steps`. With `False` the last iterate is returned unchecked (see `gaussx.CGSolver`).

Source code in src/gaussx/_strategies/_precond_cg.py

class PreconditionedCGSolver(AbstractSolverStrategy):
    """CG solver with pivoted partial Cholesky preconditioner.

    Uses `gaussx.PartialCholeskyPreconditioner`'s guarded pivoted
    partial Cholesky to build a rank-k preconditioner, then solves
    ``(σ² I + F Fᵀ)⁻¹ v`` via the Woodbury identity inside lineax CG.

    For operators of the form ``K + σ² I``, preconditioning
    dramatically reduces the number of CG iterations.

    Pass a ``preconditioner`` built once with
    `gaussx.PartialCholeskyPreconditioner.from_operator` or
    `gaussx.NystromPreconditioner.from_operator` (on ``K``, with
    ``shift=σ²``) to reuse it across solves. Otherwise a rank
    ``preconditioner_rank`` factor of ``A − shift · I`` is rebuilt from the
    system operator ``A`` at every solve, so the noise is never counted
    twice (#345).

    Attributes:
        preconditioner_rank: Rank of the partial Cholesky built per solve.
            Set to 0 to disable preconditioning (falls back to plain CG).
            Ignored when ``preconditioner`` is given.
        shift: The noise variance ``σ²`` in the system ``A = K + σ² I``,
            for the preconditioner built per solve. Must not exceed the
            noise actually in ``A``. Ignored when ``preconditioner`` is given.
        rtol: Relative tolerance for CG. ``None``: ``1e-5`` in float64,
            ``1e-3`` in float32 (see `gaussx.CGSolver`).
        atol: Absolute tolerance for CG. ``None``: ``1e-5`` in float64; in
            float32, ``sqrt(eps) * max|b_i|``, relative to the right-hand
            side (the solve runs on ``b / max|b_i|``; gh-639).
        max_steps: Maximum CG iterations.
        num_probes: Number of probe vectors for stochastic logdet.
        lanczos_order: Lanczos iterations for SLQ logdet.
        seed: Seed for probe vector generation.
        preconditioner: A prebuilt preconditioner, used for every solve.
        throw: Raise when CG does not converge within ``max_steps``. With
            ``False`` the last iterate is returned unchecked (see
            `gaussx.CGSolver`).
    """

    preconditioner_rank: int = eqx.field(static=True, default=50)
    shift: float = eqx.field(static=True, default=1.0)
    rtol: float | None = eqx.field(static=True, default=None)
    atol: float | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=1000)
    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    preconditioner: AbstractPreconditioner | None = None
    throw: bool = eqx.field(static=True, default=True)

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve A x = b via preconditioned CG.

        Args:
            operator: A PSD linear operator.
            vector: The right-hand side b.

        Returns:
            The solution x.
        """
        preconditioner = self.preconditioner
        if preconditioner is None:
            preconditioner = PartialCholeskyPreconditioner(
                rank=self.preconditioner_rank, shift=self.shift
            )
        return CGSolver(
            rtol=self.rtol,
            atol=self.atol,
            max_steps=self.max_steps,
            preconditioner=preconditioner,
            throw=self.throw,
        ).solve(operator, vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic log-determinant via Lanczos quadrature.

        Args:
            operator: A PSD linear operator.

        Returns:
            Scalar estimate of log |det(A)|.
        """
        return SLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            seed=self.seed,
        ).logdet(operator, key=key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b via preconditioned CG.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`vector`	`Float[Array, ' n']`	The right-hand side b.	required

Returns:

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

Source code in src/gaussx/_strategies/_precond_cg.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve A x = b via preconditioned CG.

    Args:
        operator: A PSD linear operator.
        vector: The right-hand side b.

    Returns:
        The solution x.
    """
    preconditioner = self.preconditioner
    if preconditioner is None:
        preconditioner = PartialCholeskyPreconditioner(
            rank=self.preconditioner_rank, shift=self.shift
        )
    return CGSolver(
        rtol=self.rtol,
        atol=self.atol,
        max_steps=self.max_steps,
        preconditioner=preconditioner,
        throw=self.throw,
    ).solve(operator, vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log-determinant via Lanczos quadrature.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required

Returns:

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

Source code in src/gaussx/_strategies/_precond_cg.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic log-determinant via Lanczos quadrature.

    Args:
        operator: A PSD linear operator.

    Returns:
        Scalar estimate of log |det(A)|.
    """
    return SLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        seed=self.seed,
    ).logdet(operator, key=key)

`MINRESSolver` ¶

MINRES solver for symmetric (possibly indefinite) systems.

Uses the Lanczos-based MINRES algorithm for the linear solve and matfree's stochastic Lanczos quadrature (SLQ) for the log-determinant. Unlike CG, MINRES only requires symmetry — it works on indefinite and singular systems.

Use cases: EP natural parameters, saddle-point systems, Laplace approximation Hessians.

The solve stops as soon as it converges and goes through lineax.linear_solve, so gradients are implicit (one more MINRES solve) rather than a tape through the iterations. Running out of max_steps raises, like gaussx.CGSolver, unless throw=False (gh-336).

Attributes:

Name	Type	Description
`rtol`	`float \| None`	Relative tolerance for MINRES. `None`: `1e-5` in float64, `1e-3` in float32 (gh-327).
`atol`	`float \| None`	Absolute tolerance for MINRES. `None`: `1e-5` in every dtype.
`max_steps`	`int`	Maximum MINRES iterations.
`shift`	`float`	Diagonal shift — solves `(A + shift * I) x = b`.
`num_probes`	`int`	Number of probe vectors for stochastic logdet.
`lanczos_order`	`int`	Order of the Lanczos decomposition for SLQ.
`seed`	`int`	Seed of the SLQ probes when `logdet` gets no `key` (the same probes on every call; see `gaussx.KeyedSolver`).
`throw`	`bool`	Raise when MINRES does not converge within `max_steps` (the default). With `False` the last iterate is returned unchecked.

Source code in src/gaussx/_strategies/_minres.py

class MINRESSolver(AbstractSolverStrategy):
    """MINRES solver for symmetric (possibly indefinite) systems.

    Uses the Lanczos-based MINRES algorithm for the linear solve
    and matfree's stochastic Lanczos quadrature (SLQ) for the
    log-determinant. Unlike CG, MINRES only requires symmetry —
    it works on indefinite and singular systems.

    Use cases: EP natural parameters, saddle-point systems,
    Laplace approximation Hessians.

    The solve stops as soon as it converges and goes through
    `lineax.linear_solve`, so gradients are implicit (one more MINRES solve)
    rather than a tape through the iterations. Running out of ``max_steps``
    raises, like `gaussx.CGSolver`, unless ``throw=False`` (gh-336).

    Attributes:
        rtol: Relative tolerance for MINRES. ``None``: ``1e-5`` in float64,
            ``1e-3`` in float32 (gh-327).
        atol: Absolute tolerance for MINRES. ``None``: ``1e-5`` in every
            dtype.
        max_steps: Maximum MINRES iterations.
        shift: Diagonal shift — solves ``(A + shift * I) x = b``.
        num_probes: Number of probe vectors for stochastic logdet.
        lanczos_order: Order of the Lanczos decomposition for SLQ.
        seed: Seed of the SLQ probes when `logdet` gets no ``key`` (the same
            probes on every call; see `gaussx.KeyedSolver`).
        throw: Raise when MINRES does not converge within ``max_steps`` (the
            default). With ``False`` the last iterate is returned unchecked.
    """

    rtol: float | None = eqx.field(static=True, default=None)
    atol: float | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=1000)
    shift: float = eqx.field(static=True, default=0.0)
    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    throw: bool = eqx.field(static=True, default=True)

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``(A + shift I) x = b`` with MINRES.

        Args:
            operator: A symmetric (possibly indefinite) linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        dtype = operator_dtype(operator, vector)
        solver = _MINRES(
            rtol=resolve_tolerance(self.rtol, dtype, 1e-5),
            atol=1e-5 if self.atol is None else self.atol,
            max_steps=self.max_steps,
        )
        if self.shift != 0.0:
            identity = lx.IdentityLinearOperator(operator.in_structure())
            shifted = operator + self.shift * identity
            if lx.is_symmetric(operator):
                shifted = lx.TaggedLinearOperator(shifted, lx.symmetric_tag)
            operator = shifted
        return lx.linear_solve(operator, vector, solver, throw=self.throw).value

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic ``log|det(A + shift I)|`` via Lanczos quadrature.

        Args:
            operator: A symmetric linear operator.
            key: PRNG key for probe vector sampling. If None,
                uses ``jax.random.PRNGKey(seed)``.

        Returns:
            Scalar estimate of ``log|det(A + shift I)|``.
        """
        return IndefiniteSLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            shift=self.shift,
            seed=self.seed,
        ).logdet(operator, key=key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve (A + shift I) x = b with MINRES.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric (possibly indefinite) linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_minres.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``(A + shift I) x = b`` with MINRES.

    Args:
        operator: A symmetric (possibly indefinite) linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    dtype = operator_dtype(operator, vector)
    solver = _MINRES(
        rtol=resolve_tolerance(self.rtol, dtype, 1e-5),
        atol=1e-5 if self.atol is None else self.atol,
        max_steps=self.max_steps,
    )
    if self.shift != 0.0:
        identity = lx.IdentityLinearOperator(operator.in_structure())
        shifted = operator + self.shift * identity
        if lx.is_symmetric(operator):
            shifted = lx.TaggedLinearOperator(shifted, lx.symmetric_tag)
        operator = shifted
    return lx.linear_solve(operator, vector, solver, throw=self.throw).value

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log|det(A + shift I)| via Lanczos quadrature.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If None, uses `jax.random.PRNGKey(seed)`.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar estimate of `log\|det(A + shift I)\|`.

Source code in src/gaussx/_strategies/_minres.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic ``log|det(A + shift I)|`` via Lanczos quadrature.

    Args:
        operator: A symmetric linear operator.
        key: PRNG key for probe vector sampling. If None,
            uses ``jax.random.PRNGKey(seed)``.

    Returns:
        Scalar estimate of ``log|det(A + shift I)|``.
    """
    return IndefiniteSLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        shift=self.shift,
        seed=self.seed,
    ).logdet(operator, key=key)

`LSMRSolver` ¶

LSMR iterative least-squares solver (Fong & Saunders 2011).

Matrix-free solver that only requires matvec and transpose-matvec. Supports Tikhonov regularization via damp parameter: minimizes ||Ax - b||^2 + damp^2 ||x||^2.

Suitable for rectangular, ill-conditioned, or regularized systems.

The undamped path delegates to lineax.LSMR (with lineax's implicit-differentiation rules). lineax's LSMR has no Tikhonov damp parameter, so damped solves use matfree's LSMR, which has a custom VJP for memory-efficient backpropagation.

Attributes:

Name	Type	Description
`atol`	`float \| None`	Absolute tolerance. `None`: `1e-6` in float64, `1e-3` in float32 (gh-327).
`btol`	`float \| None`	Relative tolerance on the residual. `None`: as `atol`.
`ctol`	`float`	Condition number tolerance (the lineax path uses `conlim = 1 / ctol`).
`max_steps`	`int`	Maximum iterations (formerly `maxiter`, a deprecated alias; gh-405).
`damp`	`float`	Tikhonov damping parameter.
`num_probes`	`int`	Number of probe vectors for stochastic logdet.
`lanczos_order`	`int`	Lanczos iterations for SLQ logdet.
`seed`	`int`	Seed for probe vector generation.
`throw`	`bool`	Raise when LSMR stops without converging (`max_steps` reached), on both the lineax and the damped matfree path. With `False` the last iterate is returned unchecked (gh-336).

Source code in src/gaussx/_strategies/_lsmr.py

class LSMRSolver(AbstractSolverStrategy):
    """LSMR iterative least-squares solver (Fong & Saunders 2011).

    Matrix-free solver that only requires matvec and transpose-matvec.
    Supports Tikhonov regularization via ``damp`` parameter:
    minimizes ``||Ax - b||^2 + damp^2 ||x||^2``.

    Suitable for rectangular, ill-conditioned, or regularized systems.

    The undamped path delegates to `lineax.LSMR` (with lineax's
    implicit-differentiation rules). lineax's LSMR has no Tikhonov
    ``damp`` parameter, so damped solves use matfree's LSMR, which has
    a custom VJP for memory-efficient backpropagation.

    Attributes:
        atol: Absolute tolerance. ``None``: ``1e-6`` in float64, ``1e-3``
            in float32 (gh-327).
        btol: Relative tolerance on the residual. ``None``: as ``atol``.
        ctol: Condition number tolerance (the lineax path uses
            ``conlim = 1 / ctol``).
        max_steps: Maximum iterations (formerly ``maxiter``, a deprecated
            alias; gh-405).
        damp: Tikhonov damping parameter.
        num_probes: Number of probe vectors for stochastic logdet.
        lanczos_order: Lanczos iterations for SLQ logdet.
        seed: Seed for probe vector generation.
        throw: Raise when LSMR stops without converging (``max_steps``
            reached), on both the lineax and the damped matfree path. With
            ``False`` the last iterate is returned unchecked (gh-336).
    """

    atol: float | None = eqx.field(static=True, default=None)
    btol: float | None = eqx.field(static=True, default=None)
    ctol: float = eqx.field(static=True, default=1e-6)
    max_steps: int = eqx.field(static=True, default=1000)
    damp: float = eqx.field(static=True, default=0.0)
    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    throw: bool = eqx.field(static=True, default=True)

    def __init__(
        self,
        *,
        atol: float | None = None,
        btol: float | None = None,
        ctol: float = 1e-6,
        max_steps: int = UNSET,
        damp: float = 0.0,
        num_probes: int = 20,
        lanczos_order: int = 30,
        seed: int = 0,
        throw: bool = True,
        maxiter: int = UNSET,
    ) -> None:
        self.atol = atol
        self.btol = btol
        self.ctol = ctol
        self.max_steps = default(
            renamed("LSMRSolver", "maxiter", maxiter, "max_steps", max_steps), 1000
        )
        self.damp = damp
        self.num_probes = num_probes
        self.lanczos_order = lanczos_order
        self.seed = seed
        self.throw = throw

    @property
    def maxiter(self) -> int:
        """Deprecated: `max_steps` (gh-405)."""
        warn_deprecated(
            "LSMRSolver.maxiter is deprecated and will be removed in gaussx 0.7.0; "
            "use .max_steps (gh-405)."
        )
        return self.max_steps

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " m"],
    ) -> Float[Array, " n"]:
        """Solve A x = b via LSMR.

        Args:
            operator: A linear operator (may be rectangular).
            vector: The right-hand side b.

        Returns:
            The (least-squares) solution x.
        """
        dtype = operator_dtype(operator, vector)
        atol = resolve_tolerance(self.atol, dtype, 1e-6)
        btol = resolve_tolerance(self.btol, dtype, 1e-6)
        if self.damp == 0.0:
            solver = lx.LSMR(
                rtol=btol,
                atol=atol,
                max_steps=self.max_steps,
                conlim=1.0 / self.ctol if self.ctol > 0 else 1e8,
            )
            return lx.linear_solve(operator, vector, solver, throw=self.throw).value

        # Tikhonov damping: not supported by lineax's LSMR, so the
        # damped path stays on matfree.
        lsmr_fn = matfree.lstsq.lsmr(
            atol=atol,
            btol=btol,
            ctol=self.ctol,
            maxiter=self.max_steps,
        )

        def vecmat(v):
            return operator.T.mv(v)

        x, stats = lsmr_fn(vecmat, vector, damp=self.damp)
        if self.throw:
            # matfree reports success but never raises; match lineax.
            x = eqx.error_if(
                x,
                jnp.logical_not(stats["success"]),
                "LSMR did not converge within `max_steps` steps. Increase "
                "`max_steps`, loosen `atol`/`btol`, or pass `throw=False`.",
            )
        return x

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic log-determinant via Lanczos quadrature.

        LSMR itself does not use this: it is SLQ on a *square* symmetric
        PSD operator, for the `AbstractSolverStrategy` interface.

        Args:
            operator: A square symmetric PSD linear operator.

        Returns:
            Scalar estimate of log |det(A)|.

        Raises:
            ValueError: If the operator is not square (gh-402).
        """
        if operator.in_size() != operator.out_size():
            raise ValueError(
                "LSMRSolver.logdet needs a square operator; got shape "
                f"({operator.out_size()}, {operator.in_size()})."
            )
        return SLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            seed=self.seed,
        ).logdet(operator, key=key)

`maxiter: int` `property` ¶

Deprecated: max_steps (gh-405).

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' m']) -> Float[Array, ' n']` ¶

Solve A x = b via LSMR.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A linear operator (may be rectangular).	required
`vector`	`Float[Array, ' m']`	The right-hand side b.	required

Returns:

Type	Description
`Float[Array, ' n']`	The (least-squares) solution x.

Source code in src/gaussx/_strategies/_lsmr.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " m"],
) -> Float[Array, " n"]:
    """Solve A x = b via LSMR.

    Args:
        operator: A linear operator (may be rectangular).
        vector: The right-hand side b.

    Returns:
        The (least-squares) solution x.
    """
    dtype = operator_dtype(operator, vector)
    atol = resolve_tolerance(self.atol, dtype, 1e-6)
    btol = resolve_tolerance(self.btol, dtype, 1e-6)
    if self.damp == 0.0:
        solver = lx.LSMR(
            rtol=btol,
            atol=atol,
            max_steps=self.max_steps,
            conlim=1.0 / self.ctol if self.ctol > 0 else 1e8,
        )
        return lx.linear_solve(operator, vector, solver, throw=self.throw).value

    # Tikhonov damping: not supported by lineax's LSMR, so the
    # damped path stays on matfree.
    lsmr_fn = matfree.lstsq.lsmr(
        atol=atol,
        btol=btol,
        ctol=self.ctol,
        maxiter=self.max_steps,
    )

    def vecmat(v):
        return operator.T.mv(v)

    x, stats = lsmr_fn(vecmat, vector, damp=self.damp)
    if self.throw:
        # matfree reports success but never raises; match lineax.
        x = eqx.error_if(
            x,
            jnp.logical_not(stats["success"]),
            "LSMR did not converge within `max_steps` steps. Increase "
            "`max_steps`, loosen `atol`/`btol`, or pass `throw=False`.",
        )
    return x

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log-determinant via Lanczos quadrature.

LSMR itself does not use this: it is SLQ on a square symmetric PSD operator, for the AbstractSolverStrategy interface.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A square symmetric PSD linear operator.	required

Returns:

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

Raises:

Type	Description
`ValueError`	If the operator is not square (gh-402).

Source code in src/gaussx/_strategies/_lsmr.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic log-determinant via Lanczos quadrature.

    LSMR itself does not use this: it is SLQ on a *square* symmetric
    PSD operator, for the `AbstractSolverStrategy` interface.

    Args:
        operator: A square symmetric PSD linear operator.

    Returns:
        Scalar estimate of log |det(A)|.

    Raises:
        ValueError: If the operator is not square (gh-402).
    """
    if operator.in_size() != operator.out_size():
        raise ValueError(
            "LSMRSolver.logdet needs a square operator; got shape "
            f"({operator.out_size()}, {operator.in_size()})."
        )
    return SLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        seed=self.seed,
    ).logdet(operator, key=key)

`BBMMSolver` ¶

CG solve and SLQ logdet with BBMM's defaults (Gardner et al. 2018).

Solve: an ordinary lineax CG solve of each right-hand side. Logdet: stochastic Lanczos quadrature (SLQ) via matfree, with lanczos_order Lanczos steps on each of num_probes probes.

The two are independent: solve_and_logdet is a convenience wrapper that calls solve and then logdet, sharing no matvecs (gh-396). The modified batched CG (mBCG) pass that does share them, running the right-hand sides and the probes through one block CG, is gaussx.inv_quad_logdet(operator, rhs, strategy=BBMMSolver(...)). It is not automatically cheaper: it runs the probe columns to convergence rather than stopping at lanczos_order Lanczos steps, so it can apply the operator to more columns than a separate solve and SLQ do.

Only the integer seed is stored. With no key, logdet draws its probes from PRNGKey(seed) at every call, so it is a deterministic function of the operator: the same probes every time (common random numbers). Pass key or use gaussx.KeyedSolver to vary them.

The options use the names shared by every iterative strategy (gh-405). The old keywords still work, with a deprecation warning: cg_tolerance=t sets rtol=atol=t, cg_max_iter is max_steps and lanczos_iter is lanczos_order.

Attributes:

Name	Type	Description
`rtol`	`float \| None`	Relative tolerance for CG. `None`: `1e-4` in float64, `1e-3` in float32 (gh-327).
`atol`	`float \| None`	Absolute tolerance for CG. `None`: `1e-4` in float64; in float32, `sqrt(eps) * max\|b_i\|`, relative to the right-hand side (the solve runs on `b / max\|b_i\|`; gh-639).
`max_steps`	`int`	Maximum CG iterations.
`lanczos_order`	`int`	Lanczos iterations for SLQ.
`num_probes`	`int`	Number of probe vectors for Hutchinson.
`seed`	`int`	Seed for probe vector generation.
`throw`	`bool`	Raise when CG does not converge within `max_steps`. With `False` the last iterate is returned unchecked (see `gaussx.CGSolver`).

Source code in src/gaussx/_strategies/_bbmm.py

class BBMMSolver(AbstractSolverStrategy):
    """CG solve and SLQ logdet with BBMM's defaults (Gardner et al. 2018).

    Solve: an ordinary lineax CG solve of each right-hand side.
    Logdet: stochastic Lanczos quadrature (SLQ) via matfree, with
    ``lanczos_order`` Lanczos steps on each of ``num_probes`` probes.

    The two are independent: `solve_and_logdet` is a convenience wrapper
    that calls `solve` and then `logdet`, sharing no matvecs (gh-396). The
    modified batched CG (mBCG) pass that *does* share them, running the
    right-hand sides and the probes through one block CG, is
    ``gaussx.inv_quad_logdet(operator, rhs, strategy=BBMMSolver(...))``. It
    is not automatically cheaper: it runs the probe columns to convergence
    rather than stopping at ``lanczos_order`` Lanczos steps, so it can apply
    the operator to more columns than a separate solve and SLQ do.

    Only the integer ``seed`` is stored. With no ``key``, `logdet` draws
    its probes from ``PRNGKey(seed)`` at every call, so it is a
    deterministic function of the operator: the same probes every time
    (common random numbers). Pass ``key`` or use `gaussx.KeyedSolver` to
    vary them.

    The options use the names shared by every iterative strategy (gh-405).
    The old keywords still work, with a deprecation warning:
    ``cg_tolerance=t`` sets ``rtol=atol=t``, ``cg_max_iter`` is
    ``max_steps`` and ``lanczos_iter`` is ``lanczos_order``.

    Attributes:
        rtol: Relative tolerance for CG. ``None``: ``1e-4`` in float64,
            ``1e-3`` in float32 (gh-327).
        atol: Absolute tolerance for CG. ``None``: ``1e-4`` in float64; in
            float32, ``sqrt(eps) * max|b_i|``, relative to the right-hand
            side (the solve runs on ``b / max|b_i|``; gh-639).
        max_steps: Maximum CG iterations.
        lanczos_order: Lanczos iterations for SLQ.
        num_probes: Number of probe vectors for Hutchinson.
        seed: Seed for probe vector generation.
        throw: Raise when CG does not converge within ``max_steps``. With
            ``False`` the last iterate is returned unchecked (see
            `gaussx.CGSolver`).
    """

    rtol: float | None = eqx.field(static=True, default=None)
    atol: float | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=1000)
    lanczos_order: int = eqx.field(static=True, default=100)
    num_probes: int = eqx.field(static=True, default=10)
    seed: int = eqx.field(static=True, default=0)
    throw: bool = eqx.field(static=True, default=True)

    def __init__(
        self,
        *,
        rtol: float | None = UNSET,
        atol: float | None = UNSET,
        max_steps: int = UNSET,
        lanczos_order: int = UNSET,
        num_probes: int = 10,
        seed: int = 0,
        throw: bool = True,
        cg_tolerance: float | None = UNSET,
        cg_max_iter: int = UNSET,
        lanczos_iter: int = UNSET,
    ) -> None:
        name = "BBMMSolver"
        rtol = renamed(name, "cg_tolerance", cg_tolerance, "rtol", rtol)
        if cg_tolerance is not UNSET:
            # cg_tolerance set both tolerances.
            if atol is not UNSET:
                raise TypeError(f"{name}: pass rtol= and atol=, not cg_tolerance=.")
            atol = cg_tolerance
        self.rtol = default(rtol, None)
        self.atol = default(atol, None)
        self.max_steps = default(
            renamed(name, "cg_max_iter", cg_max_iter, "max_steps", max_steps), 1000
        )
        self.lanczos_order = default(
            renamed(name, "lanczos_iter", lanczos_iter, "lanczos_order", lanczos_order),
            100,
        )
        self.num_probes = num_probes
        self.seed = seed
        self.throw = throw

    @property
    def cg_tolerance(self) -> float | None:
        """Deprecated: `rtol` (gh-405)."""
        warn_deprecated(
            "BBMMSolver.cg_tolerance is deprecated and will be removed in gaussx "
            "0.7.0; use .rtol (gh-405)."
        )
        return self.rtol

    @property
    def cg_max_iter(self) -> int:
        """Deprecated: `max_steps` (gh-405)."""
        warn_deprecated(
            "BBMMSolver.cg_max_iter is deprecated and will be removed in gaussx "
            "0.7.0; use .max_steps (gh-405)."
        )
        return self.max_steps

    @property
    def lanczos_iter(self) -> int:
        """Deprecated: `lanczos_order` (gh-405)."""
        warn_deprecated(
            "BBMMSolver.lanczos_iter is deprecated and will be removed in gaussx "
            "0.7.0; use .lanczos_order (gh-405)."
        )
        return self.lanczos_order

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve A x = b via CG.

        Args:
            operator: A PSD linear operator.
            vector: The right-hand side b.

        Returns:
            The solution x.
        """
        dtype = operator_dtype(operator, vector)
        atol, scale = rhs_scaling(self.atol, dtype, 1e-4, vector)
        rhs = vector if scale is None else vector / scale
        solver = lx.CG(
            rtol=resolve_tolerance(self.rtol, dtype, 1e-4),
            atol=atol,
            max_steps=self.max_steps,
        )
        x = lx.linear_solve(operator, rhs, solver, throw=self.throw).value
        return x if scale is None else x * scale

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic log-determinant via Lanczos quadrature.

        Probe vectors are generated deterministically from ``self.seed``.

        Args:
            operator: A PSD linear operator.

        Returns:
            Scalar estimate of log |det(A)|.
        """
        return SLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            seed=self.seed,
        ).logdet(operator, key=key)

    def solve_and_logdet(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> tuple[Float[Array, " n"], Float[Array, ""]]:
        """``(solve(A, b), logdet(A))`` in one call.

        A convenience wrapper: it costs exactly `solve` plus `logdet`, with
        no shared matvecs. For the shared mBCG pass use
        ``gaussx.inv_quad_logdet(operator, rhs, strategy=self)``.

        Args:
            operator: A PSD linear operator.
            vector: The right-hand side b.

        Returns:
            Tuple of (solution, log_determinant).
        """
        return self.solve(operator, vector), self.logdet(operator)

`cg_tolerance: float | None` `property` ¶

Deprecated: rtol (gh-405).

`cg_max_iter: int` `property` ¶

Deprecated: max_steps (gh-405).

`lanczos_iter: int` `property` ¶

Deprecated: lanczos_order (gh-405).

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b via CG.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`vector`	`Float[Array, ' n']`	The right-hand side b.	required

Returns:

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

Source code in src/gaussx/_strategies/_bbmm.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve A x = b via CG.

    Args:
        operator: A PSD linear operator.
        vector: The right-hand side b.

    Returns:
        The solution x.
    """
    dtype = operator_dtype(operator, vector)
    atol, scale = rhs_scaling(self.atol, dtype, 1e-4, vector)
    rhs = vector if scale is None else vector / scale
    solver = lx.CG(
        rtol=resolve_tolerance(self.rtol, dtype, 1e-4),
        atol=atol,
        max_steps=self.max_steps,
    )
    x = lx.linear_solve(operator, rhs, solver, throw=self.throw).value
    return x if scale is None else x * scale

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log-determinant via Lanczos quadrature.

Probe vectors are generated deterministically from self.seed.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required

Returns:

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

Source code in src/gaussx/_strategies/_bbmm.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic log-determinant via Lanczos quadrature.

    Probe vectors are generated deterministically from ``self.seed``.

    Args:
        operator: A PSD linear operator.

    Returns:
        Scalar estimate of log |det(A)|.
    """
    return SLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        seed=self.seed,
    ).logdet(operator, key=key)

`solve_and_logdet(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> tuple[Float[Array, ' n'], Float[Array, '']]` ¶

(solve(A, b), logdet(A)) in one call.

A convenience wrapper: it costs exactly solve plus logdet, with no shared matvecs. For the shared mBCG pass use gaussx.inv_quad_logdet(operator, rhs, strategy=self).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`vector`	`Float[Array, ' n']`	The right-hand side b.	required

Returns:

Type	Description
`tuple[Float[Array, ' n'], Float[Array, '']]`	Tuple of (solution, log_determinant).

Source code in src/gaussx/_strategies/_bbmm.py

def solve_and_logdet(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> tuple[Float[Array, " n"], Float[Array, ""]]:
    """``(solve(A, b), logdet(A))`` in one call.

    A convenience wrapper: it costs exactly `solve` plus `logdet`, with
    no shared matvecs. For the shared mBCG pass use
    ``gaussx.inv_quad_logdet(operator, rhs, strategy=self)``.

    Args:
        operator: A PSD linear operator.
        vector: The right-hand side b.

    Returns:
        Tuple of (solution, log_determinant).
    """
    return self.solve(operator, vector), self.logdet(operator)

`ComposedSolver` ¶

Mix-and-match solve and logdet from different strategies.

This lets you pair, e.g., an exact dense solve with a stochastic log-determinant estimator, or an iterative CG solve with a closed-form Kronecker log-determinant.

Accepts either fine-grained protocols (AbstractSolveStrategy, AbstractLogdetStrategy) or full solver strategies.

Attributes:

Name	Type	Description
`solve_strategy`	`AbstractSolveStrategy`	Strategy whose `.solve()` method will be used.
`logdet_strategy`	`AbstractLogdetStrategy`	Strategy whose `.logdet()` method will be used.

Examples:

solver = ComposedSolver(
    solve_strategy=DenseSolver(),
    logdet_strategy=SLQLogdet(num_probes=50, lanczos_order=30),
)

Source code in src/gaussx/_strategies/_composed.py

class ComposedSolver(AbstractSolverStrategy):
    """Mix-and-match solve and logdet from different strategies.

    This lets you pair, e.g., an exact dense solve with a stochastic
    log-determinant estimator, or an iterative CG solve with a
    closed-form Kronecker log-determinant.

    Accepts either fine-grained protocols (`AbstractSolveStrategy`,
    `AbstractLogdetStrategy`) or full solver strategies.

    Attributes:
        solve_strategy: Strategy whose ``.solve()`` method will be used.
        logdet_strategy: Strategy whose ``.logdet()`` method will be used.

    Examples:

        solver = ComposedSolver(
            solve_strategy=DenseSolver(),
            logdet_strategy=SLQLogdet(num_probes=50, lanczos_order=30),
        )
    """

    solve_strategy: AbstractSolveStrategy
    logdet_strategy: AbstractLogdetStrategy

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``A x = b`` by delegating to ``solve_strategy``.

        Args:
            operator: Linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        return self.solve_strategy.solve(operator, vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Compute ``log|det(A)|`` by delegating to ``logdet_strategy``.

        Args:
            operator: Linear operator ``A``.
            key: Optional PRNG key forwarded to stochastic estimators.

        Returns:
            Scalar log-determinant.
        """
        return self.logdet_strategy.logdet(operator, key=key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b by delegating to solve_strategy.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_composed.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``A x = b`` by delegating to ``solve_strategy``.

    Args:
        operator: Linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    return self.solve_strategy.solve(operator, vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Compute log|det(A)| by delegating to logdet_strategy.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`key`	`Array \| None`	Optional PRNG key forwarded to stochastic estimators.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_composed.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Compute ``log|det(A)|`` by delegating to ``logdet_strategy``.

    Args:
        operator: Linear operator ``A``.
        key: Optional PRNG key forwarded to stochastic estimators.

    Returns:
        Scalar log-determinant.
    """
    return self.logdet_strategy.logdet(operator, key=key)

`KeyedSolver` ¶

Bind a PRNG key to a strategy's stochastic log-determinant (gh-384).

A stochastic logdet strategy (gaussx.SLQLogdet, gaussx.CGSolver, gaussx.BBMMSolver, ...) called without a key draws its probes from PRNGKey(seed), so every call sees the same probes: common random numbers. That suits stochastic-gradient training (the objective is a smooth, fixed function of the parameters), but averaging such estimates reduces no variance, and MCMC then targets a fixed pseudo-likelihood.

Distribution methods such as MultivariateNormal.log_prob cannot take a key under the numpyro contract. Wrap the strategy instead and pass a fresh key per step. The key is a PyTree leaf, so it can change under jax.jit without recompiling:

solver = gaussx.KeyedSolver(gaussx.CGSolver(), key)
mvn = gaussx.MultivariateNormal(loc, cov, solver=solver)

The functional API (gaussx.gaussian_log_prob, gaussx.gaussian_entropy, gaussx.kl_standard_normal) takes key= directly.

Attributes:

Name	Type	Description
`strategy`	`AbstractLogdetStrategy`	The wrapped strategy. `solve` needs it to be a solve strategy as well.
`key`	`Array`	The PRNG key passed to `strategy.logdet`. A key passed to `logdet` explicitly takes precedence.

Source code in src/gaussx/_strategies/_keyed.py

class KeyedSolver(AbstractSolverStrategy):
    """Bind a PRNG key to a strategy's stochastic log-determinant (gh-384).

    A stochastic logdet strategy (`gaussx.SLQLogdet`, `gaussx.CGSolver`,
    `gaussx.BBMMSolver`, ...) called without a ``key`` draws its probes from
    ``PRNGKey(seed)``, so every call sees *the same* probes: common random
    numbers. That suits stochastic-gradient training (the objective is a
    smooth, fixed function of the parameters), but averaging such estimates
    reduces no variance, and MCMC then targets a fixed pseudo-likelihood.

    Distribution methods such as ``MultivariateNormal.log_prob`` cannot take
    a key under the numpyro contract. Wrap the strategy instead and pass a
    fresh key per step. The key is a PyTree leaf, so it can change under
    ``jax.jit`` without recompiling:

        solver = gaussx.KeyedSolver(gaussx.CGSolver(), key)
        mvn = gaussx.MultivariateNormal(loc, cov, solver=solver)

    The functional API (`gaussx.gaussian_log_prob`,
    `gaussx.gaussian_entropy`, `gaussx.kl_standard_normal`) takes
    ``key=`` directly.

    Attributes:
        strategy: The wrapped strategy. ``solve`` needs it to be a solve
            strategy as well.
        key: The PRNG key passed to ``strategy.logdet``. A key passed to
            `logdet` explicitly takes precedence.
    """

    strategy: AbstractLogdetStrategy
    key: jax.Array

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``A x = b`` with the wrapped strategy.

        Args:
            operator: Linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.

        Raises:
            TypeError: If the wrapped strategy has no ``solve``.
        """
        if not isinstance(self.strategy, AbstractSolveStrategy):
            raise TypeError(
                f"{type(self.strategy).__name__} is a logdet-only strategy; "
                "wrap a solver strategy (or a ComposedSolver) to solve."
            )
        return self.strategy.solve(operator, vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """``log|det(A)|`` from the wrapped strategy, with the bound key.

        Args:
            operator: Linear operator ``A``.
            key: Overrides the bound key when given.

        Returns:
            Scalar log-determinant.
        """
        return self.strategy.logdet(operator, key=self.key if key is None else key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b with the wrapped strategy.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Raises:

Type	Description
`TypeError`	If the wrapped strategy has no `solve`.

Source code in src/gaussx/_strategies/_keyed.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``A x = b`` with the wrapped strategy.

    Args:
        operator: Linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.

    Raises:
        TypeError: If the wrapped strategy has no ``solve``.
    """
    if not isinstance(self.strategy, AbstractSolveStrategy):
        raise TypeError(
            f"{type(self.strategy).__name__} is a logdet-only strategy; "
            "wrap a solver strategy (or a ComposedSolver) to solve."
        )
    return self.strategy.solve(operator, vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

log|det(A)| from the wrapped strategy, with the bound key.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`key`	`Array \| None`	Overrides the bound key when given.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_keyed.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """``log|det(A)|`` from the wrapped strategy, with the bound key.

    Args:
        operator: Linear operator ``A``.
        key: Overrides the bound key when given.

    Returns:
        Scalar log-determinant.
    """
    return self.strategy.logdet(operator, key=self.key if key is None else key)

`LineaxSolver` ¶

Bases: AbstractSolveStrategy

A lineax solver as a gaussx solve strategy.

solver= means a gaussx strategy everywhere except the fallback of gaussx.solve, which historically takes a lineax solver. Both gaussx.linear_solve and gaussx.dispatch_solve wrap a bare lineax.AbstractLinearSolver in this class, so either kind works in either place.

Attributes:

Name	Type	Description
`solver`	`AbstractLinearSolver`	The lineax solver, e.g. `lineax.CG(rtol=1e-6, atol=1e-6)`. Static configuration, like every strategy's.

Source code in src/gaussx/_strategies/_lineax.py

class LineaxSolver(AbstractSolveStrategy):
    """A lineax solver as a gaussx solve strategy.

    ``solver=`` means a gaussx strategy everywhere except the fallback of
    `gaussx.solve`, which historically takes a lineax solver. Both
    `gaussx.linear_solve` and `gaussx.dispatch_solve` wrap a bare
    ``lineax.AbstractLinearSolver`` in this class, so either kind works in
    either place.

    Attributes:
        solver: The lineax solver, e.g. ``lineax.CG(rtol=1e-6, atol=1e-6)``.
            Static configuration, like every strategy's.
    """

    solver: lx.AbstractLinearSolver = eqx.field(static=True)

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """``lineax.linear_solve(operator, vector, solver).value``.

        Args:
            operator: Linear operator ``A``.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        return lx.linear_solve(operator, vector, self.solver).value

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

lineax.linear_solve(operator, vector, solver).value.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Linear operator `A`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_lineax.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """``lineax.linear_solve(operator, vector, solver).value``.

    Args:
        operator: Linear operator ``A``.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    return lx.linear_solve(operator, vector, self.solver).value

`SparseCholeskySolver` ¶

Exact solve / logdet / diag_inv through a sparse Cholesky factor.

For a symmetric positive-definite SparseOperator (a GMRF precision, a Laplace Hessian Q + AᵀWA). The symbolic analysis is cached per sparsity pattern, so only the numeric factorisation runs per call; its gradients are exact (Takahashi for logdet, an adjoint solve for solve). Pass it explicitly: AutoSolver keeps choosing CG for large PSD operators, and nothing picks this strategy by size.

Attributes:

Name	Type	Description
`ordering`	`Ordering`	Fill-reducing ordering, `"rcm"`, `"natural"` or `"amd"` (CHOLMOD). See `gaussx.symbolic_cholesky`.
`backend`	`Backend`	`"jax"` or `"cholmod"` (opt-in, needs scikit-sparse).

Examples:

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

# Path-graph Laplacian plus a unit shift
n = 5
Q = gaussx.SparseOperator.from_coo(
    np.r_[np.arange(n), np.arange(1, n)],
    np.r_[np.arange(n), np.arange(n - 1)],
    jnp.r_[jnp.array([2.0, 3.0, 3.0, 3.0, 2.0]), -jnp.ones(n - 1)],
    (n, n),
    symmetric=True,
    tags=lx.positive_semidefinite_tag,
)
solver = gaussx.SparseCholeskySolver()
x = solver.solve(Q, jnp.ones(n))
ld = solver.logdet(Q)
variances = gaussx.diag_inv(Q, solver=solver)  # Takahashi

Source code in src/gaussx/_strategies/_sparse_cholesky.py

class SparseCholeskySolver(AbstractSolverStrategy):
    """Exact ``solve`` / ``logdet`` / ``diag_inv`` through a sparse Cholesky factor.

    For a symmetric positive-definite `SparseOperator` (a GMRF precision, a
    Laplace Hessian ``Q + AᵀWA``). The symbolic analysis is cached per
    sparsity pattern, so only the numeric factorisation runs per call; its
    gradients are exact (Takahashi for ``logdet``, an adjoint solve for
    ``solve``). Pass it explicitly: `AutoSolver` keeps choosing CG for large
    PSD operators, and nothing picks this strategy by size.

    Attributes:
        ordering: Fill-reducing ordering, ``"rcm"``, ``"natural"`` or
            ``"amd"`` (CHOLMOD). See `gaussx.symbolic_cholesky`.
        backend: ``"jax"`` or ``"cholmod"`` (opt-in, needs scikit-sparse).

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

        # Path-graph Laplacian plus a unit shift
        n = 5
        Q = gaussx.SparseOperator.from_coo(
            np.r_[np.arange(n), np.arange(1, n)],
            np.r_[np.arange(n), np.arange(n - 1)],
            jnp.r_[jnp.array([2.0, 3.0, 3.0, 3.0, 2.0]), -jnp.ones(n - 1)],
            (n, n),
            symmetric=True,
            tags=lx.positive_semidefinite_tag,
        )
        solver = gaussx.SparseCholeskySolver()
        x = solver.solve(Q, jnp.ones(n))
        ld = solver.logdet(Q)
        variances = gaussx.diag_inv(Q, solver=solver)  # Takahashi
        ```
    """

    ordering: Ordering = eqx.field(static=True, default="rcm")
    backend: Backend = eqx.field(static=True, default="jax")

    def factor(self, operator: lx.AbstractLinearOperator) -> SparseCholeskyFactor:
        """The sparse Cholesky factor of ``operator``.

        Args:
            operator: A `SparseOperator` (or a lineax-tagged one).

        Returns:
            The factor, on the cached symbolic analysis of its pattern.

        Raises:
            TypeError: If ``operator`` is not a `SparseOperator`.
        """
        if isinstance(operator, lx.TaggedLinearOperator):
            operator = operator.operator
        if not isinstance(operator, SparseOperator):
            raise TypeError(
                "SparseCholeskySolver needs a SparseOperator, got "
                f"{type(operator).__name__}."
            )
        symbolic = symbolic_cholesky(
            operator.pattern, ordering=self.ordering, backend=self.backend
        )
        return sparse_cholesky(operator, symbolic)

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " n"],
    ) -> Float[Array, " n"]:
        """Solve ``A x = b`` with two sparse triangular solves.

        Args:
            operator: A symmetric positive-definite `SparseOperator`.
            vector: Right-hand side ``b``, shape ``(n,)``.

        Returns:
            Solution ``x``, shape ``(n,)``.
        """
        return self.factor(operator).solve(vector)

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Exact ``log|A| = 2 Σ log L_jj``.

        Args:
            operator: A symmetric positive-definite `SparseOperator`.
            key: Unused; present for protocol compatibility.

        Returns:
            Scalar log-determinant.
        """
        return self.factor(operator).logdet()

    def diag_inv(self, operator: lx.AbstractLinearOperator) -> Float[Array, " n"]:
        """Exact ``diag(A⁻¹)`` (marginal variances) by one Takahashi sweep.

        Args:
            operator: A symmetric positive-definite `SparseOperator`.

        Returns:
            Shape ``(n,)``.
        """
        return self.factor(operator).diag_inv()

`factor(operator: lx.AbstractLinearOperator) -> SparseCholeskyFactor` ¶

The sparse Cholesky factor of operator.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A `SparseOperator` (or a lineax-tagged one).	required

Returns:

Type	Description
`SparseCholeskyFactor`	The factor, on the cached symbolic analysis of its pattern.

Raises:

Type	Description
`TypeError`	If `operator` is not a `SparseOperator`.

Source code in src/gaussx/_strategies/_sparse_cholesky.py

def factor(self, operator: lx.AbstractLinearOperator) -> SparseCholeskyFactor:
    """The sparse Cholesky factor of ``operator``.

    Args:
        operator: A `SparseOperator` (or a lineax-tagged one).

    Returns:
        The factor, on the cached symbolic analysis of its pattern.

    Raises:
        TypeError: If ``operator`` is not a `SparseOperator`.
    """
    if isinstance(operator, lx.TaggedLinearOperator):
        operator = operator.operator
    if not isinstance(operator, SparseOperator):
        raise TypeError(
            "SparseCholeskySolver needs a SparseOperator, got "
            f"{type(operator).__name__}."
        )
    symbolic = symbolic_cholesky(
        operator.pattern, ordering=self.ordering, backend=self.backend
    )
    return sparse_cholesky(operator, symbolic)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' n']) -> Float[Array, ' n']` ¶

Solve A x = b with two sparse triangular solves.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric positive-definite `SparseOperator`.	required
`vector`	`Float[Array, ' n']`	Right-hand side `b`, shape `(n,)`.	required

Returns:

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

Source code in src/gaussx/_strategies/_sparse_cholesky.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " n"],
) -> Float[Array, " n"]:
    """Solve ``A x = b`` with two sparse triangular solves.

    Args:
        operator: A symmetric positive-definite `SparseOperator`.
        vector: Right-hand side ``b``, shape ``(n,)``.

    Returns:
        Solution ``x``, shape ``(n,)``.
    """
    return self.factor(operator).solve(vector)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Exact log|A| = 2 Σ log L_jj.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric positive-definite `SparseOperator`.	required
`key`	`Array \| None`	Unused; present for protocol compatibility.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_sparse_cholesky.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Exact ``log|A| = 2 Σ log L_jj``.

    Args:
        operator: A symmetric positive-definite `SparseOperator`.
        key: Unused; present for protocol compatibility.

    Returns:
        Scalar log-determinant.
    """
    return self.factor(operator).logdet()

`diag_inv(operator: lx.AbstractLinearOperator) -> Float[Array, ' n']` ¶

Exact diag(A⁻¹) (marginal variances) by one Takahashi sweep.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric positive-definite `SparseOperator`.	required

Returns:

Type	Description
`Float[Array, ' n']`	Shape `(n,)`.

Source code in src/gaussx/_strategies/_sparse_cholesky.py

def diag_inv(self, operator: lx.AbstractLinearOperator) -> Float[Array, " n"]:
    """Exact ``diag(A⁻¹)`` (marginal variances) by one Takahashi sweep.

    Args:
        operator: A symmetric positive-definite `SparseOperator`.

    Returns:
        Shape ``(n,)``.
    """
    return self.factor(operator).diag_inv()

Tolerances and float32¶

The iterative strategies (CGSolver, PreconditionedCGSolver, BBMMSolver, MINRESSolver, LSMRSolver) leave their tolerances unset by default and resolve them from the operator's dtype at solve time: the historical constants in float64 (1e-5; 1e-4 for BBMM; 1e-6 for LSMR) and 1e-3 in float32. A float32 CG solve cannot reach a relative residual of 1e-5 once the condition number passes about \(10^3\). It runs out of steps, or breaks down, and returns an iterate worse than zero. Set rtol/atol explicitly to override.

CGSolver, PreconditionedCGSolver, BBMMSolver and AutoSolver raise when CG exhausts its step budget. With throw=False they return the last iterate unchecked instead. Use that only where the caller checks the result, because an unconverged CG iterate can be far from the solution.

Probe keys and common random numbers¶

A stochastic log-determinant (SLQ, used by CGSolver, PreconditionedCGSolver, BBMMSolver, MINRESSolver, LSMRSolver and a large PSD AutoSolver) called without a key draws its probes from PRNGKey(seed), so every call sees the same probes. These are common random numbers: right for stochastic-gradient training, where the objective becomes a fixed smooth function of the parameters, but averaging such estimates reduces no variance, and MCMC then targets a fixed pseudo-likelihood. Pass key= to gaussian_log_prob, gaussian_entropy or kl_standard_normal. Distribution methods cannot take one, so wrap their strategy in KeyedSolver, whose key is a PyTree leaf, or change seed.

Sketch-and-precondition least squares¶

SketchAndPrecondLSMR solves tall (optionally ridge-damped) least squares \(\min_x \|Ax - b\|^2 + \delta^2\|x\|^2\), \(m \gg n\), in a number of LSMR steps that does not grow with \(m\) or \(\kappa(A)\). A sketch \(S \in \mathbb{R}^{d \times m}\) (Sketching) of \(A\) gives \([SA;\ \delta I] = QR\), and \(M = R^{-1}\) is a right preconditioner with

\[ \kappa(\tilde A M) \le \frac{1+\varepsilon}{1-\varepsilon}, \qquad \tilde A = [A;\ \delta I], \]

when \(S\) is an \(\varepsilon\)-embedding for \(\operatorname{range}(A)\). The sketch-and-solve estimate \(x_0 = R^{-1}Q_{1:d}^\top Sb\) is the warm start; sketch_and_solve returns it on its own (an \(O(\varepsilon)\) approximation). The dense QR costs \(O(dn^2)\), so this targets \(n \lesssim 10^4\); for larger \(n\), solve the normal equations with PreconditionedCGSolver and a NystromPreconditioner on \(A^\top A\) with shift=δ².

# Gauss–Newton step on a (10⁶ pixels × 300 parameters) Jacobian J_op
step = gx.SketchAndPrecondLSMR(damp=1e-3).solve(J_op, -residuals)

Structured linear algebra and Gaussian primitives for JAX.

`SketchAndPrecondLSMR` ¶

Sketch-and-precondition LSMR for tall least squares (Rokhlin & Tygert, 2008).

Solves

\[ \min_x \|Ax - b\|^2 + \delta^2 \|x\|^2, \qquad A \in \mathbb{R}^{m \times n},\ m \gg n, \]

i.e. plain least squares on the augmented system \(\tilde A = [A;\ \delta I]\), \(\tilde b = [b;\ 0]\) (\(\tilde A = A\) when \(\delta = 0\)). A sketch \(S \in \mathbb{R}^{d \times m}\), \(d = \lceil \gamma n \rceil\), gives the QR factorisation \([SA;\ \delta I] = QR\), and \(M = R^{-1}\) is a right preconditioner: if \(S\) is an \(\varepsilon\)-embedding for \(\operatorname{range}(A)\), then \(\|[SA;\ \delta I]\,My\| = \|y\|\) and

\[ \|\tilde A M y\| \in \Big[\tfrac{1}{1+\varepsilon}, \tfrac{1}{1-\varepsilon}\Big]\|y\| \quad\Longrightarrow\quad \kappa(\tilde A M) \le \frac{1+\varepsilon}{1-\varepsilon}, \]

so LSMR on \(\tilde A M\) contracts the error by about \((\kappa - 1)/(\kappa + 1)\) per step, independently of \(m\) and of \(\kappa(A)\). The default \(\gamma = 4\) gives \(\varepsilon \approx \sqrt{1/\gamma} = 0.5\) for a Gaussian sketch (and comparable for the sparse sign sketch), so \(\kappa \lesssim 3\) and LSMR converges in a few tens of steps (Meng, Saunders & Mahoney, 2014; Epperly, 2024).

S = sample(sketch, d = ceil(gamma n), m; PRNGKey(seed))
Q, R = qr([S A; δ I])                    # O(d n²)
x₀ = R⁻¹ Q[:d]ᵀ (S b)                    # sketch-and-solve warm start
y = LSMR(Ã R⁻¹, [b; 0] − Ã x₀)           # undamped, lineax
return x₀ + R⁻¹ y

The inner solve is lineax.LSMR on the composed operator \(\tilde A M\) (the undamped path of LSMRSolver), so gradients use lineax's implicit differentiation. Only \(A\) is sketched; the ridge block is exact.

With damp == 0, \(A\) must have full column rank (otherwise the solution is not unique and \(R\) is singular; throw=True raises); damp > 0 always has a unique solution. A matrix-free \(A\) is sketched with \(d\) transpose-matvecs, a batch of rows of \(S\) at a time, so the extra memory is \(O(m)\) per batch, not \(O(dm)\).

Scope. The dense QR of the \(d \times n\) sketch costs \(O(d n^2)\) and \(O(dn)\) memory, so this targets \(n \lesssim 10^4\) (e.g. Gauss-Newton steps on a Jacobian with \(10^6\) pixels and a few hundred parameters). For larger \(n\), solve the normal equations \((A^\top A + \delta^2 I)x = A^\top b\) with PreconditionedCGSolver and a NystromPreconditioner built from \(A^\top A\) with shift=δ².

Attributes:

Name	Type	Description
`sketch`	`SketchName`	`"sparse_sign"` (default; `SparseSignSketch`), `"srht"` (`SRHTSketch`) or `"gaussian"` (`GaussianSketch`).
`sampling_factor`	`float`	\(\gamma \ge 1\); the sketch has \(d = \min(\lceil \gamma n \rceil, m)\) rows.
`nnz`	`int`	Non-zeros per column of the sparse sign sketch.
`damp`	`float`	Ridge parameter \(\delta \ge 0\).
`atol`	`float \| None`	LSMR tolerance on the normal-equations residual. `None`: `1e-6` in float64, `1e-3` in float32, as `LSMRSolver` (gh-327). The preconditioned system has \(\kappa \lesssim 3\), so `1e-6` is reachable in float32 too and costs only a few more steps; set it explicitly when float32 accuracy matters.
`btol`	`float \| None`	LSMR relative tolerance on the residual. `None`: as `atol`.
`max_steps`	`int`	Maximum LSMR iterations.
`seed`	`int`	The sketch is drawn from `jax.random.PRNGKey(seed)`; it is also the probe seed of `logdet`.
`throw`	`bool`	Raise when LSMR stops without converging, as `LSMRSolver`.

References

Rokhlin, V. & Tygert, M. (2008). A fast randomized algorithm for overdetermined linear least-squares regression. PNAS, 105(36), 13212-13217.

Meng, X., Saunders, M. A. & Mahoney, M. W. (2014). LSRN: A parallel iterative solver for strongly over- or underdetermined systems. SIAM J. Sci. Comput., 36(2), C95-C118.

Epperly, E. N. (2024). Fast and forward stable randomized algorithms for linear least-squares problems. SIAM J. Matrix Anal. Appl., 45(4), 1782-1804.

Examples:

>>> import jax.numpy as jnp, jax.random as jr, lineax as lx
>>> import gaussx as gx
>>> A = jr.normal(jr.key(0), (2000, 10))
>>> b = jr.normal(jr.key(1), (2000,))
>>> solver = gx.SketchAndPrecondLSMR(atol=1e-6, btol=1e-6)
>>> x = solver.solve(lx.MatrixLinearOperator(A), b)
>>> bool(jnp.allclose(x, jnp.linalg.lstsq(A, b)[0], atol=1e-4))
True

Source code in src/gaussx/_strategies/_sketch_precond.py

class SketchAndPrecondLSMR(AbstractSolverStrategy):
    r"""Sketch-and-precondition LSMR for tall least squares (Rokhlin & Tygert, 2008).

    Solves

    $$
    \min_x \|Ax - b\|^2 + \delta^2 \|x\|^2, \qquad A \in \mathbb{R}^{m
    \times n},\ m \gg n,
    $$

    i.e. plain least squares on the augmented system $\tilde A = [A;\
    \delta I]$, $\tilde b = [b;\ 0]$ ($\tilde A = A$ when $\delta = 0$). A
    sketch $S \in \mathbb{R}^{d \times m}$, $d = \lceil \gamma n \rceil$,
    gives the QR factorisation $[SA;\ \delta I] = QR$, and $M = R^{-1}$ is a
    right preconditioner: if $S$ is an $\varepsilon$-embedding for
    $\operatorname{range}(A)$, then $\|[SA;\ \delta I]\,My\| = \|y\|$ and

    $$
    \|\tilde A M y\| \in \Big[\tfrac{1}{1+\varepsilon},
    \tfrac{1}{1-\varepsilon}\Big]\|y\|
    \quad\Longrightarrow\quad
    \kappa(\tilde A M) \le \frac{1+\varepsilon}{1-\varepsilon},
    $$

    so LSMR on $\tilde A M$ contracts the error by about
    $(\kappa - 1)/(\kappa + 1)$ per step, independently of $m$ and of
    $\kappa(A)$. The default $\gamma = 4$ gives $\varepsilon \approx
    \sqrt{1/\gamma} = 0.5$ for a Gaussian sketch (and comparable for the
    sparse sign sketch), so $\kappa \lesssim 3$ and LSMR converges in a few
    tens of steps (Meng, Saunders & Mahoney, 2014; Epperly, 2024).

    ```text
    S = sample(sketch, d = ceil(gamma n), m; PRNGKey(seed))
    Q, R = qr([S A; δ I])                    # O(d n²)
    x₀ = R⁻¹ Q[:d]ᵀ (S b)                    # sketch-and-solve warm start
    y = LSMR(Ã R⁻¹, [b; 0] − Ã x₀)           # undamped, lineax
    return x₀ + R⁻¹ y
    ```

    The inner solve is `lineax.LSMR` on the composed operator $\tilde A M$
    (the undamped path of `LSMRSolver`), so gradients use lineax's implicit
    differentiation. Only $A$ is sketched; the ridge block is exact.

    With ``damp == 0``, $A$ must have full column rank (otherwise the
    solution is not unique and $R$ is singular; ``throw=True`` raises);
    ``damp > 0`` always has a unique solution. A matrix-free $A$ is
    sketched with $d$ transpose-matvecs, a batch of rows of $S$ at a time,
    so the extra memory is $O(m)$ per batch, not $O(dm)$.

    **Scope.** The dense QR of the $d \times n$ sketch costs $O(d n^2)$ and
    $O(dn)$ memory, so this targets $n \lesssim 10^4$ (e.g. Gauss-Newton
    steps on a Jacobian with $10^6$ pixels and a few hundred parameters).
    For larger $n$, solve the normal equations $(A^\top A + \delta^2 I)x =
    A^\top b$ with `PreconditionedCGSolver` and a `NystromPreconditioner`
    built from $A^\top A$ with ``shift=δ²``.

    Attributes:
        sketch: ``"sparse_sign"`` (default; `SparseSignSketch`), ``"srht"``
            (`SRHTSketch`) or ``"gaussian"`` (`GaussianSketch`).
        sampling_factor: $\gamma \ge 1$; the sketch has $d = \min(\lceil
            \gamma n \rceil, m)$ rows.
        nnz: Non-zeros per column of the sparse sign sketch.
        damp: Ridge parameter $\delta \ge 0$.
        atol: LSMR tolerance on the normal-equations residual. ``None``:
            ``1e-6`` in float64, ``1e-3`` in float32, as `LSMRSolver`
            (gh-327). The preconditioned system has $\kappa \lesssim 3$, so
            ``1e-6`` is reachable in float32 too and costs only a few more
            steps; set it explicitly when float32 accuracy matters.
        btol: LSMR relative tolerance on the residual. ``None``: as
            ``atol``.
        max_steps: Maximum LSMR iterations.
        seed: The sketch is drawn from ``jax.random.PRNGKey(seed)``; it is
            also the probe seed of `logdet`.
        throw: Raise when LSMR stops without converging, as `LSMRSolver`.

    References:
        Rokhlin, V. & Tygert, M. (2008). A fast randomized algorithm for
        overdetermined linear least-squares regression. *PNAS*, 105(36),
        13212-13217.

        Meng, X., Saunders, M. A. & Mahoney, M. W. (2014). LSRN: A parallel
        iterative solver for strongly over- or underdetermined systems.
        *SIAM J. Sci. Comput.*, 36(2), C95-C118.

        Epperly, E. N. (2024). Fast and forward stable randomized algorithms
        for linear least-squares problems. *SIAM J. Matrix Anal. Appl.*,
        45(4), 1782-1804.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr, lineax as lx
        >>> import gaussx as gx
        >>> A = jr.normal(jr.key(0), (2000, 10))
        >>> b = jr.normal(jr.key(1), (2000,))
        >>> solver = gx.SketchAndPrecondLSMR(atol=1e-6, btol=1e-6)
        >>> x = solver.solve(lx.MatrixLinearOperator(A), b)
        >>> bool(jnp.allclose(x, jnp.linalg.lstsq(A, b)[0], atol=1e-4))
        True
    """

    sketch: SketchName = eqx.field(static=True, default="sparse_sign")
    sampling_factor: float = eqx.field(static=True, default=4.0)
    nnz: int = eqx.field(static=True, default=8)
    damp: float = eqx.field(static=True, default=0.0)
    atol: float | None = eqx.field(static=True, default=None)
    btol: float | None = eqx.field(static=True, default=None)
    max_steps: int = eqx.field(static=True, default=100)
    seed: int = eqx.field(static=True, default=0)
    throw: bool = eqx.field(static=True, default=True)

    def __check_init__(self) -> None:
        if self.sketch not in ("sparse_sign", "srht", "gaussian"):
            raise ValueError(
                "SketchAndPrecondLSMR.sketch must be 'sparse_sign', 'srht' or "
                f"'gaussian', got {self.sketch!r}."
            )
        if self.sampling_factor < 1:
            raise ValueError(
                "SketchAndPrecondLSMR needs sampling_factor >= 1, got "
                f"{self.sampling_factor}."
            )
        if self.damp < 0:
            raise ValueError(f"SketchAndPrecondLSMR needs damp >= 0, got {self.damp}.")

    def _sample_sketch(self, m: int, n: int) -> AbstractSketch:
        d = min(math.ceil(self.sampling_factor * n), m)
        key = jax.random.PRNGKey(self.seed)
        if self.sketch == "sparse_sign":
            return SparseSignSketch.sample(key, d, m, nnz=self.nnz)
        if self.sketch == "srht":
            return SRHTSketch.sample(key, d, m)
        return GaussianSketch.sample(key, d, m)

    def _solve(
        self, operator: lx.AbstractLinearOperator, vector: Float[Array, " m"]
    ) -> tuple[Float[Array, " n"], lx.Solution]:
        """The solution and lineax's inner LSMR solution (with its stats)."""
        m, n = operator.out_size(), operator.in_size()
        if self.damp == 0.0 and m < n:
            raise ValueError(
                "SketchAndPrecondLSMR needs a tall or square operator when "
                f"damp == 0, got shape ({m}, {n})."
            )
        sketch = self._sample_sketch(m, n)
        Q_top, R = _sketch_qr(operator, sketch, self.damp)
        x0 = _warm_start(Q_top, R, sketch, vector)
        if self.damp == 0.0 and self.throw:
            x0 = _check_full_rank(R, x0)
        damp = self.damp

        def precondition(y: Float[Array, " n"]) -> Float[Array, " n"]:
            return jsl.solve_triangular(R, y, lower=False)

        def augmented(x: Float[Array, " n"]) -> Float[Array, " mn"]:
            Ax = operator.mv(x)
            if damp == 0.0:
                return Ax
            return jnp.concatenate([Ax, damp * x])

        dtype = operator_dtype(operator, vector)
        preconditioned = lx.FunctionLinearOperator(
            lambda y: augmented(precondition(y)),
            jax.ShapeDtypeStruct((n,), R.dtype),
        )
        residual = (
            vector if damp == 0.0 else jnp.concatenate([vector, jnp.zeros_like(x0)])
        ) - augmented(x0)
        atol = resolve_tolerance(self.atol, dtype, 1e-6)
        btol = resolve_tolerance(self.btol, dtype, 1e-6)
        solution = lx.linear_solve(
            preconditioned,
            residual,
            lx.LSMR(rtol=btol, atol=atol, max_steps=self.max_steps),
            throw=self.throw,
        )
        return x0 + precondition(solution.value), solution

    def solve(
        self,
        operator: lx.AbstractLinearOperator,
        vector: Float[Array, " m"],
    ) -> Float[Array, " n"]:
        r"""Least-squares solution $\arg\min_x \|Ax - b\|^2 + \delta^2\|x\|^2$.

        Args:
            operator: Operator $A$, shape ``(m, n)``; tall ($m \ge n$) unless
                ``damp > 0``. Matrix-free operators are sketched with $d$
                transpose-matvecs.
            vector: Right-hand side $b$, shape ``(m,)``.

        Returns:
            The solution $x$, shape ``(n,)``.

        Raises:
            ValueError: If ``damp == 0`` and $m < n$.
            equinox.EquinoxRuntimeError: With ``throw=True``, if ``damp == 0``
                and $A$ has numerically dependent columns (the solution is
                not unique; pass ``damp > 0``), or LSMR does not converge.
        """
        return self._solve(operator, vector)[0]

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        r"""Stochastic log-determinant, delegated to `LSMRSolver`'s SLQ path.

        Present for the `AbstractSolverStrategy` protocol only: the sketch
        plays no part. For the log-determinant of a square system, use a
        square-system strategy (`SLQLogdet`, `CGSolver`, `DenseSolver`).

        Args:
            operator: A square symmetric PSD linear operator.
            key: PRNG key for the probes; ``None`` means
                ``PRNGKey(seed)``.

        Returns:
            Scalar estimate of $\log|\det A|$.
        """
        return LSMRSolver(seed=self.seed).logdet(operator, key=key)

`solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' m']) -> Float[Array, ' n']` ¶

Least-squares solution \(\arg\min_x \|Ax - b\|^2 + \delta^2\|x\|^2\).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Operator \(A\), shape `(m, n)`; tall (\(m \ge n\)) unless `damp > 0`. Matrix-free operators are sketched with \(d\) transpose-matvecs.	required
`vector`	`Float[Array, ' m']`	Right-hand side \(b\), shape `(m,)`.	required

Returns:

Type	Description
`Float[Array, ' n']`	The solution \(x\), shape `(n,)`.

Raises:

Type	Description
`ValueError`	If `damp == 0` and \(m < n\).
`EquinoxRuntimeError`	With `throw=True`, if `damp == 0` and \(A\) has numerically dependent columns (the solution is not unique; pass `damp > 0`), or LSMR does not converge.

Source code in src/gaussx/_strategies/_sketch_precond.py

def solve(
    self,
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " m"],
) -> Float[Array, " n"]:
    r"""Least-squares solution $\arg\min_x \|Ax - b\|^2 + \delta^2\|x\|^2$.

    Args:
        operator: Operator $A$, shape ``(m, n)``; tall ($m \ge n$) unless
            ``damp > 0``. Matrix-free operators are sketched with $d$
            transpose-matvecs.
        vector: Right-hand side $b$, shape ``(m,)``.

    Returns:
        The solution $x$, shape ``(n,)``.

    Raises:
        ValueError: If ``damp == 0`` and $m < n$.
        equinox.EquinoxRuntimeError: With ``throw=True``, if ``damp == 0``
            and $A$ has numerically dependent columns (the solution is
            not unique; pass ``damp > 0``), or LSMR does not converge.
    """
    return self._solve(operator, vector)[0]

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log-determinant, delegated to LSMRSolver's SLQ path.

Present for the AbstractSolverStrategy protocol only: the sketch plays no part. For the log-determinant of a square system, use a square-system strategy (SLQLogdet, CGSolver, DenseSolver).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A square symmetric PSD linear operator.	required
`key`	`Array \| None`	PRNG key for the probes; `None` means `PRNGKey(seed)`.	`None`

Returns:

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

Source code in src/gaussx/_strategies/_sketch_precond.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    r"""Stochastic log-determinant, delegated to `LSMRSolver`'s SLQ path.

    Present for the `AbstractSolverStrategy` protocol only: the sketch
    plays no part. For the log-determinant of a square system, use a
    square-system strategy (`SLQLogdet`, `CGSolver`, `DenseSolver`).

    Args:
        operator: A square symmetric PSD linear operator.
        key: PRNG key for the probes; ``None`` means
            ``PRNGKey(seed)``.

    Returns:
        Scalar estimate of $\log|\det A|$.
    """
    return LSMRSolver(seed=self.seed).logdet(operator, key=key)

`sketch_and_solve(operator: lx.AbstractLinearOperator, vector: Float[Array, ' m'], *, sketch: AbstractSketch, damp: float = 0.0) -> Float[Array, ' n']` ¶

Approximate least squares from a sketch.

Solves the sketched problem \(\min_x \|S(Ax - b)\|^2 + \delta^2\|x\|^2\) exactly, through the QR factorisation \([SA;\ \delta I] = QR\) of the \(d \times n\) sketch (stacked on \(\delta I\) when damp \(= \delta > 0\)):

\[ \hat x = R^{-1} Q_{1:d}^\top S b . \]

If \(S\) is an \(\varepsilon\)-subspace embedding for \(\operatorname{range}([A\ b])\), the residual is within a constant factor of the optimum \(x^\star\) (Woodruff, 2014, §2):

\[ \|A\hat x - b\| \le \tfrac{1}{1-\varepsilon}\|S(A\hat x - b)\| \le \tfrac{1}{1-\varepsilon}\|S(Ax^\star - b)\| \le \frac{1+\varepsilon}{1-\varepsilon}\,\|Ax^\star - b\| . \]

The solution itself is only \(O(\varepsilon)\)-accurate, so this is a function, not a solver strategy; SketchAndPrecondLSMR uses it as its warm start and then solves to tolerance.

SA = S A                        # d × n, matrix-free via sketch_operator
Q, R = qr([SA; δ I])            # reduced QR, O(d n²)
x̂ = R⁻¹ Q[:d]ᵀ (S b)

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	Tall operator \(A\), shape `(m, n)`.	required
`vector`	`Float[Array, ' m']`	Right-hand side \(b\), shape `(m,)`.	required
`sketch`	`AbstractSketch`	A sampled sketch with `in_size == m` and, for a unique solution when `damp == 0`, `out_size >= n`.	required
`damp`	`float`	Ridge parameter \(\delta \ge 0\).	`0.0`

Returns:

Type	Description
`Float[Array, ' n']`	The sketch-and-solve estimate \(\hat x\), shape `(n,)`.

Raises:

Type	Description
`ValueError`	If `damp < 0` or the sketch's `in_size` is not \(m\).
`EquinoxRuntimeError`	If `damp == 0` and \(SA\) is numerically rank-deficient (\(A\) has dependent columns).

References

Woodruff, D. P. (2014). Sketching as a tool for numerical linear algebra. Foundations and Trends in Theoretical Computer Science, 10(1-2), 1-157.

Examples:

>>> import jax.numpy as jnp, jax.random as jr, lineax as lx
>>> import gaussx as gx
>>> A = jr.normal(jr.key(0), (2000, 10))
>>> b = A @ jnp.ones(10) + 0.1 * jr.normal(jr.key(1), (2000,))
>>> S = gx.GaussianSketch.sample(jr.key(2), d=200, m=2000)
>>> x = gx.sketch_and_solve(lx.MatrixLinearOperator(A), b, sketch=S)
>>> bool(jnp.max(jnp.abs(x - 1.0)) < 0.05)
True

Source code in src/gaussx/_strategies/_sketch_precond.py

def sketch_and_solve(
    operator: lx.AbstractLinearOperator,
    vector: Float[Array, " m"],
    *,
    sketch: AbstractSketch,
    damp: float = 0.0,
) -> Float[Array, " n"]:
    r"""Approximate least squares from a sketch.

    Solves the sketched problem $\min_x \|S(Ax - b)\|^2 + \delta^2\|x\|^2$
    exactly, through the QR factorisation
    $[SA;\ \delta I] = QR$ of the $d \times n$ sketch (stacked on
    $\delta I$ when ``damp`` $= \delta > 0$):

    $$
    \hat x = R^{-1} Q_{1:d}^\top S b .
    $$

    If $S$ is an $\varepsilon$-subspace embedding for
    $\operatorname{range}([A\ b])$, the residual is within a constant factor
    of the optimum $x^\star$ (Woodruff, 2014, §2):

    $$
    \|A\hat x - b\| \le \tfrac{1}{1-\varepsilon}\|S(A\hat x - b)\|
    \le \tfrac{1}{1-\varepsilon}\|S(Ax^\star - b)\|
    \le \frac{1+\varepsilon}{1-\varepsilon}\,\|Ax^\star - b\| .
    $$

    The solution itself is only $O(\varepsilon)$-accurate, so this is a
    function, not a solver strategy; `SketchAndPrecondLSMR` uses it as its
    warm start and then solves to tolerance.

    ```text
    SA = S A                        # d × n, matrix-free via sketch_operator
    Q, R = qr([SA; δ I])            # reduced QR, O(d n²)
    x̂ = R⁻¹ Q[:d]ᵀ (S b)
    ```

    Args:
        operator: Tall operator $A$, shape ``(m, n)``.
        vector: Right-hand side $b$, shape ``(m,)``.
        sketch: A sampled sketch with ``in_size == m`` and, for a unique
            solution when ``damp == 0``, ``out_size >= n``.
        damp: Ridge parameter $\delta \ge 0$.

    Returns:
        The sketch-and-solve estimate $\hat x$, shape ``(n,)``.

    Raises:
        ValueError: If ``damp < 0`` or the sketch's ``in_size`` is not $m$.
        equinox.EquinoxRuntimeError: If ``damp == 0`` and $SA$ is
            numerically rank-deficient ($A$ has dependent columns).

    References:
        Woodruff, D. P. (2014). Sketching as a tool for numerical linear
        algebra. *Foundations and Trends in Theoretical Computer Science*,
        10(1-2), 1-157.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr, lineax as lx
        >>> import gaussx as gx
        >>> A = jr.normal(jr.key(0), (2000, 10))
        >>> b = A @ jnp.ones(10) + 0.1 * jr.normal(jr.key(1), (2000,))
        >>> S = gx.GaussianSketch.sample(jr.key(2), d=200, m=2000)
        >>> x = gx.sketch_and_solve(lx.MatrixLinearOperator(A), b, sketch=S)
        >>> bool(jnp.max(jnp.abs(x - 1.0)) < 0.05)
        True
    """
    if damp < 0:
        raise ValueError(f"sketch_and_solve needs damp >= 0, got {damp}.")
    Q_top, R = _sketch_qr(operator, sketch, damp)
    x = _warm_start(Q_top, R, sketch, vector)
    return x if damp != 0.0 else _check_full_rank(R, x)

Logdet strategies¶

Dense eigendecomposition for exactness; stochastic Lanczos quadrature (SLQ) for \(O(n^2 \cdot \text{rank})\) estimates on large PSD (or symmetric-indefinite) operators.

NystromLogdet (Wenger et al., 2022) is the variance-reduced SLQ for a covariance-form \(K + \sigma^2 I\). The log-determinant of the Nyström preconditioner \(P = \hat K + \sigma^2 I\) is exact, and SLQ only estimates \(\log|P^{-1/2}(K + \sigma^2 I)P^{-1/2}|\), whose operator is close to the identity, so the probe variance collapses. Use plain SLQLogdet for precision-form GMRF systems.

Structured linear algebra and Gaussian primitives for JAX.

`DenseLogdet` ¶

Dense log-determinant via gaussx structural dispatch.

Delegates to gaussx.logdet which automatically selects the best algorithm based on operator structure (Diagonal, BlockDiag, Kronecker, LowRankUpdate, or dense fallback).

Source code in src/gaussx/_strategies/_slq_logdet.py

class DenseLogdet(AbstractLogdetStrategy):
    """Dense log-determinant via gaussx structural dispatch.

    Delegates to `gaussx.logdet` which automatically selects
    the best algorithm based on operator structure (Diagonal,
    BlockDiag, Kronecker, LowRankUpdate, or dense fallback).
    """

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Compute ``log |det(A)|`` via structural dispatch.

        Args:
            operator: A linear operator.
            key: Ignored (deterministic).

        Returns:
            Scalar ``log |det(A)|``.
        """

        return _logdet(operator)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

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

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A linear operator.	required
`key`	`Array \| None`	Ignored (deterministic).	`None`

Returns:

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

Source code in src/gaussx/_strategies/_slq_logdet.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Compute ``log |det(A)|`` via structural dispatch.

    Args:
        operator: A linear operator.
        key: Ignored (deterministic).

    Returns:
        Scalar ``log |det(A)|``.
    """

    return _logdet(operator)

`SLQLogdet` ¶

Stochastic log-determinant via Lanczos quadrature (SLQ).

Estimates log det(A) for PSD operators using stochastic trace estimation: logdet(A) = tr(log(A)). Uses matfree's Lanczos decomposition with sign-flip ("Rademacher") probe vectors by default.

With no key, logdet uses PRNGKey(seed), so every call sees the same probes: common random numbers, which suit stochastic-gradient training but give no variance reduction when estimates are averaged and a fixed pseudo-likelihood under MCMC. Pass a fresh key (or wrap a strategy in gaussx.KeyedSolver) to decorrelate calls.

Attributes:

Name	Type	Description
`num_probes`	`int`	Number of probe vectors for Hutchinson estimator.
`lanczos_order`	`int`	Order of the Lanczos decomposition.
`seed`	`int`	Seed for probe vector generation (used when no `key` is passed to `logdet`).
`sampler`	`SamplerName`	Probe distribution (`"signs"`, `"normal"`, `"sphere"`).

Source code in src/gaussx/_strategies/_slq_logdet.py

class SLQLogdet(AbstractLogdetStrategy):
    """Stochastic log-determinant via Lanczos quadrature (SLQ).

    Estimates ``log det(A)`` for PSD operators using stochastic
    trace estimation: ``logdet(A) = tr(log(A))``.  Uses matfree's
    Lanczos decomposition with sign-flip ("Rademacher") probe vectors
    by default.

    With no ``key``, `logdet` uses ``PRNGKey(seed)``, so every call sees the
    same probes: common random numbers, which suit stochastic-gradient
    training but give no variance reduction when estimates are averaged and
    a fixed pseudo-likelihood under MCMC. Pass a fresh ``key`` (or wrap a
    strategy in `gaussx.KeyedSolver`) to decorrelate calls.

    Attributes:
        num_probes: Number of probe vectors for Hutchinson estimator.
        lanczos_order: Order of the Lanczos decomposition.
        seed: Seed for probe vector generation (used when no
            ``key`` is passed to `logdet`).
        sampler: Probe distribution (``"signs"``, ``"normal"``,
            ``"sphere"``).
    """

    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    sampler: SamplerName = eqx.field(static=True, default="signs")

    def _integrand(self, n: int):
        return _logdet_integrand(min(self.lanczos_order, n))

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        """Stochastic ``log det(A)`` via Lanczos quadrature.

        Args:
            operator: A PSD linear operator.
            key: PRNG key for probe vector sampling.  If ``None``,
                uses ``jax.random.PRNGKey(self.seed)``.

        Returns:
            Scalar estimate of ``log det(A)``.
        """
        _check_symmetric(operator, "SLQLogdet")
        if key is None:
            key = jax.random.PRNGKey(self.seed)

        n = operator.in_size()
        point, _ = _slq_estimators(
            self._integrand(n),
            n,
            self.num_probes,
            self.sampler,
            operator.in_structure().dtype,
        )
        return point(operator.mv, key)

    def logdet_and_error(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> tuple[Float[Array, ""], Float[Array, ""]]:
        """Stochastic ``log det(A)`` with its standard error.

        Args:
            operator: A PSD linear operator.
            key: PRNG key for probe vector sampling.  If ``None``,
                uses ``jax.random.PRNGKey(self.seed)``.

        Returns:
            Tuple ``(estimate, standard_error)`` where the standard
            error is the standard error of the mean across probes.
        """
        _check_symmetric(operator, "SLQLogdet")
        if key is None:
            key = jax.random.PRNGKey(self.seed)

        n = operator.in_size()
        _, with_sem = _slq_estimators(
            self._integrand(n),
            n,
            self.num_probes,
            self.sampler,
            operator.in_structure().dtype,
        )
        return with_sem(operator.mv, key)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log det(A) via Lanczos quadrature.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If `None`, uses `jax.random.PRNGKey(self.seed)`.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar estimate of `log det(A)`.

Source code in src/gaussx/_strategies/_slq_logdet.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Stochastic ``log det(A)`` via Lanczos quadrature.

    Args:
        operator: A PSD linear operator.
        key: PRNG key for probe vector sampling.  If ``None``,
            uses ``jax.random.PRNGKey(self.seed)``.

    Returns:
        Scalar estimate of ``log det(A)``.
    """
    _check_symmetric(operator, "SLQLogdet")
    if key is None:
        key = jax.random.PRNGKey(self.seed)

    n = operator.in_size()
    point, _ = _slq_estimators(
        self._integrand(n),
        n,
        self.num_probes,
        self.sampler,
        operator.in_structure().dtype,
    )
    return point(operator.mv, key)

`logdet_and_error(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> tuple[Float[Array, ''], Float[Array, '']]` ¶

Stochastic log det(A) with its standard error.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A PSD linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If `None`, uses `jax.random.PRNGKey(self.seed)`.	`None`

Returns:

Type	Description
`Float[Array, '']`	Tuple `(estimate, standard_error)` where the standard
`Float[Array, '']`	error is the standard error of the mean across probes.

Source code in src/gaussx/_strategies/_slq_logdet.py

def logdet_and_error(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    """Stochastic ``log det(A)`` with its standard error.

    Args:
        operator: A PSD linear operator.
        key: PRNG key for probe vector sampling.  If ``None``,
            uses ``jax.random.PRNGKey(self.seed)``.

    Returns:
        Tuple ``(estimate, standard_error)`` where the standard
        error is the standard error of the mean across probes.
    """
    _check_symmetric(operator, "SLQLogdet")
    if key is None:
        key = jax.random.PRNGKey(self.seed)

    n = operator.in_size()
    _, with_sem = _slq_estimators(
        self._integrand(n),
        n,
        self.num_probes,
        self.sampler,
        operator.in_structure().dtype,
    )
    return with_sem(operator.mv, key)

`IndefiniteSLQLogdet` ¶

Stochastic log|det(A)| for symmetric (possibly indefinite) operators.

Like SLQLogdet but uses log(|lambda|) as the matrix function, so it works on indefinite and negative-definite matrices. Supports a diagonal shift (A + shift * I).

Attributes:

Name	Type	Description
`num_probes`	`int`	Number of probe vectors for Hutchinson estimator.
`lanczos_order`	`int`	Order of the Lanczos decomposition.
`shift`	`float`	Diagonal shift applied before computing the logdet.
`seed`	`int`	Seed for probe vector generation.
`sampler`	`SamplerName`	Probe distribution (`"signs"`, `"normal"`, `"sphere"`).

Source code in src/gaussx/_strategies/_slq_logdet.py

class IndefiniteSLQLogdet(AbstractLogdetStrategy):
    """Stochastic ``log|det(A)|`` for symmetric (possibly indefinite) operators.

    Like `SLQLogdet` but uses ``log(|lambda|)`` as the matrix
    function, so it works on indefinite and negative-definite matrices.
    Supports a diagonal shift ``(A + shift * I)``.

    Attributes:
        num_probes: Number of probe vectors for Hutchinson estimator.
        lanczos_order: Order of the Lanczos decomposition.
        shift: Diagonal shift applied before computing the logdet.
        seed: Seed for probe vector generation.
        sampler: Probe distribution (``"signs"``, ``"normal"``,
            ``"sphere"``).
    """

    num_probes: int = eqx.field(static=True, default=20)
    lanczos_order: int = eqx.field(static=True, default=30)
    shift: float = eqx.field(static=True, default=0.0)
    seed: int = eqx.field(static=True, default=0)
    sampler: SamplerName = eqx.field(static=True, default="signs")

    def _shifted_matvec(self, operator: lx.AbstractLinearOperator):
        shift = self.shift

        def matvec(v):
            return operator.mv(v) + shift * v

        return matvec

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        r"""Stochastic ``log|det(A + shift I)|`` via Lanczos quadrature.

        Args:
            operator: A symmetric linear operator.
            key: PRNG key for probe vector sampling.  If ``None``,
                uses ``jax.random.PRNGKey(self.seed)``.

        Returns:
            Scalar estimate of ``log|det(A + shift I)|``.
        """
        _check_symmetric(operator, "IndefiniteSLQLogdet")
        if key is None:
            key = jax.random.PRNGKey(self.seed)

        n = operator.in_size()
        order = min(self.lanczos_order, n)
        point, _ = _slq_estimators(
            _logabsdet_integrand(order),
            n,
            self.num_probes,
            self.sampler,
            operator.in_structure().dtype,
        )
        return point(self._shifted_matvec(operator), key)

    def logdet_and_error(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> tuple[Float[Array, ""], Float[Array, ""]]:
        """Stochastic ``log|det(A + shift I)|`` with its standard error.

        Args:
            operator: A symmetric linear operator.
            key: PRNG key for probe vector sampling.  If ``None``,
                uses ``jax.random.PRNGKey(self.seed)``.

        Returns:
            Tuple ``(estimate, standard_error)``.
        """
        _check_symmetric(operator, "IndefiniteSLQLogdet")
        if key is None:
            key = jax.random.PRNGKey(self.seed)

        n = operator.in_size()
        order = min(self.lanczos_order, n)
        _, with_sem = _slq_estimators(
            _logabsdet_integrand(order),
            n,
            self.num_probes,
            self.sampler,
            operator.in_structure().dtype,
        )
        return with_sem(self._shifted_matvec(operator), key)

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Stochastic log|det(A + shift I)| via Lanczos quadrature.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If `None`, uses `jax.random.PRNGKey(self.seed)`.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar estimate of `log\|det(A + shift I)\|`.

Source code in src/gaussx/_strategies/_slq_logdet.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    r"""Stochastic ``log|det(A + shift I)|`` via Lanczos quadrature.

    Args:
        operator: A symmetric linear operator.
        key: PRNG key for probe vector sampling.  If ``None``,
            uses ``jax.random.PRNGKey(self.seed)``.

    Returns:
        Scalar estimate of ``log|det(A + shift I)|``.
    """
    _check_symmetric(operator, "IndefiniteSLQLogdet")
    if key is None:
        key = jax.random.PRNGKey(self.seed)

    n = operator.in_size()
    order = min(self.lanczos_order, n)
    point, _ = _slq_estimators(
        _logabsdet_integrand(order),
        n,
        self.num_probes,
        self.sampler,
        operator.in_structure().dtype,
    )
    return point(self._shifted_matvec(operator), key)

`logdet_and_error(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> tuple[Float[Array, ''], Float[Array, '']]` ¶

Stochastic log|det(A + shift I)| with its standard error.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	A symmetric linear operator.	required
`key`	`Array \| None`	PRNG key for probe vector sampling. If `None`, uses `jax.random.PRNGKey(self.seed)`.	`None`

Returns:

Type	Description
`tuple[Float[Array, ''], Float[Array, '']]`	Tuple `(estimate, standard_error)`.

Source code in src/gaussx/_strategies/_slq_logdet.py

def logdet_and_error(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    """Stochastic ``log|det(A + shift I)|`` with its standard error.

    Args:
        operator: A symmetric linear operator.
        key: PRNG key for probe vector sampling.  If ``None``,
            uses ``jax.random.PRNGKey(self.seed)``.

    Returns:
        Tuple ``(estimate, standard_error)``.
    """
    _check_symmetric(operator, "IndefiniteSLQLogdet")
    if key is None:
        key = jax.random.PRNGKey(self.seed)

    n = operator.in_size()
    order = min(self.lanczos_order, n)
    _, with_sem = _slq_estimators(
        _logabsdet_integrand(order),
        n,
        self.num_probes,
        self.sampler,
        operator.in_structure().dtype,
    )
    return with_sem(self._shifted_matvec(operator), key)

`NystromLogdet` ¶

Nyström-preconditioned log-determinant of \(A + \mu I\) (Wenger et al., 2022).

For a covariance-form system \(A + \mu I\) (\(A\) PSD, e.g. \(K + \sigma^2 I\)) and the rank-\(k\) randomized Nyström approximation \(\hat A = U\hat\Lambda U^\top\) of \(A\) (gaussx.randomized_nystrom), the preconditioner \(P = \hat A + \mu I\) splits the log-determinant as

\[ \log|A + \mu I| = \log|P| + \log\big|P^{-1/2}(A + \mu I)P^{-1/2}\big|, \qquad \log|P| = \sum_{i=1}^{k}\log(\hat\lambda_i + \mu) + (n - k)\log\mu . \]

The first term is exact. The second is SLQ (gaussx.SLQLogdet) on the preconditioned operator, whose eigenvalues cluster at one once \(k\) exceeds the effective dimension \(d_\text{eff}(\mu) = \operatorname{tr}(A (A + \mu I)^{-1})\); SLQ's variance scales with \(\|\log M\|_F^2\), so it collapses as \(M \to I\), and its Lanczos quadrature converges in few steps. With

\[ P^{-1/2} = U\big((\hat\Lambda + \mu I)^{-1/2} - \mu^{-1/2} I\big)U^\top + \mu^{-1/2} I, \]

every matvec with \(M\) costs one matvec with \(A + \mu I\) and \(O(nk)\).

key₁, key₂ = split(key)
c = 10 n eps μ                               # above the rounding of A + μI − μI
U, Λ̂ = randomized_nystrom(A + μI − (μ − c)I, rank; key₁);  Λ̂ ← max(Λ̂ − c, 0)
log|P| = Σ log(λ̂ᵢ + μ) + (n − k) log μ
M = P^{-1/2} (A + μI) P^{-1/2}
return log|P| + SLQ(M; num_probes, lanczos_order; key₂)

Covariance form only: \(\hat A\) captures the top of \(A\)'s spectrum, as for gaussx.NystromPreconditioner. For a precision-form GMRF system \(Q + A^\top W A\) the hard directions are the smallest eigenvalues; use SLQLogdet (or an exact sparse factorisation) there.

operator is the full system \(A + \mu I\), as for every other logdet strategy; the strategy sketches its PSD part as the matrix-free operator − shift·I (plus \(c = 10\,n\,\epsilon\,\mu\), removed again from \(\hat\Lambda\), so that the rounding of the subtraction cannot make it indefinite), so shift must be the \(\mu\) that operator contains (#345). shift is a PyTree leaf, so it can be a traced, learned noise variance.

Attributes:

Name	Type	Description
`shift`	`float \| Float[Array, '']`	\(\mu > 0\), e.g. the noise variance \(\sigma^2\).
`rank`	`int`	Nyström rank \(k\) (`rank` matvecs), clamped to \(n\). Aim for \(k \gtrsim d_\text{eff}(\mu)\).
`oversample`	`int`	Extra sketch columns (see `gaussx.randomized_nystrom`).
`num_probes`	`int`	SLQ probe vectors for the remainder.
`lanczos_order`	`int`	Lanczos steps per probe.
`seed`	`int`	`logdet` without a `key` uses `PRNGKey(seed)` for both the sketch and the probes (common random numbers, as `SLQLogdet`).
`sampler`	`SamplerName`	SLQ probe distribution (`"signs"`, `"normal"`, `"sphere"`).

References

Wenger, J., Pleiss, G., Hennig, P., Cunningham, J. P. & Gardner, J. R. (2022). Preconditioning for scalable Gaussian process hyperparameter optimization. ICML, PMLR 162, 23751-23780.

Frangella, Z., Tropp, J. A. & Udell, M. (2023). Randomized Nyström preconditioning. SIAM J. Matrix Anal. Appl., 44(2), 718-752.

Examples:

>>> import einx, jax.numpy as jnp, lineax as lx
>>> import gaussx as gx
>>> x = jnp.linspace(0.0, 10.0, 200)
>>> K = jnp.exp(-0.5 * einx.subtract("i, j -> i j", x, x) ** 2)
>>> A = lx.MatrixLinearOperator(
...     K + 0.1 * jnp.eye(200), lx.positive_semidefinite_tag
... )
>>> est = gx.NystromLogdet(shift=0.1, rank=40).logdet(A)
>>> exact = jnp.linalg.slogdet(A.as_matrix())[1]
>>> bool(jnp.abs(est - exact) < 1e-2 * jnp.abs(exact))
True

Source code in src/gaussx/_strategies/_nystrom_logdet.py

class NystromLogdet(AbstractLogdetStrategy):
    r"""Nyström-preconditioned log-determinant of $A + \mu I$ (Wenger et al., 2022).

    For a covariance-form system $A + \mu I$ ($A$ PSD, e.g. $K + \sigma^2 I$)
    and the rank-$k$ randomized Nyström approximation $\hat A = U\hat\Lambda
    U^\top$ of $A$ (`gaussx.randomized_nystrom`), the preconditioner
    $P = \hat A + \mu I$ splits the log-determinant as

    $$
    \log|A + \mu I| = \log|P| + \log\big|P^{-1/2}(A + \mu I)P^{-1/2}\big|,
    \qquad
    \log|P| = \sum_{i=1}^{k}\log(\hat\lambda_i + \mu) + (n - k)\log\mu .
    $$

    The first term is exact. The second is SLQ (`gaussx.SLQLogdet`) on the
    preconditioned operator, whose eigenvalues cluster at one once $k$
    exceeds the effective dimension $d_\text{eff}(\mu) = \operatorname{tr}(A
    (A + \mu I)^{-1})$; SLQ's variance scales with $\|\log M\|_F^2$, so it
    collapses as $M \to I$, and its Lanczos quadrature converges in few
    steps. With

    $$
    P^{-1/2} = U\big((\hat\Lambda + \mu I)^{-1/2} - \mu^{-1/2} I\big)U^\top
    + \mu^{-1/2} I,
    $$

    every matvec with $M$ costs one matvec with $A + \mu I$ and $O(nk)$.

    ```text
    key₁, key₂ = split(key)
    c = 10 n eps μ                               # above the rounding of A + μI − μI
    U, Λ̂ = randomized_nystrom(A + μI − (μ − c)I, rank; key₁);  Λ̂ ← max(Λ̂ − c, 0)
    log|P| = Σ log(λ̂ᵢ + μ) + (n − k) log μ
    M = P^{-1/2} (A + μI) P^{-1/2}
    return log|P| + SLQ(M; num_probes, lanczos_order; key₂)
    ```

    **Covariance form only**: $\hat A$ captures the *top* of $A$'s
    spectrum, as for `gaussx.NystromPreconditioner`. For a precision-form
    GMRF system $Q + A^\top W A$ the hard directions are the smallest
    eigenvalues; use `SLQLogdet` (or an exact sparse factorisation) there.

    ``operator`` is the full system $A + \mu I$, as for every other logdet
    strategy; the strategy sketches its PSD part as the matrix-free
    ``operator − shift·I`` (plus $c = 10\,n\,\epsilon\,\mu$, removed again
    from $\hat\Lambda$, so that the rounding of the subtraction cannot make
    it indefinite), so ``shift`` must be the $\mu$ that ``operator``
    contains (#345). ``shift`` is a PyTree leaf, so it can be
    a traced, learned noise variance.

    Attributes:
        shift: $\mu > 0$, e.g. the noise variance $\sigma^2$.
        rank: Nyström rank $k$ (``rank`` matvecs), clamped to $n$. Aim for
            $k \gtrsim d_\text{eff}(\mu)$.
        oversample: Extra sketch columns (see `gaussx.randomized_nystrom`).
        num_probes: SLQ probe vectors for the remainder.
        lanczos_order: Lanczos steps per probe.
        seed: ``logdet`` without a ``key`` uses ``PRNGKey(seed)`` for both
            the sketch and the probes (common random numbers, as
            `SLQLogdet`).
        sampler: SLQ probe distribution (``"signs"``, ``"normal"``,
            ``"sphere"``).

    References:
        Wenger, J., Pleiss, G., Hennig, P., Cunningham, J. P. & Gardner,
        J. R. (2022). Preconditioning for scalable Gaussian process
        hyperparameter optimization. *ICML*, PMLR 162, 23751-23780.

        Frangella, Z., Tropp, J. A. & Udell, M. (2023). Randomized Nyström
        preconditioning. *SIAM J. Matrix Anal. Appl.*, 44(2), 718-752.

    Examples:
        >>> import einx, jax.numpy as jnp, lineax as lx
        >>> import gaussx as gx
        >>> x = jnp.linspace(0.0, 10.0, 200)
        >>> K = jnp.exp(-0.5 * einx.subtract("i, j -> i j", x, x) ** 2)
        >>> A = lx.MatrixLinearOperator(
        ...     K + 0.1 * jnp.eye(200), lx.positive_semidefinite_tag
        ... )
        >>> est = gx.NystromLogdet(shift=0.1, rank=40).logdet(A)
        >>> exact = jnp.linalg.slogdet(A.as_matrix())[1]
        >>> bool(jnp.abs(est - exact) < 1e-2 * jnp.abs(exact))
        True
    """

    shift: float | Float[Array, ""]
    rank: int = eqx.field(static=True, default=50)
    oversample: int = eqx.field(static=True, default=0)
    num_probes: int = eqx.field(static=True, default=10)
    lanczos_order: int = eqx.field(static=True, default=30)
    seed: int = eqx.field(static=True, default=0)
    sampler: SamplerName = eqx.field(static=True, default="signs")

    def __check_init__(self) -> None:
        if isinstance(self.shift, (int, float)) and self.shift <= 0:
            raise ValueError(f"NystromLogdet needs shift > 0, got {self.shift}.")
        if self.rank < 1:
            raise ValueError(f"NystromLogdet needs rank >= 1, got {self.rank}.")

    def logdet(
        self,
        operator: lx.AbstractLinearOperator,
        *,
        key: jax.Array | None = None,
    ) -> Float[Array, ""]:
        r"""Estimate $\log|A + \mu I|$.

        Args:
            operator: The symmetric PSD system $A + \mu I$, ``shift`` $= \mu$
                included; may be matrix-free.
            key: PRNG key for the sketch and the probes. ``None`` means
                ``jax.random.PRNGKey(seed)``.

        Returns:
            Scalar estimate of $\log|A + \mu I|$.
        """
        _check_symmetric(operator, "NystromLogdet")
        if key is None:
            key = jax.random.PRNGKey(self.seed)
        sketch_key, probe_key = jax.random.split(key)
        n = operator.in_size()
        structure = operator.in_structure()
        mu = jnp.asarray(self.shift, dtype=structure.dtype)

        # `operator - mu I` carries rounding noise of ~eps * mu, indefinite
        # once A is small next to mu (e.g. a zero kernel amplitude), which
        # breaks the Nystrom Cholesky (gh-486 review). Sketch A + c I with c
        # well above that noise, then shift the eigenvalues back.
        guard = 10 * n * jnp.finfo(structure.dtype).eps * mu
        psd_part = lx.FunctionLinearOperator(
            lambda v: operator.mv(v) - (mu - guard) * v,
            structure,
            lx.positive_semidefinite_tag,
        )
        approx = randomized_nystrom(
            psd_part, min(self.rank, n), oversample=self.oversample, key=sketch_key
        )
        U, eigenvalues = approx.U, jnp.maximum(approx.d - guard, 0.0)
        k = eigenvalues.shape[0]
        logdet_preconditioner = jnp.sum(jnp.log(eigenvalues + mu)) + (n - k) * jnp.log(
            mu
        )
        inv_sqrt_mu = 1.0 / jnp.sqrt(mu)
        correction = 1.0 / jnp.sqrt(eigenvalues + mu) - inv_sqrt_mu

        def inv_sqrt_preconditioner(v: Float[Array, " n"]) -> Float[Array, " n"]:
            coefficients = correction * einsum(U, v, "n k, n -> k")
            return inv_sqrt_mu * v + einsum(U, coefficients, "n k, k -> n")

        preconditioned = lx.FunctionLinearOperator(
            lambda v: inv_sqrt_preconditioner(operator.mv(inv_sqrt_preconditioner(v))),
            structure,
            (lx.symmetric_tag, lx.positive_semidefinite_tag),
        )
        remainder = SLQLogdet(
            num_probes=self.num_probes,
            lanczos_order=self.lanczos_order,
            seed=self.seed,
            sampler=self.sampler,
        ).logdet(preconditioned, key=probe_key)
        return logdet_preconditioner + remainder

`logdet(operator: lx.AbstractLinearOperator, *, key: jax.Array | None = None) -> Float[Array, '']` ¶

Estimate \(\log|A + \mu I|\).

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	The symmetric PSD system \(A + \mu I\), `shift` \(= \mu\) included; may be matrix-free.	required
`key`	`Array \| None`	PRNG key for the sketch and the probes. `None` means `jax.random.PRNGKey(seed)`.	`None`

Returns:

Type	Description
`Float[Array, '']`	Scalar estimate of \(\log\|A + \mu I\|\).

Source code in src/gaussx/_strategies/_nystrom_logdet.py

def logdet(
    self,
    operator: lx.AbstractLinearOperator,
    *,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    r"""Estimate $\log|A + \mu I|$.

    Args:
        operator: The symmetric PSD system $A + \mu I$, ``shift`` $= \mu$
            included; may be matrix-free.
        key: PRNG key for the sketch and the probes. ``None`` means
            ``jax.random.PRNGKey(seed)``.

    Returns:
        Scalar estimate of $\log|A + \mu I|$.
    """
    _check_symmetric(operator, "NystromLogdet")
    if key is None:
        key = jax.random.PRNGKey(self.seed)
    sketch_key, probe_key = jax.random.split(key)
    n = operator.in_size()
    structure = operator.in_structure()
    mu = jnp.asarray(self.shift, dtype=structure.dtype)

    # `operator - mu I` carries rounding noise of ~eps * mu, indefinite
    # once A is small next to mu (e.g. a zero kernel amplitude), which
    # breaks the Nystrom Cholesky (gh-486 review). Sketch A + c I with c
    # well above that noise, then shift the eigenvalues back.
    guard = 10 * n * jnp.finfo(structure.dtype).eps * mu
    psd_part = lx.FunctionLinearOperator(
        lambda v: operator.mv(v) - (mu - guard) * v,
        structure,
        lx.positive_semidefinite_tag,
    )
    approx = randomized_nystrom(
        psd_part, min(self.rank, n), oversample=self.oversample, key=sketch_key
    )
    U, eigenvalues = approx.U, jnp.maximum(approx.d - guard, 0.0)
    k = eigenvalues.shape[0]
    logdet_preconditioner = jnp.sum(jnp.log(eigenvalues + mu)) + (n - k) * jnp.log(
        mu
    )
    inv_sqrt_mu = 1.0 / jnp.sqrt(mu)
    correction = 1.0 / jnp.sqrt(eigenvalues + mu) - inv_sqrt_mu

    def inv_sqrt_preconditioner(v: Float[Array, " n"]) -> Float[Array, " n"]:
        coefficients = correction * einsum(U, v, "n k, n -> k")
        return inv_sqrt_mu * v + einsum(U, coefficients, "n k, k -> n")

    preconditioned = lx.FunctionLinearOperator(
        lambda v: inv_sqrt_preconditioner(operator.mv(inv_sqrt_preconditioner(v))),
        structure,
        (lx.symmetric_tag, lx.positive_semidefinite_tag),
    )
    remainder = SLQLogdet(
        num_probes=self.num_probes,
        lanczos_order=self.lanczos_order,
        seed=self.seed,
        sampler=self.sampler,
    ).logdet(preconditioned, key=probe_key)
    return logdet_preconditioner + remainder

Preconditioners¶

Approximate inverses \(M^{-1} \approx A^{-1}\) that accelerate the iterative solvers above. Pass them via the preconditioner= argument of linear_solve, CGSolver, or PreconditionedCGSolver.

For a covariance-form system \(K + \sigma^2 I\), build NystromPreconditioner or PartialCholeskyPreconditioner once with from_operator(K, rank, shift=σ²): the operator is the PSD part \(K\) only and the noise is passed separately, so it is never counted twice. With a Nyström rank \(\ell \gtrsim 2\lceil 1.5\,d_{\text{eff}}(\sigma^2)\rceil + 1\), \(d_{\text{eff}}(\mu) = \operatorname{tr}\big(K(K + \mu I)^{-1}\big)\), CG needs a number of iterations that does not grow with \(n\) (Frangella, Tropp & Udell, 2023).

Structured linear algebra and Gaussian primitives for JAX.

`AbstractPreconditioner` ¶

Bases: Module

Protocol for preconditioners producing an approximate inverse.

Subclasses implement as_operator, returning a PSD lineax operator that applies M^{-1} (or None to disable preconditioning).

Source code in src/gaussx/_preconditioners/_base.py

class AbstractPreconditioner(eqx.Module):
    """Protocol for preconditioners producing an approximate inverse.

    Subclasses implement `as_operator`, returning a PSD lineax operator
    that applies ``M^{-1}`` (or ``None`` to disable preconditioning).
    """

    @abc.abstractmethod
    def as_operator(
        self,
        operator: lx.AbstractLinearOperator | None = None,
    ) -> lx.AbstractLinearOperator | None:
        """Return ``M^{-1}`` as a PSD lineax operator.

        Args:
            operator: The system operator ``A`` being solved. Used by
                data-dependent preconditioners (e.g. partial-Cholesky); ignored
                by static ones.

        Returns:
            A positive-semidefinite operator applying ``M^{-1}``, or ``None`` to
            indicate that no preconditioning should be applied.
        """
        ...

    def __call__(
        self,
        vector: Float[Array, " n"],
        operator: lx.AbstractLinearOperator | None = None,
    ) -> Float[Array, " n"]:
        """Apply ``M^{-1}`` to *vector* (identity when disabled)."""
        op = self.as_operator(operator)
        if op is None:
            return vector
        return op.mv(vector)

`as_operator(operator: lx.AbstractLinearOperator | None = None) -> lx.AbstractLinearOperator | None` `abstractmethod` ¶

Return M^{-1} as a PSD lineax operator.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator \| None`	The system operator `A` being solved. Used by data-dependent preconditioners (e.g. partial-Cholesky); ignored by static ones.	`None`

Returns:

Type	Description
`AbstractLinearOperator \| None`	A positive-semidefinite operator applying `M^{-1}`, or `None` to
`AbstractLinearOperator \| None`	indicate that no preconditioning should be applied.

Source code in src/gaussx/_preconditioners/_base.py

@abc.abstractmethod
def as_operator(
    self,
    operator: lx.AbstractLinearOperator | None = None,
) -> lx.AbstractLinearOperator | None:
    """Return ``M^{-1}`` as a PSD lineax operator.

    Args:
        operator: The system operator ``A`` being solved. Used by
            data-dependent preconditioners (e.g. partial-Cholesky); ignored
            by static ones.

    Returns:
        A positive-semidefinite operator applying ``M^{-1}``, or ``None`` to
        indicate that no preconditioning should be applied.
    """
    ...

`JacobiPreconditioner` ¶

Diagonal preconditioner M^{-1} = diag(1 / |diag(A)|).

The cheapest preconditioner: scales each coordinate by the reciprocal of the corresponding diagonal entry of A. Effective when A is diagonally dominant.

Cost. Without an explicit diagonal it is read from the system operator at every solve. A stored matrix and the structured operators (gaussx.diag: BlockDiag, Kronecker, LowRankUpdate, ..., also inside sums and scalings) give it exactly for free. A matrix-free part, such as a bare lineax.FunctionLinearOperator, is never materialised, but its diagonal costs one matvec per row (n matvecs, O(n) memory), since Jacobi needs it exactly (gh-361). Pass diagonal= when it is known in closed form, e.g. a kernel's variance.

A negative diagonal entry is inverted in absolute value, so the operator really is positive semi-definite, as CG requires (the usual choice for indefinite Jacobi). A zero entry is left unscaled (0), with a zero gradient (gh-402).

Attributes:

Name	Type	Description
`diagonal`	`Float[Array, ' n'] \| None`	The diagonal of `A`. When `None`, it is extracted from the operator passed to `as_operator` (see Cost).

Source code in src/gaussx/_preconditioners/_jacobi.py

class JacobiPreconditioner(AbstractPreconditioner):
    """Diagonal preconditioner ``M^{-1} = diag(1 / |diag(A)|)``.

    The cheapest preconditioner: scales each coordinate by the reciprocal of
    the corresponding diagonal entry of ``A``. Effective when ``A`` is
    diagonally dominant.

    **Cost.** Without an explicit ``diagonal`` it is read from the system
    operator at every solve. A stored matrix and the structured operators
    (`gaussx.diag`: `BlockDiag`, `Kronecker`, `LowRankUpdate`, ..., also
    inside sums and scalings) give it exactly for free. A matrix-free part,
    such as a bare `lineax.FunctionLinearOperator`, is never materialised,
    but its diagonal costs one matvec per row (``n`` matvecs, O(n) memory),
    since Jacobi needs it exactly (gh-361). Pass ``diagonal=`` when it is
    known in closed form, e.g. a kernel's variance.

    A negative diagonal entry is inverted in absolute value, so the
    operator really is positive semi-definite, as CG requires (the usual
    choice for indefinite Jacobi). A zero entry is left unscaled (``0``),
    with a zero gradient (gh-402).

    Attributes:
        diagonal: The diagonal of ``A``. When ``None``, it is extracted from the
            operator passed to `as_operator` (see **Cost**).
    """

    diagonal: Float[Array, " n"] | None = None

    def as_operator(
        self,
        operator: lx.AbstractLinearOperator | None = None,
    ) -> lx.AbstractLinearOperator:
        """Return ``diag(1 / d)`` as a PSD operator."""
        d = self.diagonal
        if d is None:
            if operator is None:
                raise ValueError(
                    "JacobiPreconditioner needs either an explicit `diagonal` "
                    "or an operator to extract one from."
                )

            d = matrix_free_diag(operator, estimate=False)
        # Double where: the untaken 1/0 branch would make the gradient NaN.
        nonzero = d != 0.0
        safe = jnp.where(nonzero, d, 1.0)
        inv = jnp.where(nonzero, 1.0 / jnp.abs(safe), 0.0)
        return lx.TaggedLinearOperator(
            lx.DiagonalLinearOperator(inv), lx.positive_semidefinite_tag
        )

`as_operator(operator: lx.AbstractLinearOperator | None = None) -> lx.AbstractLinearOperator` ¶

Return diag(1 / d) as a PSD operator.

Source code in src/gaussx/_preconditioners/_jacobi.py

def as_operator(
    self,
    operator: lx.AbstractLinearOperator | None = None,
) -> lx.AbstractLinearOperator:
    """Return ``diag(1 / d)`` as a PSD operator."""
    d = self.diagonal
    if d is None:
        if operator is None:
            raise ValueError(
                "JacobiPreconditioner needs either an explicit `diagonal` "
                "or an operator to extract one from."
            )

        d = matrix_free_diag(operator, estimate=False)
    # Double where: the untaken 1/0 branch would make the gradient NaN.
    nonzero = d != 0.0
    safe = jnp.where(nonzero, d, 1.0)
    inv = jnp.where(nonzero, 1.0 / jnp.abs(safe), 0.0)
    return lx.TaggedLinearOperator(
        lx.DiagonalLinearOperator(inv), lx.positive_semidefinite_tag
    )

`OperatorPreconditioner` ¶

Use an externally supplied approximate inverse as a preconditioner.

Attributes:

Name	Type	Description
`approx_inverse`	`AbstractLinearOperator \| Callable`	The approximate inverse `M^{-1}`, either as a lineax operator or as a callable `v -> M^{-1} v`.
`in_structure`	`object`	Input structure for the callable form. When `None` it is taken from the system operator passed to `as_operator`.

Source code in src/gaussx/_preconditioners/_operator.py

class OperatorPreconditioner(AbstractPreconditioner):
    """Use an externally supplied approximate inverse as a preconditioner.

    Attributes:
        approx_inverse: The approximate inverse ``M^{-1}``, either as a lineax
            operator or as a callable ``v -> M^{-1} v``.
        in_structure: Input structure for the callable form. When ``None`` it is
            taken from the system operator passed to `as_operator`.
    """

    approx_inverse: lx.AbstractLinearOperator | Callable
    in_structure: object = eqx.field(static=True, default=None)

    def as_operator(
        self,
        operator: lx.AbstractLinearOperator | None = None,
    ) -> lx.AbstractLinearOperator:
        """Return the approximate inverse as a PSD lineax operator."""
        ai = self.approx_inverse
        if isinstance(ai, lx.AbstractLinearOperator):
            if lx.is_positive_semidefinite(ai):
                return ai
            return lx.TaggedLinearOperator(ai, lx.positive_semidefinite_tag)

        structure = self.in_structure
        if structure is None:
            if operator is None:
                raise ValueError(
                    "OperatorPreconditioner with a callable approximate inverse "
                    "needs `in_structure` or a system operator to infer it."
                )
            structure = operator.out_structure()
        return lx.FunctionLinearOperator(ai, structure, lx.positive_semidefinite_tag)

`as_operator(operator: lx.AbstractLinearOperator | None = None) -> lx.AbstractLinearOperator` ¶

Return the approximate inverse as a PSD lineax operator.

Source code in src/gaussx/_preconditioners/_operator.py

def as_operator(
    self,
    operator: lx.AbstractLinearOperator | None = None,
) -> lx.AbstractLinearOperator:
    """Return the approximate inverse as a PSD lineax operator."""
    ai = self.approx_inverse
    if isinstance(ai, lx.AbstractLinearOperator):
        if lx.is_positive_semidefinite(ai):
            return ai
        return lx.TaggedLinearOperator(ai, lx.positive_semidefinite_tag)

    structure = self.in_structure
    if structure is None:
        if operator is None:
            raise ValueError(
                "OperatorPreconditioner with a callable approximate inverse "
                "needs `in_structure` or a system operator to infer it."
            )
        structure = operator.out_structure()
    return lx.FunctionLinearOperator(ai, structure, lx.positive_semidefinite_tag)

`NystromPreconditioner` ¶

Randomized Nyström preconditioner for A + μ I.

For a system A + μ I with A PSD (e.g. K + σ² I), it builds a rank-l randomized Nyström approximation Â = U Λ̂ Uᵀ of the PSD part A alone (gaussx.randomized_nystrom) and applies (Frangella, Tropp & Udell, 2023)

\[ P^{-1}x = (\hat\lambda_\ell + \mu)\,U(\hat\Lambda + \mu I)^{-1}U^\top x + (x - UU^\top x). \]

On range(U) the preconditioned operator is ≈ λ̂_l + μ; on the complement it is A + μ I itself, with eigenvalues in [μ, λ_{l+1} + μ]. So

\[ \kappa\big(P^{-1/2}(A+\mu I)P^{-1/2}\big) \le \frac{\hat\lambda_\ell + \mu + \|A - \hat A\|}{\mu}, \]

which is O(1) once l ≳ d_eff(μ) = tr(A (A + μ I)⁻¹), the effective dimension. With l = 2⌈1.5 d_eff(μ)⌉ + 1 the expected condition number is below 28, so CG needs a number of iterations that does not grow with n.

Build it once with from_operator on the PSD part A and an explicit shift=μ, like gaussx.PartialCholeskyPreconditioner: A must not already contain μ, or the noise is counted twice (#345). The factors are stored as arrays, so as_operator ignores its argument and repeated solves cost no further matvecs.

Covariance form only. Â captures the top of A's spectrum, which is right for K + σ² I. It is not a good preconditioner for a precision-form GMRF system Q + Aᵀ W A, whose hard directions are its smallest eigenvalues; use an exact sparse factorisation or Jacobi-preconditioned CG there.

Up to gaussx 0.4 this class was a randomized Rayleigh-Ritz projection onto a random subspace, which made CG slower below full rank (#354). For a Rayleigh-Ritz eigendecomposition use gaussx.randomized_eigh(op, rank, n_power_iter=0); it projects onto orth(AΩ) rather than orth(Ω), which is more accurate, so it is not numerically identical to the old preconditioner.

Attributes:

Name	Type	Description
`basis`	`Float[Array, 'n l']`	Orthonormal `U`, shape `(n, l)`.
`eigenvalues`	`Float[Array, ' l']`	`Λ̂`, descending, shape `(l,)`.
`shift`	`Float[Array, '']`	`μ` (e.g. the noise variance `σ²`).

Examples:

A GP with n = 200k and a Matérn-3/2 kernel is solved by preconditioned CG in tens of iterations, not thousands::

K_op = kl.to_operator(kernel, X, implicit=True)
P = gx.NystromPreconditioner.from_operator(
    K_op, rank=500, shift=noise_var, key=key
)
A_op = K_op + lx.DiagonalLinearOperator(jnp.full(n, noise_var))
alpha = gx.PreconditionedCGSolver(preconditioner=P).solve(A_op, y)

A small runnable version:

>>> import einx, jax.numpy as jnp, jax.random as jr, lineax as lx
>>> import gaussx as gx
>>> x = jnp.linspace(0.0, 10.0, 100)
>>> K = jnp.exp(-0.5 * einx.subtract("i, j -> i j", x, x) ** 2)
>>> psd = lx.positive_semidefinite_tag
>>> P = gx.NystromPreconditioner.from_operator(
...     lx.MatrixLinearOperator(K, psd), rank=30, shift=0.1, key=jr.key(0)
... )
>>> A = lx.MatrixLinearOperator(K + 0.1 * jnp.eye(100), psd)
>>> alpha = gx.PreconditionedCGSolver(preconditioner=P).solve(A, jnp.ones(100))

Source code in src/gaussx/_preconditioners/_nystrom.py

class NystromPreconditioner(AbstractPreconditioner):
    r"""Randomized Nyström preconditioner for ``A + μ I``.

    For a system ``A + μ I`` with ``A`` PSD (e.g. ``K + σ² I``), it builds a
    rank-``l`` randomized Nyström approximation ``Â = U Λ̂ Uᵀ`` of the PSD
    part ``A`` alone (`gaussx.randomized_nystrom`) and applies (Frangella,
    Tropp & Udell, 2023)

    $$
    P^{-1}x = (\hat\lambda_\ell + \mu)\,U(\hat\Lambda + \mu I)^{-1}U^\top x
    + (x - UU^\top x).
    $$

    On ``range(U)`` the preconditioned operator is ``≈ λ̂_l + μ``; on the
    complement it is ``A + μ I`` itself, with eigenvalues in
    ``[μ, λ_{l+1} + μ]``. So

    $$
    \kappa\big(P^{-1/2}(A+\mu I)P^{-1/2}\big)
    \le \frac{\hat\lambda_\ell + \mu + \|A - \hat A\|}{\mu},
    $$

    which is ``O(1)`` once ``l ≳ d_eff(μ) = tr(A (A + μ I)⁻¹)``, the
    effective dimension. With ``l = 2⌈1.5 d_eff(μ)⌉ + 1`` the expected
    condition number is below 28, so CG needs a number of iterations that
    does not grow with ``n``.

    Build it once with `from_operator` on the PSD part ``A`` and an explicit
    ``shift=μ``, like `gaussx.PartialCholeskyPreconditioner`: ``A`` must not
    already contain ``μ``, or the noise is counted twice (#345). The
    factors are stored as arrays, so `as_operator` ignores its argument and
    repeated solves cost no further matvecs.

    **Covariance form only.** ``Â`` captures the *top* of ``A``'s spectrum,
    which is right for ``K + σ² I``. It is not a good preconditioner for a
    precision-form GMRF system ``Q + Aᵀ W A``, whose hard directions are its
    smallest eigenvalues; use an exact sparse factorisation or
    Jacobi-preconditioned CG there.

    Up to gaussx 0.4 this class was a randomized Rayleigh-Ritz projection
    onto a random subspace, which made CG *slower* below full rank (#354).
    For a Rayleigh-Ritz eigendecomposition use
    ``gaussx.randomized_eigh(op, rank, n_power_iter=0)``; it projects onto
    ``orth(AΩ)`` rather than ``orth(Ω)``, which is more accurate, so it is not
    numerically identical to the old preconditioner.

    Attributes:
        basis: Orthonormal ``U``, shape ``(n, l)``.
        eigenvalues: ``Λ̂``, descending, shape ``(l,)``.
        shift: ``μ`` (e.g. the noise variance ``σ²``).

    Examples:
        A GP with n = 200k and a Matérn-3/2 kernel is solved by preconditioned
        CG in tens of iterations, not thousands::

            K_op = kl.to_operator(kernel, X, implicit=True)
            P = gx.NystromPreconditioner.from_operator(
                K_op, rank=500, shift=noise_var, key=key
            )
            A_op = K_op + lx.DiagonalLinearOperator(jnp.full(n, noise_var))
            alpha = gx.PreconditionedCGSolver(preconditioner=P).solve(A_op, y)

        A small runnable version:

        >>> import einx, jax.numpy as jnp, jax.random as jr, lineax as lx
        >>> import gaussx as gx
        >>> x = jnp.linspace(0.0, 10.0, 100)
        >>> K = jnp.exp(-0.5 * einx.subtract("i, j -> i j", x, x) ** 2)
        >>> psd = lx.positive_semidefinite_tag
        >>> P = gx.NystromPreconditioner.from_operator(
        ...     lx.MatrixLinearOperator(K, psd), rank=30, shift=0.1, key=jr.key(0)
        ... )
        >>> A = lx.MatrixLinearOperator(K + 0.1 * jnp.eye(100), psd)
        >>> alpha = gx.PreconditionedCGSolver(preconditioner=P).solve(A, jnp.ones(100))
    """

    basis: Float[Array, "n l"]
    eigenvalues: Float[Array, " l"]
    shift: Float[Array, ""]

    @classmethod
    def from_operator(
        cls,
        operator: lx.AbstractLinearOperator,
        rank: int = 50,
        *,
        shift: float | Float[Array, ""],
        oversample: int = 0,
        key: jax.Array | None = None,
    ) -> NystromPreconditioner:
        """Sketch the PSD part once and store the preconditioner.

        Args:
            operator: The PSD part ``A`` (e.g. a kernel matrix ``K``),
                **not** ``A + μ I``. May be matrix-free.
            rank: Sketch size ``l`` (``rank`` matvecs of ``A``), clamped to
                the operator size. Aim for ``l ≳ 2⌈1.5 d_eff(μ)⌉ + 1``.
            shift: ``μ``, e.g. the noise variance ``σ²``; the preconditioner
                targets ``A + μ I``. Must be positive.
            oversample: Extra sketch columns, discarded after the
                approximation (see `gaussx.randomized_nystrom`).
            key: PRNG key for the test matrix. ``None`` means
                ``jax.random.PRNGKey(0)``.

        Returns:
            A built `NystromPreconditioner`.

        Raises:
            ValueError: If ``rank < 1`` or a concrete ``shift`` is not
                positive.
        """
        if isinstance(shift, (int, float)) and shift <= 0:
            raise ValueError(f"shift must be positive, got {shift}.")
        approx = randomized_nystrom(operator, rank, oversample=oversample, key=key)
        return cls(
            basis=approx.U,
            eigenvalues=approx.d,
            shift=jnp.asarray(shift, dtype=approx.d.dtype),
        )

    def as_operator(
        self,
        operator: lx.AbstractLinearOperator | None = None,
    ) -> lx.AbstractLinearOperator:
        """Return ``P⁻¹`` as a PSD operator; *operator* is ignored."""
        U, mu = self.basis, self.shift
        # P⁻¹ = I + U diag((λ̂_l + μ)/(λ̂ + μ) − 1) Uᵀ.
        scale = (self.eigenvalues[-1] + mu) / (self.eigenvalues + mu) - 1.0
        structure = jax.ShapeDtypeStruct((U.shape[0],), U.dtype)

        def matvec(x: Float[Array, " n"]) -> Float[Array, " n"]:
            coeffs = einsum(U, x, "n l, n -> l")
            return x + einsum(U, scale * coeffs, "n l, l -> n")

        return lx.FunctionLinearOperator(
            matvec, structure, lx.positive_semidefinite_tag
        )

`from_operator(operator: lx.AbstractLinearOperator, rank: int = 50, *, shift: float | Float[Array, ''], oversample: int = 0, key: jax.Array | None = None) -> NystromPreconditioner` `classmethod` ¶

Sketch the PSD part once and store the preconditioner.

Parameters:

Name	Type	Description	Default
`operator`	`AbstractLinearOperator`	The PSD part `A` (e.g. a kernel matrix `K`), not `A + μ I`. May be matrix-free.	required
`rank`	`int`	Sketch size `l` (`rank` matvecs of `A`), clamped to the operator size. Aim for `l ≳ 2⌈1.5 d_eff(μ)⌉ + 1`.	`50`
`shift`	`float \| Float[Array, '']`	`μ`, e.g. the noise variance `σ²`; the preconditioner targets `A + μ I`. Must be positive.	required
`oversample`	`int`	Extra sketch columns, discarded after the approximation (see `gaussx.randomized_nystrom`).	`0`
`key`	`Array \| None`	PRNG key for the test matrix. `None` means `jax.random.PRNGKey(0)`.	`None`

Returns:

Type	Description
`NystromPreconditioner`	A built `NystromPreconditioner`.

Raises:

Type	Description
`ValueError`	If `rank < 1` or a concrete `shift` is not positive.

Source code in src/gaussx/_preconditioners/_nystrom.py

@classmethod
def from_operator(
    cls,
    operator: lx.AbstractLinearOperator,
    rank: int = 50,
    *,
    shift: float | Float[Array, ""],
    oversample: int = 0,
    key: jax.Array | None = None,
) -> NystromPreconditioner:
    """Sketch the PSD part once and store the preconditioner.

    Args:
        operator: The PSD part ``A`` (e.g. a kernel matrix ``K``),
            **not** ``A + μ I``. May be matrix-free.
        rank: Sketch size ``l`` (``rank`` matvecs of ``A``), clamped to
            the operator size. Aim for ``l ≳ 2⌈1.5 d_eff(μ)⌉ + 1``.
        shift: ``μ``, e.g. the noise variance ``σ²``; the preconditioner
            targets ``A + μ I``. Must be positive.
        oversample: Extra sketch columns, discarded after the
            approximation (see `gaussx.randomized_nystrom`).
        key: PRNG key for the test matrix. ``None`` means
            ``jax.random.PRNGKey(0)``.

    Returns:
        A built `NystromPreconditioner`.

    Raises:
        ValueError: If ``rank < 1`` or a concrete ``shift`` is not
            positive.
    """
    if isinstance(shift, (int, float)) and shift <= 0:
        raise ValueError(f"shift must be positive, got {shift}.")
    approx = randomized_nystrom(operator, rank, oversample=oversample, key=key)
    return cls(
        basis=approx.U,
        eigenvalues=approx.d,
        shift=jnp.asarray(shift, dtype=approx.d.dtype),
    )

`as_operator(operator: lx.AbstractLinearOperator | None = None) -> lx.AbstractLinearOperator` ¶

Return P⁻¹ as a PSD operator; operator is ignored.

Source code in src/gaussx/_preconditioners/_nystrom.py

def as_operator(
    self,
    operator: lx.AbstractLinearOperator | None = None,
) -> lx.AbstractLinearOperator:
    """Return ``P⁻¹`` as a PSD operator; *operator* is ignored."""
    U, mu = self.basis, self.shift
    # P⁻¹ = I + U diag((λ̂_l + μ)/(λ̂ + μ) − 1) Uᵀ.
    scale = (self.eigenvalues[-1] + mu) / (self.eigenvalues + mu) - 1.0
    structure = jax.ShapeDtypeStruct((U.shape[0],), U.dtype)

    def matvec(x: Float[Array, " n"]) -> Float[Array, " n"]:
        coeffs = einsum(U, x, "n l, n -> l")
        return x + einsum(U, scale * coeffs, "n l, l -> n")

    return lx.FunctionLinearOperator(
        matvec, structure, lx.positive_semidefinite_tag
    )

`PartialCholeskyPreconditioner` ¶