API reference

Generated from the source. See Operators for the classes and Configuration for settings.

Creation

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

as_linop(a: Any) -> LinearOperator

Convert object to LinearOperator.

Source code in linox/api.py

def as_linop(a: Any) -> LinearOperator:
    """Convert object to LinearOperator."""
    return _array_module.as_linop(a)

todense(a: LinearlyOperatorLike) -> jax.Array

Convert operator to dense matrix.

Source code in linox/api.py

def todense(a: LinearlyOperatorLike) -> jax.Array:
    """Convert operator to dense matrix."""
    return _array_module.todense(a)

eye(N: int, M: int | None = None, k: int = 0, dtype: Any = None) -> LinearOperator

Return a 2-D array with ones on the diagonal and zeros elsewhere.

Args: N: Number of rows in the output. M: Number of columns in the output. If None, defaults to N. k: Index of the diagonal: 0 (the default) refers to the main diagonal, a positive value refers to an upper diagonal, and a negative value to a lower diagonal. dtype: Data-type of the returned array.

Source code in linox/api.py

def eye(N: int, M: int | None = None, k: int = 0, dtype: Any = None) -> LinearOperator:
    """Return a 2-D array with ones on the diagonal and zeros elsewhere.

    Args:
        N: Number of rows in the output.
        M: Number of columns in the output. If None, defaults to N.
        k: Index of the diagonal: 0 (the default) refers to the main diagonal,
            a positive value refers to an upper diagonal, and a negative value
            to a lower diagonal.
        dtype: Data-type of the returned array.
    """
    if M is None:
        M = N

    if k == 0 and N == M:
        return Identity(N, dtype=dtype)

    # For off-diagonals or non-square, just use Matrix wrapping dense eye for now
    return Matrix(jnp.eye(N, M, k=k, dtype=dtype))

zeros(dim: Int, shape: tuple[Int, Int] | None = None) -> LinearOperator

Create a zero operator.

Args: dim: Dimension or number of rows. shape: Optional shape tuple (rows, cols).

Source code in linox/api.py

def zeros(dim: Int, shape: tuple[Int, Int] | None = None) -> LinearOperator:
    """Create a zero operator.

    Args:
        dim: Dimension or number of rows.
        shape: Optional shape tuple (rows, cols).
    """
    if shape is None:
        return Zero((dim, dim))
    return Zero(shape)

ones(dim: Int, shape: tuple[Int, Int] | None = None) -> LinearOperator

Create a ones operator.

Source code in linox/api.py

def ones(dim: Int, shape: tuple[Int, Int] | None = None) -> LinearOperator:
    """Create a ones operator."""
    if shape is None:
        return Ones((dim, dim))
    return Ones(shape)

diag(v: jax.Array) -> LinearOperator

Create a diagonal operator from a vector.

Source code in linox/api.py

def diag(v: jax.Array) -> LinearOperator:
    """Create a diagonal operator from a vector."""
    return Diagonal(v)

kron(a: LinearlyOperatorLike, b: LinearlyOperatorLike) -> LinearOperator

Compute the Kronecker product of two linear operators.

Source code in linox/api.py

def kron(a: LinearlyOperatorLike, b: LinearlyOperatorLike) -> LinearOperator:
    """Compute the Kronecker product of two linear operators."""
    return Kronecker(ensure_linop(a), ensure_linop(b))

block_diag(*opers: LinearlyOperatorLike) -> LinearOperator

Construct a block diagonal operator from input operators.

Source code in linox/api.py

def block_diag(*opers: LinearlyOperatorLike) -> LinearOperator:
    """Construct a block diagonal operator from input operators."""
    return BlockDiagonal(*(ensure_linop(op) for op in opers))

bmat(blocks: list[list[LinearlyOperatorLike]]) -> LinearOperator

Construct a block matrix from a list of lists of operators.

Source code in linox/api.py

def bmat(blocks: list[list[LinearlyOperatorLike]]) -> LinearOperator:
    """Construct a block matrix from a list of lists of operators."""
    linop_blocks = [[ensure_linop(op) for op in row] for row in blocks]
    return BlockMatrix(linop_blocks)

toeplitz(c: jax.Array, r: jax.Array | None = None) -> LinearOperator

Construct a Toeplitz operator from column c and optional row r.

If r is None, assumes symmetric Toeplitz (r = c).

Source code in linox/api.py

def toeplitz(c: jax.Array, r: jax.Array | None = None) -> LinearOperator:
    """Construct a Toeplitz operator from column c and optional row r.

    If r is None, assumes symmetric Toeplitz (r = c).
    """
    if r is not None:
        msg = "Asymmetric Toeplitz not yet supported via simple wrapper."
        raise NotImplementedError(msg)
    return Toeplitz(c)

Solving

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

solve(a: LinearlyOperatorLike, b: jax.Array, method: str = 'auto', *, throw: bool = True, return_info: bool = False, residual_rtol: float = 1e-05, **kwargs) -> jax.Array | tuple[jax.Array, Solution]

Solve linear system Ax = b.

Args: a: Linear operator. b: Right-hand side vector/matrix. method: Solver method ("exact", "lsmr", "cg", "auto"). "cg" uses linox's own preconditioned conjugate gradients, which requires a symmetric positive-definite operator and accepts a preconditioner= operator. throw: Raise :class:~linox.linalg.solution.LinearSolveError when the solve fails (the default). Pass False to accept whatever the solver produced. Under jax.jit the outcome is a tracer and cannot be raised at trace time, so the failure is reported by a runtime callback instead -- branch on info.result if you need to handle it inside the computation. return_info: Also return a :class:~linox.linalg.solution.Solution carrying the outcome code and solver diagnostics. residual_rtol: Relative residual above which a square solve is judged to have failed. A singular direct solve typically returns finite, enormous values rather than NaN, so the residual is the only reliable detector.

Returns:

Type	Description
``x``, or ``(x, info)`` when ``return_info=True``.

Raises:

Type	Description
LinearSolveError	If the solve failed and throw=True.

Source code in linox/api.py

def solve(
    a: LinearlyOperatorLike,
    b: jax.Array,
    method: str = "auto",
    *,
    throw: bool = True,
    return_info: bool = False,
    residual_rtol: float = 1e-5,
    **kwargs,
) -> jax.Array | tuple[jax.Array, Solution]:
    """Solve linear system Ax = b.

    Args:
        a: Linear operator.
        b: Right-hand side vector/matrix.
        method: Solver method ("exact", "lsmr", "cg", "auto"). ``"cg"`` uses
            linox's own preconditioned conjugate gradients, which requires a
            symmetric positive-definite operator and accepts a
            ``preconditioner=`` operator.
        throw: Raise :class:`~linox.linalg.solution.LinearSolveError` when the
            solve fails (the default). Pass ``False`` to accept whatever the
            solver produced. Under ``jax.jit`` the outcome is a tracer and
            cannot be raised at trace time, so the failure is reported by a
            runtime callback instead -- branch on ``info.result`` if you need
            to handle it inside the computation.
        return_info: Also return a :class:`~linox.linalg.solution.Solution`
            carrying the outcome code and solver diagnostics.
        residual_rtol: Relative residual above which a square solve is judged
            to have failed. A singular direct solve typically returns finite,
            enormous values rather than NaN, so the residual is the only
            reliable detector.

    Returns
    -------
        ``x``, or ``(x, info)`` when ``return_info=True``.

    Raises
    ------
    LinearSolveError
        If the solve failed and ``throw=True``.
    """
    op = ensure_linop(a)
    b = jnp.asarray(b)

    m = config.resolve_method("solve", op, method)

    stats: dict[str, jax.Array] = {}
    result: _RESULTS | jax.Array | None = None
    # Tolerance for the residual sanity check, or None to skip it because the
    # solver reported its own outcome.
    check_rtol: float | None = residual_rtol

    if m == "exact" or _has_structured_solver(op):
        x = _lsolve_impl(op, b, **kwargs)
    elif m == "lsmr":
        from linox.linalg.approx.lsmr import lsmr_solve

        x, info = lsmr_solve(op, b, **kwargs)
        stats = dict(info)
        # LSMR reports its own termination code, and stops at *its* tolerance
        # rather than machine precision. Second-guessing that with a tighter
        # residual threshold would flag perfectly good converged solves.
        result = _lsmr_result(info["istop"])
        check_rtol = None
    elif m in {"cg", "conjugate_gradient"}:
        from linox.linalg.approx.cg import CG_CONVERGED, cg_solve

        x, info = cg_solve(op, b, **kwargs)
        stats = dict(info)
        # Unlike `jax.scipy.sparse.linalg.cg`, this reports whether it
        # converged, so the loose residual guard is no longer needed.
        result = jnp.where(
            jnp.asarray(info["istop"]) == CG_CONVERGED,
            jnp.int32(_RESULTS.successful),
            jnp.int32(_RESULTS.max_steps_reached),
        )
        check_rtol = None
    else:
        x = _lsolve_impl(op, b, **kwargs)

    # Residual check for square systems. Rectangular ones legitimately have a
    # nonzero residual (that is least squares, not failure), so skip them.
    if check_rtol is not None and is_square(op):
        result, residual = _residual_result(op, x, b, rtol=check_rtol)
        stats = {**stats, "residual": residual}

    if result is None:
        result = _RESULTS.successful

    _check_result(result, throw=throw, detail=f"operator: {op}")

    if return_info:
        return x, Solution(value=x, result=result, stats=stats)
    return x

psolve(a: LinearOperator, b: jax.Array, rtol=1e-08) -> jax.Array

Solve Ax = b using pseudo-inverse. See lpsolve for implementation details.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def psolve(a: LinearOperator, b: jax.Array, rtol=1e-8) -> jax.Array:
    """Solve Ax = b using pseudo-inverse. See lpsolve for implementation details."""
    return lpsolve(a, b, rtol)

inverse(a: LinearlyOperatorLike, method: str = 'auto', **kwargs) -> LinearOperator

Compute the inverse of a linear operator.

Args: a: Linear operator. method: "exact" (lazy inverse) or "auto".

Source code in linox/api.py

def inverse(a: LinearlyOperatorLike, method: str = "auto", **kwargs) -> LinearOperator:
    """Compute the inverse of a linear operator.

    Args:
        a: Linear operator.
        method: "exact" (lazy inverse) or "auto".
    """
    op = ensure_linop(a)
    m = config.resolve_method("inverse", op, method)

    if m == "exact":
        return linverse(op)

    if m == "approx":
        # Default to LSMR for approximate inverse
        m = "lsmr"

    # Return lazy inverse with specified solver method
    return InverseLinearOperator(op, method=m, solver_options=kwargs)

pinverse(a: LinearlyOperatorLike, method: str = 'auto', **kwargs) -> LinearOperator

Compute the pseudo-inverse of a linear operator.

Source code in linox/api.py

def pinverse(a: LinearlyOperatorLike, method: str = "auto", **kwargs) -> LinearOperator:
    """Compute the pseudo-inverse of a linear operator."""
    op = ensure_linop(a)
    # Similar method resolution could apply
    return lpinverse(op)

lu_factor(a: LinearOperator, overwrite_a: bool = False) -> tuple[jax.Array, jax.Array]

LU factorization of a linear operator.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def lu_factor(
    a: LinearOperator,
    overwrite_a: bool = False,
) -> tuple[jax.Array, jax.Array]:
    """LU factorization of a linear operator."""
    _warn(f"Linear operator {a} is densed for lu_factor computation.")
    return jax.scipy.linalg.lu_factor(a._todense(), overwrite_a=overwrite_a)

lu_solve(a: LinearOperator, b: jax.Array) -> jax.Array

Solve the linear system Ax = b given the LU factorization of A.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def lu_solve(a: LinearOperator, b: jax.Array) -> jax.Array:
    """Solve the linear system Ax = b given the LU factorization of A."""
    if a.shape[-1] != _rhs_rows(b):
        msg = f"Shape mismatch: {a.shape} and {b.shape}"
        raise ValueError(msg)
    lu, piv = lu_factor(a)
    _warn(f"Linear operator {a} is densed for lu_solve computation.")
    return jax.scipy.linalg.lu_solve((lu, piv), b, overwrite_b=False)

Outcome reporting for linear solves.

Before this module, :func:linox.solve returned a bare array with no way to tell whether it meant anything. A singular system produced finite, plausible, wildly wrong numbers -- no exception, no warning, no NaN -- which then propagated silently into whatever the caller did next.

The types here follow the shape of lineax's Solution/RESULTS: solves report an outcome code, and by default a failed solve raises rather than handing back garbage.

LinearSolveError

Bases: RuntimeError

Raised when a linear solve fails and throw=True (the default).

Source code in linox/linalg/solution.py

class LinearSolveError(RuntimeError):
    """Raised when a linear solve fails and ``throw=True`` (the default)."""

    def __init__(self, result: RESULTS, detail: str = "") -> None:
        self.result = result
        message = f"{RESULTS(result).name}: {RESULTS(result).message}"
        if detail:
            message = f"{message}\n{detail}"
        super().__init__(message)

RESULTS

Bases: IntEnum

Outcome of a linear solve.

Source code in linox/linalg/solution.py

class RESULTS(enum.IntEnum):
    """Outcome of a linear solve."""

    successful = 0
    max_steps_reached = 1
    singular = 2
    breakdown = 3
    stagnation = 4
    conlim = 5
    nonfinite_input = 6
    nonfinite_output = 7

    @property
    def message(self) -> str:
        """Human-readable explanation of this outcome."""
        return _MESSAGES[self]

message: str property

Human-readable explanation of this outcome.

Solution dataclass

The outcome of a linear solve.

Attributes:

Name	Type	Description
value	Array	The solution array. Meaningful only when result is RESULTS.successful.
result	RESULTS | Array	The outcome code. May be a traced array under jax.jit, in which case compare it against :class:RESULTS members with jnp ops rather than Python ==.
stats	dict[str, Any]	Solver-specific diagnostics, e.g. num_steps, istop, residual.

Source code in linox/linalg/solution.py

@dataclass(frozen=True)
class Solution:
    """The outcome of a linear solve.

    Attributes
    ----------
    value:
        The solution array. Meaningful only when ``result`` is
        ``RESULTS.successful``.
    result:
        The outcome code. May be a traced array under ``jax.jit``, in which
        case compare it against :class:`RESULTS` members with ``jnp`` ops
        rather than Python ``==``.
    stats:
        Solver-specific diagnostics, e.g. ``num_steps``, ``istop``,
        ``residual``.
    """

    value: jax.Array
    result: RESULTS | jax.Array
    stats: dict[str, Any] = field(default_factory=dict)

    @property
    def successful(self) -> bool | jax.Array:
        """Whether the solve succeeded."""
        return self.result == RESULTS.successful

successful: bool | jax.Array property

Whether the solve succeeded.

check_result(result: RESULTS | jax.Array, *, throw: bool, detail: str = '') -> None

Raise (eager) or emit a runtime error message (traced) on failure.

Under jax.jit the outcome is a tracer, so there is nothing to raise at trace time. In that case a runtime callback reports the failure when it actually occurs, and the caller can still branch on Solution.result inside the computation.

Source code in linox/linalg/solution.py

def check_result(result: RESULTS | jax.Array, *, throw: bool, detail: str = "") -> None:
    """Raise (eager) or emit a runtime error message (traced) on failure.

    Under ``jax.jit`` the outcome is a tracer, so there is nothing to raise at
    trace time. In that case a runtime callback reports the failure when it
    actually occurs, and the caller can still branch on ``Solution.result``
    inside the computation.
    """
    if not throw:
        return

    try:
        code = int(result)
    except (jax.errors.ConcretizationTypeError, TypeError):
        _report_under_trace(result, detail)
        return

    if code != RESULTS.successful:
        raise LinearSolveError(RESULTS(code), detail)

residual_result(operator: Any, solution: jax.Array, rhs: jax.Array, *, rtol: float = 1e-05) -> tuple[RESULTS | jax.Array, jax.Array]

Classify a solve by its relative residual ||Ax - b|| / ||b||.

This is the only reliable detector for a direct solve against a singular operator: the output is typically finite and enormous rather than NaN, so a finiteness check alone misses it. Costs one extra matvec, which is cheap next to the factorisation it is validating.

Source code in linox/linalg/solution.py

def residual_result(
    operator: Any,
    solution: jax.Array,
    rhs: jax.Array,
    *,
    rtol: float = 1e-5,
) -> tuple[RESULTS | jax.Array, jax.Array]:
    """Classify a solve by its relative residual ``||Ax - b|| / ||b||``.

    This is the only reliable detector for a direct solve against a singular
    operator: the output is typically finite and enormous rather than NaN, so
    a finiteness check alone misses it. Costs one extra matvec, which is cheap
    next to the factorisation it is validating.
    """
    residual = jnp.linalg.norm(jnp.asarray(operator @ solution) - rhs) / jnp.maximum(jnp.linalg.norm(rhs), jnp.finfo(jnp.asarray(rhs).dtype).tiny)

    nonfinite = ~jnp.all(jnp.isfinite(jnp.asarray(solution)))
    failed = nonfinite | (residual > rtol)

    result = jnp.where(
        nonfinite,
        jnp.int32(RESULTS.nonfinite_output),
        jnp.where(failed, jnp.int32(RESULTS.singular), jnp.int32(RESULTS.successful)),
    )
    return result, residual

Decompositions

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

eigh(a: LinearlyOperatorLike, k: Int | None = None, subset_by_index: tuple[Int, Int] | None = None, method: str = 'auto', **kwargs) -> tuple[jax.Array, LinearOperator] | jax.Array

Compute eigenvalues and eigenvectors.

Args: a: Linear operator. k: Number of eigenvalues (for approx/partial). subset_by_index: Range of indices (start, end) for eigenvalues. method: "exact" or "lanczos".

Source code in linox/api.py

def eigh(
    a: LinearlyOperatorLike,
    k: Int | None = None,
    subset_by_index: tuple[Int, Int] | None = None,
    method: str = "auto",
    **kwargs,
) -> tuple[jax.Array, LinearOperator] | jax.Array:
    """Compute eigenvalues and eigenvectors.

    Args:
        a: Linear operator.
        k: Number of eigenvalues (for approx/partial).
        subset_by_index: Range of indices (start, end) for eigenvalues.
        method: "exact" or "lanczos".
    """
    m = config.resolve_method("eigh", ensure_linop(a), method)
    return _spectral_module.eigh(ensure_linop(a), k=k, subset_by_index=subset_by_index, method=m, **kwargs)

svd(a: LinearlyOperatorLike, **kwargs) -> tuple[jax.Array, jax.Array, jax.Array]

Singular Value Decomposition.

Source code in linox/api.py

def svd(a: LinearlyOperatorLike, **kwargs) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Singular Value Decomposition."""
    return _svd_impl(ensure_linop(a), **kwargs)

qr(a: LinearOperator) -> tuple[jax.Array, jax.Array]

QR decomposition of a linear operator. See lqr for implementation details.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def qr(a: LinearOperator) -> tuple[jax.Array, jax.Array]:
    """QR decomposition of a linear operator. See lqr for implementation details."""
    return lqr(a)

cholesky(a: LinearOperator) -> jax.Array

Cholesky decomposition of a linear operator. See lcholesky for implementation details.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def cholesky(a: LinearOperator) -> jax.Array:
    """Cholesky decomposition of a linear operator. See lcholesky for implementation details."""
    return lcholesky(a)

sqrt(a: LinearlyOperatorLike, method: str = 'auto', **kwargs) -> LinearOperator

Matrix square root factor.

Returns an operator S satisfying S @ S.T == a. Note this is a factor, not necessarily the symmetric principal square root: the exact path returns whatever structured factorisation is available for the operator (a Cholesky factor for a dense :class:Matrix, the elementwise root for a :class:Diagonal, and so on). Request method="lanczos" to get the principal square root a**(1/2) via a Krylov method.

Args: a: Linear operator. method: One of "auto", "exact", "approx", "lanczos".

Source code in linox/api.py

def sqrt(a: LinearlyOperatorLike, method: str = "auto", **kwargs) -> LinearOperator:
    """Matrix square root factor.

    Returns an operator ``S`` satisfying ``S @ S.T == a``. Note this is a
    *factor*, not necessarily the symmetric principal square root: the exact
    path returns whatever structured factorisation is available for the
    operator (a Cholesky factor for a dense :class:`Matrix`, the elementwise
    root for a :class:`Diagonal`, and so on). Request ``method="lanczos"``
    to get the principal square root ``a**(1/2)`` via a Krylov method.

    Args:
        a: Linear operator.
        method: One of ``"auto"``, ``"exact"``, ``"approx"``, ``"lanczos"``.
    """
    op = ensure_linop(a)
    m = config.resolve_method("sqrt", op, method)

    if m in {"lanczos", "approx"}:
        # An explicit approximate request is honoured as given. Only when
        # `auto` resolution picked the Krylov path do we prefer a structured
        # exact factorisation if one happens to exist.
        if method == "auto":
            try:
                return _lsqrt_impl(op)
            except NotImplementedError:
                pass
        return _functions_module.sqrt(op, method="lanczos", **kwargs)

    return _lsqrt_impl(op)

Traces and determinants

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

trace(a: LinearlyOperatorLike, method: str = 'auto', **kwargs) -> jax.Array

Compute the trace of a linear operator.

Args: a: Linear operator. method: Computation method ("auto", "exact", "hutchinson").

Source code in linox/api.py

def trace(a: LinearlyOperatorLike, method: str = "auto", **kwargs) -> jax.Array:
    """Compute the trace of a linear operator.

    Args:
        a: Linear operator.
        method: Computation method ("auto", "exact", "hutchinson").
    """
    op = ensure_linop(a)
    m = config.resolve_method("trace", op, method)

    # `auto` picks Hutchinson for large operators, but that needs a PRNG key.
    # Without one, fall back to the exact path rather than failing a plain
    # `trace(a)` call purely because the operator is big.
    if m == "hutchinson" and method == "auto" and kwargs.get("key") is None:
        m = "exact"

    if m == "hutchinson":
        return _trace_module.trace(op, method="hutchinson", **kwargs)
    return _trace_module.trace(op, **kwargs)

det(a: LinearlyOperatorLike) -> jax.Array

Compute determinant.

Source code in linox/api.py

def det(a: LinearlyOperatorLike) -> jax.Array:
    """Compute determinant."""
    from linox.linalg.determinants import det as _det

    return _det(ensure_linop(a))

slogdet(a: LinearlyOperatorLike, method: str = 'auto', **kwargs) -> tuple[jax.Array, jax.Array]

Compute sign and log of determinant.

Args: a: Linear operator. method: Computation method ("auto", "exact", "slq").

Source code in linox/api.py

def slogdet(a: LinearlyOperatorLike, method: str = "auto", **kwargs) -> tuple[jax.Array, jax.Array]:
    """Compute sign and log of determinant.

    Args:
        a: Linear operator.
        method: Computation method ("auto", "exact", "slq").
    """
    from linox.linalg.determinants import slogdet as _slogdet

    op = ensure_linop(a)
    m = config.resolve_method("slogdet", op, method)

    # As in `trace`: SLQ needs a PRNG key, so an `auto` resolution that lands
    # there without one falls back to the exact path.
    if m == "slq" and method == "auto" and kwargs.get("key") is None:
        m = "exact"

    return _slogdet(op, method=m, **kwargs)

logdet(a: LinearlyOperatorLike) -> jax.Array

Compute log of determinant.

Source code in linox/api.py

def logdet(a: LinearlyOperatorLike) -> jax.Array:
    """Compute log of determinant."""
    from linox.linalg.determinants import logdet as _logdet

    return _logdet(ensure_linop(a))

diagonal(a: LinearOperator) -> jax.Array

Extract the diagonal of an operator as a :class:jax.Array.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def diagonal(a: LinearOperator) -> jax.Array:
    """Extract the diagonal of an operator as a :class:`jax.Array`."""
    _warn(f"Linear operator {a} is densed for diagonal computation.")
    dense_matrix = a._todense()
    if len(a.shape) <= 2:
        return jnp.diag(dense_matrix)
    n = dense_matrix.shape[-1]
    diag_indices = jnp.arange(n)
    return dense_matrix[..., diag_indices, diag_indices]

Matrix functions

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

exp(a: LinearlyOperatorLike, **kwargs) -> LinearOperator

Matrix exponential.

Source code in linox/api.py

def exp(a: LinearlyOperatorLike, **kwargs) -> LinearOperator:
    """Matrix exponential."""
    return _functions_module.exp(ensure_linop(a), **kwargs)

log(a: LinearlyOperatorLike, **kwargs) -> LinearOperator

Matrix logarithm.

Source code in linox/api.py

def log(a: LinearlyOperatorLike, **kwargs) -> LinearOperator:
    """Matrix logarithm."""
    return _functions_module.log(ensure_linop(a), **kwargs)

pow(a: LinearlyOperatorLike, p: float, **kwargs) -> LinearOperator

Matrix power.

Source code in linox/api.py

def pow(a: LinearlyOperatorLike, p: float, **kwargs) -> LinearOperator:
    """Matrix power."""
    return _functions_module.pow(ensure_linop(a), p, **kwargs)

Properties

Linox API - Public Functional Interface.

This module provides the main functional entry points for the Linox library.

is_square(a: LinearOperator) -> bool

Whether the operator has equal row and column counts.

Source code in linox/operators/arithmetic.py

def is_square(a: LinearOperator) -> bool:
    """Whether the operator has equal row and column counts."""
    return a.shape[-1] == a.shape[-2]

is_symmetric(a: LinearOperator, *, rtol: float = 1e-05, atol: float = 1e-08, key: jax.Array | None = None, num_probes: int = 1) -> bool

Check if a linear operator is symmetric without densifying.

Uses randomized probing: generates random vectors x and checks if Ax ≈ A^T x. This avoids densifying the full matrix.

Args: a: Linear operator to check rtol: Relative tolerance for comparison atol: Absolute tolerance for comparison key: Random key for generating test vectors (default: uses key 0) num_probes: Number of random vectors to test (default: 1)

Returns:

Type	Description
True if the operator appears symmetric within tolerance

Source code in linox/operators/arithmetic.py

def is_symmetric(
    a: LinearOperator,
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
    key: jax.Array | None = None,
    num_probes: int = 1,
) -> bool:
    """Check if a linear operator is symmetric without densifying.

    Uses randomized probing: generates random vectors x and checks if Ax ≈ A^T x.
    This avoids densifying the full matrix.

    Args:
        a: Linear operator to check
        rtol: Relative tolerance for comparison
        atol: Absolute tolerance for comparison
        key: Random key for generating test vectors (default: uses key 0)
        num_probes: Number of random vectors to test (default: 1)

    Returns
    -------
        True if the operator appears symmetric within tolerance
    """
    if not is_square(a):
        return False

    if key is None:
        key = jax.random.PRNGKey(0)

    n = a.shape[-1]

    for i in range(num_probes):
        # Generate random normalized vector
        probe_key = jax.random.fold_in(key, i)
        x = jax.random.normal(probe_key, (n,), dtype=a.dtype)
        x /= jnp.linalg.norm(x)

        # Compute Ax and A^T x
        v1 = a @ x
        v2 = a.T @ x

        # Check if v1 ≈ v2
        if not jnp.allclose(v1, v2, rtol=rtol, atol=atol):
            return False

    return True

is_hermitian(a: LinearOperator, *, rtol: float = 1e-05, atol: float = 1e-08, key: jax.Array | None = None, num_probes: int = 1) -> bool

Check if a linear operator is Hermitian without densifying.

Uses randomized probing: generates random vectors x and checks if Ax ≈ (A^H x) where A^H is the conjugate transpose. This avoids densifying the full matrix.

For real matrices, this is equivalent to checking symmetry.

Args: a: Linear operator to check rtol: Relative tolerance for comparison atol: Absolute tolerance for comparison key: Random key for generating test vectors (default: uses key 0) num_probes: Number of random vectors to test (default: 1)

Returns:

Type	Description
True if the operator appears Hermitian within tolerance

Source code in linox/operators/arithmetic.py

def is_hermitian(
    a: LinearOperator,
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
    key: jax.Array | None = None,
    num_probes: int = 1,
) -> bool:
    """Check if a linear operator is Hermitian without densifying.

    Uses randomized probing: generates random vectors x and checks if Ax ≈ (A^H x)
    where A^H is the conjugate transpose. This avoids densifying the full matrix.

    For real matrices, this is equivalent to checking symmetry.

    Args:
        a: Linear operator to check
        rtol: Relative tolerance for comparison
        atol: Absolute tolerance for comparison
        key: Random key for generating test vectors (default: uses key 0)
        num_probes: Number of random vectors to test (default: 1)

    Returns
    -------
        True if the operator appears Hermitian within tolerance
    """
    if not is_square(a):
        return False

    if key is None:
        key = jax.random.PRNGKey(0)

    n = a.shape[-1]

    for i in range(num_probes):
        # Generate random normalized vector
        probe_key = jax.random.fold_in(key, i)
        x = jax.random.normal(probe_key, (n,), dtype=a.dtype)
        if jnp.issubdtype(a.dtype, jnp.complexfloating):
            # For complex operators, add imaginary part
            x_imag = jax.random.normal(probe_key, (n,), dtype=a.dtype)
            x += 1j * x_imag
        x /= jnp.linalg.norm(x)

        # Compute Ax and A^H x (conjugate transpose)
        v1 = a @ x
        a.T @ jnp.conj(x)

        # For Hermitian: <Ax, x> = <x, Ax> = conj(<Ax, x>)
        # Equivalently: Ax should equal conj(A^T conj(x))
        # Or more directly: check if <v1, x> ≈ conj(<v2, x>)
        # But simpler: check if v1 ≈ conj(v2) when x is real
        # Actually, let's use the proper test: A x = conj(A^T conj(x))
        v2_hermitian = jnp.conj(a.T @ jnp.conj(x))

        if not jnp.allclose(v1, v2_hermitian, rtol=rtol, atol=atol):
            return False

    return True

symmetrize(a: LinearOperator) -> ArithmeticType

Return the symmetric part (a + a.T) / 2 of an operator.

Source code in linox/operators/arithmetic.py

def symmetrize(a: LinearOperator) -> ArithmeticType:
    """Return the symmetric part ``(a + a.T) / 2`` of an operator."""
    return 0.5 * (a + a.transpose())

congruence_transform(A: ArithmeticType, B: ArithmeticType) -> LinearOperator

Return the congruence transform A @ B @ A.T.

Source code in linox/operators/arithmetic.py

@plum.dispatch
def congruence_transform(A: ArithmeticType, B: ArithmeticType) -> LinearOperator:  # noqa: F811
    """Return the congruence transform ``A @ B @ A.T``."""
    return CongruenceTransform(A, B)

allclose(a: LinearOperatorLike, b: LinearOperatorLike, rtol: float = 1e-05, atol: float = 1e-08) -> bool

Check if two linear operators are close to each other.

Args: a: First linear operator. b: Second linear operator. rtol: Relative tolerance. atol: Absolute tolerance.

Returns:

Type	Description
Whether the two linear operators are close to each other.

Source code in linox/utils/array.py

def allclose(
    a: LinearOperatorLike,
    b: LinearOperatorLike,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> bool:
    """Check if two linear operators are close to each other.

    Args:
        a: First linear operator.
        b: Second linear operator.
        rtol: Relative tolerance.
        atol: Absolute tolerance.

    Returns
    -------
        Whether the two linear operators are close to each other.
    """
    a_dense = todense(a)
    b_dense = todense(b)
    return jnp.allclose(a_dense, b_dense, rtol=rtol, atol=atol)

validate(op: LinearOperator, *, mode: Literal['default', 'debug'] = 'default', rtol: float = 1e-05, atol: float = 1e-08, num_probes: int = 5, key: jax.Array | None = None) -> bool

Validate a LinearOperator and its children recursively.

Performs structural validation by default. In debug mode, also performs numerical probe-based validation to check symmetry and PSD promises.

Parameters:

Name	Type	Description	Default
op	LinearOperator	The operator to validate.	required
mode	('default', 'debug')	Validation mode. "default" performs only cheap structural checks. "debug" also performs expensive numerical probes.	"default"
rtol	float	Relative tolerance for numerical checks. Default is 1e-5.	1e-05
atol	float	Absolute tolerance for numerical checks. Default is 1e-8.	1e-08
num_probes	int	Number of random probes for numerical validation. Default is 5.	5
key	Array	JAX random key for probe generation. If None, uses a default key.	None

Returns:

Type	Description
bool	True if validation passes.

Raises:

Type	Description
ValidationError	If validation fails, with details about the failure and hints.

See Also

linox.config : For setting LINOX_DEBUG=1 to enable debug validation globally.

Examples:

>>> import linox as lo
>>> A = lo.Matrix(jnp.eye(3))
>>> lo.validate(A)  # Cheap structural validation
True

>>> # Debug mode includes probe-based checks
>>> lo.validate(A, mode="debug")
True

Source code in linox/utils/validation.py

def validate(
    op: LinearOperator,
    *,
    mode: Literal["default", "debug"] = "default",
    rtol: float = 1e-5,
    atol: float = 1e-8,
    num_probes: int = 5,
    key: jax.Array | None = None,
) -> bool:
    """Validate a LinearOperator and its children recursively.

    Performs structural validation by default. In debug mode, also performs
    numerical probe-based validation to check symmetry and PSD promises.

    Parameters
    ----------
    op : LinearOperator
        The operator to validate.
    mode : {"default", "debug"}, optional
        Validation mode. "default" performs only cheap structural checks.
        "debug" also performs expensive numerical probes.
    rtol : float, optional
        Relative tolerance for numerical checks. Default is 1e-5.
    atol : float, optional
        Absolute tolerance for numerical checks. Default is 1e-8.
    num_probes : int, optional
        Number of random probes for numerical validation. Default is 5.
    key : jax.Array, optional
        JAX random key for probe generation. If None, uses a default key.

    Returns
    -------
    bool
        True if validation passes.

    Raises
    ------
    ValidationError
        If validation fails, with details about the failure and hints.

    See Also
    --------
    linox.config : For setting LINOX_DEBUG=1 to enable debug validation globally.

    Examples
    --------
    >>> import linox as lo
    >>> A = lo.Matrix(jnp.eye(3))
    >>> lo.validate(A)  # Cheap structural validation
    True

    >>> # Debug mode includes probe-based checks
    >>> lo.validate(A, mode="debug")
    True
    """
    if key is None:
        key = jax.random.key(42)

    # Structural validation (always performed)
    _validate_structural(op)

    # Numerical validation (debug mode only)
    if mode == "debug":
        _validate_numerical(op, rtol=rtol, atol=atol, num_probes=num_probes, key=key)

    return True

Matrix-free algorithms

Preconditioned conjugate gradients for symmetric positive-definite systems.

jax.scipy.sparse.linalg.cg reports no convergence information, which left :func:linox.solve unable to say whether a CG solve had actually succeeded -- it had to fall back to a loose residual guard. This implementation reports a termination code in the same style as :mod:linox.linalg.approx.lsmr, and accepts a preconditioner.

Only matrix-vector products are used, so a matrix-free operator stays matrix-free.

cg_solve(A: LinearOperatorLike, b: jax.Array, *, preconditioner: LinearOperatorLike | None = None, rtol: float = 1e-06, atol: float = 0.0, maxiter: int | None = None, x0: jax.Array | None = None, track_iterations: bool = False) -> tuple[jax.Array, dict]

Solve A x = b for symmetric positive-definite A.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	Symmetric positive-definite operator. Only A @ v is used.	required
b	Array	Right-hand side, shape (n,).	required
preconditioner	LinearOperatorLike | None	Operator approximating A^{-1}, applied as M @ r. Must itself be symmetric positive-definite for the method to remain valid. None means no preconditioning.	None
rtol	float	Convergence is declared when ||r|| <= max(rtol * ||b||, atol).	1e-06
atol	float	Convergence is declared when ||r|| <= max(rtol * ||b||, atol).	1e-06
maxiter	int | None	Iteration cap. Defaults to 10 * n, matching SciPy and the LSMR implementation here.	None
x0	Array | None	Initial guess. Defaults to zeros.	None
track_iterations	bool	Report the exact iteration count as info["itn"], at the cost of reverse-mode differentiability. See the note below; either mode costs exactly one CG run.	False

Returns:

Name	Type	Description
x	Array	The solution.
info	dict	istop (1 converged, 2 not converged), normr and the convergence threshold atol_eff.

Notes on differentiability

The iteration is a lax.while_loop, which has no reverse-mode rule, and its counter is only observable from inside it. Those two facts pull in opposite directions, so both modes exist and each costs one CG run:

• track_iterations=False (default) routes the solution through :func:jax.lax.custom_linear_solve, which supplies the adjoint while_loop lacks -- for symmetric A the cotangent is itself a solve against A, which is how jax.scipy.sparse.linalg.cg manages it too. jax.grad works; the loop runs inside the callable, so itn is not observable.
• track_iterations=True runs the loop directly, so info["itn"] is exact. Reverse-mode differentiation then raises, while_loop having no VJP; forward mode and jax.jit are unaffected.

Notes

The iteration guards p^T A p <= 0, which cannot happen for a positive definite A and therefore signals that the operator is not positive definite -- or is too ill-conditioned to behave as though it is. The guard halts the recurrence rather than dividing and returning NaN; the caller sees non-convergence and a large residual.

Source code in linox/linalg/approx/cg.py

def cg_solve(
    A: LinearOperatorLike,
    b: jax.Array,
    *,
    preconditioner: LinearOperatorLike | None = None,
    rtol: float = 1e-6,
    atol: float = 0.0,
    maxiter: int | None = None,
    x0: jax.Array | None = None,
    track_iterations: bool = False,
) -> tuple[jax.Array, dict]:
    r"""Solve ``A x = b`` for symmetric positive-definite ``A``.

    Parameters
    ----------
    A:
        Symmetric positive-definite operator. Only ``A @ v`` is used.
    b:
        Right-hand side, shape ``(n,)``.
    preconditioner:
        Operator approximating ``A^{-1}``, applied as ``M @ r``. Must itself be
        symmetric positive-definite for the method to remain valid. ``None``
        means no preconditioning.
    rtol, atol:
        Convergence is declared when
        ``||r|| <= max(rtol * ||b||, atol)``.
    maxiter:
        Iteration cap. Defaults to ``10 * n``, matching SciPy and the LSMR
        implementation here.
    x0:
        Initial guess. Defaults to zeros.
    track_iterations:
        Report the exact iteration count as ``info["itn"]``, at the cost of
        reverse-mode differentiability. See the note below; either mode costs
        exactly one CG run.

    Returns
    -------
    x:
        The solution.
    info:
        ``istop`` (1 converged, 2 not converged), ``normr`` and the
        convergence threshold ``atol_eff``.

    Notes on differentiability
    --------------------------
    The iteration is a ``lax.while_loop``, which has no reverse-mode rule, and
    its counter is only observable from inside it. Those two facts pull in
    opposite directions, so both modes exist and each costs one CG run:

    * ``track_iterations=False`` (default) routes the solution through
      :func:`jax.lax.custom_linear_solve`, which supplies the adjoint
      ``while_loop`` lacks -- for symmetric ``A`` the cotangent is itself a
      solve against ``A``, which is how ``jax.scipy.sparse.linalg.cg`` manages
      it too. ``jax.grad`` works; the loop runs inside the callable, so ``itn``
      is not observable.
    * ``track_iterations=True`` runs the loop directly, so ``info["itn"]`` is
      exact. Reverse-mode differentiation then raises, ``while_loop`` having no
      VJP; forward mode and ``jax.jit`` are unaffected.

    Notes
    -----
    The iteration guards ``p^T A p <= 0``, which cannot happen for a positive
    definite ``A`` and therefore signals that the operator is not positive
    definite -- or is too ill-conditioned to behave as though it is. The guard
    halts the recurrence rather than dividing and returning NaN; the caller
    sees non-convergence and a large residual.
    """
    A = as_linop(A)
    b = jnp.asarray(b)
    n = A.shape[-1]

    if b.ndim != 1:
        msg = f"cg_solve expects a vector right-hand side, got shape {b.shape}. Solve each column separately, or use jax.vmap."
        raise ValueError(msg)

    dtype = jnp.result_type(A.dtype, b.dtype)
    b = b.astype(dtype)
    maxiter = 10 * n if maxiter is None else maxiter

    if preconditioner is None:

        def apply_M(residual: jax.Array) -> jax.Array:
            return residual

    else:
        M = as_linop(preconditioner)

        def apply_M(residual: jax.Array) -> jax.Array:
            return M @ residual

    # `x` is the initial guess; the residual and search direction are built
    # inside `run`, which is what `custom_linear_solve` actually invokes.
    x = jnp.zeros((n,), dtype=dtype) if x0 is None else jnp.asarray(x0, dtype=dtype)

    # A zero right-hand side has x = 0 as its exact solution; the threshold
    # then falls back to `atol` so the loop exits immediately rather than
    # chasing a relative tolerance against zero.
    atol_eff = jnp.maximum(rtol * jnp.linalg.norm(b), atol)

    def cond(state):
        _x, r, _z, _rz, itn, istop = state
        return (istop == 0) & (itn < maxiter) & (jnp.linalg.norm(r) > atol_eff)

    def body(state):
        x, r, p, rz, itn, istop = state
        Ap = A @ p
        pAp = jnp.vdot(p, Ap)

        # Guard the division rather than producing NaN.
        broken = ~jnp.isfinite(pAp) | (pAp <= 0)
        safe_pAp = jnp.where(broken, jnp.ones_like(pAp), pAp)
        alpha = rz / safe_pAp

        x_new = x + alpha * p
        r_new = r - alpha * Ap
        z_new = apply_M(r_new)
        rz_new = jnp.vdot(r_new, z_new)
        beta = rz_new / jnp.where(rz == 0, jnp.ones_like(rz), rz)
        p_new = z_new + beta * p

        istop_new = jnp.where(broken, CG_NOT_CONVERGED, istop)
        keep = ~broken
        return (
            jnp.where(keep, x_new, x),
            jnp.where(keep, r_new, r),
            jnp.where(keep, p_new, p),
            jnp.where(keep, rz_new, rz),
            itn + 1,
            istop_new,
        )

    def run(_matvec, rhs: jax.Array) -> jax.Array:
        """One CG run, returning only the solution.

        `custom_linear_solve` passes the matvec as the first argument; we
        close over `A` directly instead, so it is unused.
        """
        r0 = rhs - A @ x
        z0 = apply_M(r0)
        init = (x, r0, z0, jnp.vdot(r0, z0), jnp.asarray(0), jnp.asarray(0))
        solution, *_ = jax.lax.while_loop(cond, body, init)
        return solution

    extra: dict[str, jax.Array] = {}
    if track_iterations:
        # Run the loop directly: `itn` is observable, reverse mode is not.
        r0 = b - A @ x
        z0 = apply_M(r0)
        init = (x, r0, z0, jnp.vdot(r0, z0), jnp.asarray(0), jnp.asarray(0))
        x, _r, _z, _rz, itn, _istop = jax.lax.while_loop(cond, body, init)
        extra["itn"] = itn
    else:
        x = jax.lax.custom_linear_solve(lambda v: A @ v, b, run, symmetric=True)

    # One extra matvec, computed identically in both modes so the reported
    # outcome never depends on which was chosen.
    normr = jnp.linalg.norm(b - A @ x)
    istop = jnp.where(normr <= atol_eff, CG_CONVERGED, CG_NOT_CONVERGED)

    return x, {"istop": istop, "normr": normr, "atol_eff": atol_eff, **extra}

Lanczos tridiagonalization and the Krylov methods built on it.

lanczos_eigh(A: LinearOperatorLike, v0: ArrayLike, num_iters: int, k: int | None = None, which: str = 'LM', reortho: bool = True) -> tuple[jax.Array, jax.Array]

Compute a few eigenvalues/eigenvectors using Lanczos method.

Source code in linox/linalg/approx/lanczos.py

def lanczos_eigh(
    A: LinearOperatorLike,
    v0: ArrayLike,
    num_iters: int,
    k: int | None = None,
    which: str = "LM",
    reortho: bool = True,
) -> tuple[jax.Array, jax.Array]:
    """Compute a few eigenvalues/eigenvectors using Lanczos method."""
    if k is None:
        k = num_iters

    # Perform Lanczos tridiagonalization
    Q, alpha, beta = lanczos_tridiag(A, v0, num_iters, reortho=reortho)

    # Construct tridiagonal matrix
    T = jnp.diag(alpha)
    if beta.size > 0:
        T = T + jnp.diag(beta, k=1) + jnp.diag(beta, k=-1)

    # Dense eigenvalue decomposition
    eig_vals, eig_vecs = jnp.linalg.eigh(T)

    # Select k eigenvalues based on 'which'
    if which == "LM":
        idx = jnp.argsort(jnp.abs(eig_vals))[::-1][:k]
    elif which == "LA":
        idx = jnp.argsort(eig_vals)[::-1][:k]
    elif which == "SA":
        idx = jnp.argsort(eig_vals)[:k]
    else:
        msg = f"Invalid 'which' parameter: {which}"
        raise ValueError(msg)

    # Project eigenvectors back
    eigenvalues = eig_vals[idx]
    eigenvectors = Q @ eig_vecs[:, idx]

    return eigenvalues, eigenvectors

lanczos_matrix_function(A: LinearOperatorLike, v: ArrayLike, func: callable, num_iters: int, reortho: bool = True) -> jax.Array

Approximate f(A)v using Lanczos tridiagonalization.

Source code in linox/linalg/approx/lanczos.py

def lanczos_matrix_function(
    A: LinearOperatorLike,
    v: ArrayLike,
    func: callable,
    num_iters: int,
    reortho: bool = True,
) -> jax.Array:
    """Approximate f(A)v using Lanczos tridiagonalization."""
    v = jnp.asarray(v)
    v_norm = jnp.linalg.norm(v)
    v_normalized = v / v_norm

    # Perform Lanczos tridiagonalization
    Q, alpha, beta = lanczos_tridiag(A, v_normalized, num_iters, reortho=reortho)

    # Construct tridiagonal matrix
    T = jnp.diag(alpha)
    if beta.size > 0:
        T = T + jnp.diag(beta, k=1) + jnp.diag(beta, k=-1)

    eigvals, eigvecs = jnp.linalg.eigh(T)

    # Only the first column of f(T) is ever used, and it is the Gauss
    # quadrature sum  sum_i (e1^T u_i) f(theta_i) u_i.  Forming f(T) in full
    # and then multiplying by e1 is both wasteful and numerically fragile.
    #
    # Fragile because Lanczos *breaks down* once the Krylov space is
    # exhausted: for an operator with a degenerate spectrum (Identity being
    # the extreme case) beta hits zero after one step, and the remaining rows
    # of T are numerical noise whose eigenvalues sit at or below zero. Those
    # spurious modes carry essentially no quadrature weight, but evaluating
    # `func` on them produces +/-inf for functions like log, and the
    # subsequent `0 * inf` turns the entire result into NaN.
    #
    # So: drop the modes that carry no weight, and never evaluate `func` on
    # their eigenvalues at all -- substituting a safe value inside the
    # `where` rather than masking afterwards, since `jnp.where` still
    # evaluates both branches.
    weights = eigvecs[0, :]
    eps = jnp.finfo(eigvals.dtype).eps
    significant = jnp.abs(weights) > eps * jnp.maximum(jnp.max(jnp.abs(weights)), 1.0)

    safe_eigvals = jnp.where(significant, eigvals, jnp.ones_like(eigvals))
    contributions = jnp.where(significant, weights * func(safe_eigvals), 0.0)

    # Project back: v_norm * Q @ (f(T) e1)
    return v_norm * (Q @ (eigvecs @ contributions))

lanczos_solve_sqrt(A: LinearOperatorLike, b: ArrayLike, tol=1e-05, min_eta=1e-14, max_iter=10, overwrite_b=False) -> jax.Array

Build a low-rank inverse factor for a PSD operator using CG/Lanczos.

Returns a skinny matrix D whose columns are A-conjugate directions (normalized by sqrt of the Rayleigh quotient), such that D @ D.T ≈ A^{-1} on the generated Krylov subspace. This acts like an "inverse sqrt" factor usable in Kronecker products, preconditioners, or low-rank approximations. Note this is not the symmetric A^{-1/2}; it is a factor whose Gram approximates A^{-1}.

Parameters:

Name	Type	Description	Default
A	array-like or linear operator supporting `A @ x`	Positive semi-definite operator.	required
b	array	Start vector for the Krylov process (will be normalized).	required
tol	float	Relative tolerance for residual norm stopping.	1e-05
min_eta	float	Minimum step Rayleigh quotient to continue (guard against breakdown).	1e-14
max_iter	int	Maximum number of Lanczos/CG iterations (columns in the factor).	10
overwrite_b	bool	If True, may reuse the buffer of b for the search direction.	False

Source code in linox/linalg/approx/lanczos.py

def lanczos_solve_sqrt(
    A: LinearOperatorLike,
    b: ArrayLike,
    tol=1e-5,
    min_eta=1e-14,
    max_iter=10,
    overwrite_b=False,
) -> jax.Array:
    """Build a low-rank inverse factor for a PSD operator using CG/Lanczos.

    Returns a skinny matrix D whose columns are A-conjugate directions
    (normalized by sqrt of the Rayleigh quotient), such that
        D @ D.T ≈ A^{-1}
    on the generated Krylov subspace. This acts like an "inverse sqrt"
    factor usable in Kronecker products, preconditioners, or low-rank
    approximations. Note this is not the symmetric A^{-1/2}; it is a
    factor whose Gram approximates A^{-1}.

    Parameters
    ----------
    A : array-like or linear operator supporting `A @ x`
        Positive semi-definite operator.
    b : array
        Start vector for the Krylov process (will be normalized).
    tol : float
        Relative tolerance for residual norm stopping.
    min_eta : float
        Minimum step Rayleigh quotient to continue (guard against breakdown).
    max_iter : int
        Maximum number of Lanczos/CG iterations (columns in the factor).
    overwrite_b : bool
        If True, may reuse the buffer of `b` for the search direction.
    """

    @jax.jit
    def _step(values):
        ds, rs, rs_norm_sq, p, eta, k = values
        # Compute search direction
        true_fn = lambda _p: rs[:, k] + rs_norm_sq[k] / rs_norm_sq[k - 1] * _p  # noqa: E731
        false_fn = lambda _p: _p  # noqa: E731
        p = jax.lax.cond(k > 0, true_fn, false_fn, p)

        # Compute modified Lanzcos vector
        w = A @ p
        eta = p @ w
        ds = ds.at[:, k].set(p / jnp.sqrt(eta))

        # Update residual
        mu = rs_norm_sq[k] / eta
        rs_prev_k = rs  # rs[:, :k]
        rs = rs.at[:, k + 1].set(rs[:, k] - mu * w)

        # Full reorthogonalization of residual (double Gram-Schmidt)
        rs = rs.at[:, k + 1].set(rs[:, k + 1] - rs_prev_k @ ((rs_prev_k.T @ rs[:, k + 1]) / rs_norm_sq))
        rs = rs.at[:, k + 1].set(rs[:, k + 1] - rs_prev_k @ ((rs_prev_k.T @ rs[:, k + 1]) / rs_norm_sq))

        rs_norm_sq = rs_norm_sq.at[k + 1].set(rs[:, k + 1].T @ rs[:, k + 1])

        return ds, rs, rs_norm_sq, p, eta, k + 1

    def _cond_fun(values):
        _ds, _, rs_norm_sq, _, _eta, k = values
        return (rs_norm_sq[k] > sqtol) & (k < max_iter)

    # Initialization
    b /= jnp.linalg.norm(b, 2)
    ds = jnp.zeros((b.size, max_iter))
    rs = jnp.zeros((b.size, max_iter + 1))
    rs_norm_sq = jnp.ones_like(rs, shape=max_iter + 1)

    # Initialize loop variables
    sqtol = tol**2
    min_eta = min_eta
    eta = jnp.inf
    rs = rs.at[:, 0].set(b)
    p = b if overwrite_b else b.copy()

    # Lanczos iterations
    ds, _, _, _, _, k = jax.lax.while_loop(_cond_fun, _step, (ds, rs, rs_norm_sq, p, eta, 0))

    return ds[:, :k]

lanczos_tridiag(A: LinearOperatorLike, v0: ArrayLike, num_iters: int, reortho: bool = True) -> tuple[jax.Array, jax.Array, jax.Array]

Lanczos tridiagonalization for symmetric operators.

Computes a tridiagonal reduction of a symmetric linear operator using the Lanczos algorithm. Returns the orthonormal Lanczos vectors Q and the tridiagonal matrix T such that A ≈ Q T Q^T on the Krylov subspace.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	Symmetric linear operator or matrix.	required
v0	ArrayLike	Initial vector for the Krylov process.	required
num_iters	int	Number of Lanczos iterations.	required
reortho	bool	Whether to perform full reorthogonalization. Default is True.	True

Returns:

Name	Type	Description
Q	(Array, shape(n, num_iters))	Orthonormal Lanczos vectors.
alpha	(Array, shape(num_iters))	Diagonal elements.
beta	(Array, shape(num_iters - 1))	Off-diagonal elements.

Source code in linox/linalg/approx/lanczos.py

def lanczos_tridiag(
    A: LinearOperatorLike,
    v0: ArrayLike,
    num_iters: int,
    reortho: bool = True,
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Lanczos tridiagonalization for symmetric operators.

    Computes a tridiagonal reduction of a symmetric linear operator using the
    Lanczos algorithm. Returns the orthonormal Lanczos vectors Q and the
    tridiagonal matrix T such that A ≈ Q T Q^T on the Krylov subspace.

    Parameters
    ----------
    A : LinearOperatorLike
        Symmetric linear operator or matrix.
    v0 : ArrayLike
        Initial vector for the Krylov process.
    num_iters : int
        Number of Lanczos iterations.
    reortho : bool, optional
        Whether to perform full reorthogonalization. Default is True.

    Returns
    -------
    Q : jax.Array, shape (n, num_iters)
        Orthonormal Lanczos vectors.
    alpha : jax.Array, shape (num_iters,)
        Diagonal elements.
    beta : jax.Array, shape (num_iters-1,)
        Off-diagonal elements.
    """
    v0 = jnp.asarray(v0)
    n = v0.size

    # Normalize initial vector
    beta_0 = jnp.linalg.norm(v0)
    v0 /= beta_0

    # Pre-allocate arrays
    Q = jnp.zeros((n, num_iters))
    alpha = jnp.zeros(num_iters)
    beta = jnp.zeros(num_iters - 1) if num_iters > 1 else jnp.zeros(0)

    # Initialize first vector
    Q = Q.at[:, 0].set(v0)

    def lanczos_step(k, carry):
        Q_curr, alpha_curr, beta_curr = carry

        # Matrix-vector product
        v = Q_curr[:, k]
        w = A @ v

        # Compute diagonal element
        alpha_k = jnp.dot(w, v)
        alpha_curr = alpha_curr.at[k].set(alpha_k)

        # Update w (three-term recurrence)
        w -= alpha_k * v
        # Subtract previous vector (if k > 0)
        prev_contrib = lax.cond(
            k > 0,
            lambda: beta_curr[k - 1] * Q_curr[:, k - 1],
            lambda: jnp.zeros_like(w),
        )
        w -= prev_contrib

        # Reorthogonalization (full Gram-Schmidt)
        if reortho:

            def reorth_body(j, w_state):
                proj = jnp.dot(w_state, Q_curr[:, j])
                return w_state - proj * Q_curr[:, j]

            w = lax.cond(
                k > 0,
                lambda w_val: lax.fori_loop(0, k, reorth_body, w_val),
                lambda w_val: w_val,
                w,
            )

        # Compute off-diagonal element
        beta_k = jnp.linalg.norm(w)

        # Store next vector
        Q_next = lax.cond(
            k < num_iters - 1,
            lambda: Q_curr.at[:, k + 1].set(w / (beta_k + 1e-16)),
            lambda: Q_curr,
        )

        # Store beta
        beta_next = lax.cond(
            k < num_iters - 1,
            lambda: beta_curr.at[k].set(beta_k),
            lambda: beta_curr,
        )

        return (Q_next, alpha_curr, beta_next)

    # Run Lanczos iterations
    Q, alpha, beta = lax.fori_loop(0, num_iters, lanczos_step, (Q, alpha, beta))

    return Q, alpha, beta

Arnoldi and general Krylov methods.

This module implements Arnoldi iteration for general matrices. For symmetric matrices, use lanczos.py.

arnoldi_iteration(A: LinearOperatorLike, v0: ArrayLike, num_iters: int) -> tuple[jax.Array, jax.Array]

Arnoldi iteration for general (non-symmetric) operators.

Computes a Hessenberg reduction of a general linear operator using the Arnoldi iteration. Returns the orthonormal Arnoldi vectors Q and the upper Hessenberg matrix H such that A ≈ Q H Q^T on the Krylov subspace.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	General linear operator or matrix.	required
v0	ArrayLike	Initial vector for the Krylov process.	required
num_iters	int	Number of Arnoldi iterations.	required

Returns:

Name	Type	Description
Q	(Array, shape(n, num_iters))	Orthonormal Arnoldi vectors.
H	(Array, shape(num_iters + 1, num_iters))	Upper Hessenberg matrix.

Source code in linox/linalg/approx/arnoldi.py

def arnoldi_iteration(
    A: LinearOperatorLike,
    v0: ArrayLike,
    num_iters: int,
) -> tuple[jax.Array, jax.Array]:
    """Arnoldi iteration for general (non-symmetric) operators.

    Computes a Hessenberg reduction of a general linear operator using the
    Arnoldi iteration. Returns the orthonormal Arnoldi vectors Q and the
    upper Hessenberg matrix H such that A ≈ Q H Q^T on the Krylov subspace.

    Parameters
    ----------
    A : LinearOperatorLike
        General linear operator or matrix.
    v0 : ArrayLike
        Initial vector for the Krylov process.
    num_iters : int
        Number of Arnoldi iterations.

    Returns
    -------
    Q : jax.Array, shape (n, num_iters)
        Orthonormal Arnoldi vectors.
    H : jax.Array, shape (num_iters+1, num_iters)
        Upper Hessenberg matrix.
    """
    v0 = jnp.asarray(v0)
    n = v0.size

    # Normalize initial vector
    beta_0 = jnp.linalg.norm(v0)
    v0 /= beta_0

    # Pre-allocate arrays
    Q = jnp.zeros((n, num_iters))
    H = jnp.zeros((num_iters + 1, num_iters))

    # Initialize first vector
    Q = Q.at[:, 0].set(v0)

    def arnoldi_step(k, carry):
        Q_curr, H_curr = carry

        # Matrix-vector product
        v = Q_curr[:, k]
        w = A @ v

        # Full Gram-Schmidt orthogonalization using fori_loop
        def gs_body(j, state):
            w_state, H_state = state
            h_jk = jnp.dot(w_state, Q_curr[:, j])
            H_state = H_state.at[j, k].set(h_jk)
            w_state -= h_jk * Q_curr[:, j]
            return (w_state, H_state)

        w, H_curr = lax.fori_loop(0, k + 1, gs_body, (w, H_curr))

        # Compute residual norm
        h_next = jnp.linalg.norm(w)
        H_curr = H_curr.at[k + 1, k].set(h_next)

        # Store next vector (if not last iteration)
        Q_next = lax.cond(
            k < num_iters - 1,
            lambda: Q_curr.at[:, k + 1].set(w / (h_next + 1e-16)),
            lambda: Q_curr,
        )

        return (Q_next, H_curr)

    # Run Arnoldi iterations using fori_loop
    Q, H = lax.fori_loop(0, num_iters, arnoldi_step, (Q, H))

    return Q, H

arnoldi_matrix_function(A: LinearOperatorLike, v: ArrayLike, func: callable, num_iters: int) -> jax.Array

Approximate f(A)v using Arnoldi iteration.

Source code in linox/linalg/approx/arnoldi.py

def arnoldi_matrix_function(
    A: LinearOperatorLike,
    v: ArrayLike,
    func: callable,
    num_iters: int,
) -> jax.Array:
    """Approximate f(A)v using Arnoldi iteration."""
    v = jnp.asarray(v)
    v_norm = jnp.linalg.norm(v)
    v_normalized = v / v_norm

    # Perform Arnoldi iteration
    Q, H = arnoldi_iteration(A, v_normalized, num_iters)

    # H is (num_iters+1, num_iters), discard last row for square matrix approximation
    H_square = H[:-1, :]

    # Apply function to Hessenberg matrix
    w, V = jnp.linalg.eig(H_square)
    fH = V @ jnp.diag(func(w)) @ jnp.linalg.inv(V)

    # Extract first row
    e1 = jnp.zeros(num_iters)
    e1 = e1.at[0].set(1.0)

    # Project back
    result = v_norm * (Q @ (fH @ e1))

    return result.real

Hutchinson's stochastic trace estimator.

Approximate trace and diagonal using stochastic probes (Hutchinson's method). Optimized for batched execution via matrix-matrix multiplication (A @ Z).

hutchinson_diagonal(A: LinearOperatorLike, key: jax.Array, num_samples: int = 100, distribution: str = 'rademacher') -> tuple[jax.Array, jax.Array]

Estimate diagonal of a linear operator using Hutchinson's method.

Computes: diag(A) ≈ (1/M) * sum(v_i ⊙ (A v_i))

Args: A: LinearOperator or array (n, n) key: PRNG key num_samples: Number of probes distribution: 'rademacher' (default) or 'normal'

Returns:

Type	Description
(estimate, std_error) each of shape (n,)

Source code in linox/linalg/approx/hutchinson.py

def hutchinson_diagonal(
    A: LinearOperatorLike,
    key: jax.Array,
    num_samples: int = 100,
    distribution: str = "rademacher",
) -> tuple[jax.Array, jax.Array]:
    """Estimate diagonal of a linear operator using Hutchinson's method.

    Computes: diag(A) ≈ (1/M) * sum(v_i ⊙ (A v_i))

    Args:
        A: LinearOperator or array (n, n)
        key: PRNG key
        num_samples: Number of probes
        distribution: 'rademacher' (default) or 'normal'

    Returns
    -------
        (estimate, std_error) each of shape (n,)
    """
    n = A.shape[-1]
    Z = _generate_probes(key, n, num_samples, distribution)  # (n, num_samples)

    AZ = A @ Z

    samples = Z * AZ  # (n, samples)

    # Mean over samples axis=1
    mean = jnp.mean(samples, axis=1)
    std = jnp.std(samples, axis=1, ddof=1) / jnp.sqrt(num_samples)

    return mean, std

hutchinson_trace(A: LinearOperatorLike, key: jax.Array, num_samples: int = 100, distribution: str = 'rademacher') -> tuple[jax.Array, jax.Array]

Estimate trace of a linear operator using Hutchinson's method.

Computes Monte Carlo estimate: trace(A) ≈ (1/M) * sum(v_i^T A v_i) Uses batched execution (A @ Z) for efficiency.

Args: A: LinearOperator or array (n, n) key: PRNG key num_samples: Number of random probes distribution: 'rademacher' (default) or 'normal'

Returns:

Type	Description
(estimate, std_error)

Source code in linox/linalg/approx/hutchinson.py

def hutchinson_trace(
    A: LinearOperatorLike,
    key: jax.Array,
    num_samples: int = 100,
    distribution: str = "rademacher",
) -> tuple[jax.Array, jax.Array]:
    """Estimate trace of a linear operator using Hutchinson's method.

    Computes Monte Carlo estimate: trace(A) ≈ (1/M) * sum(v_i^T A v_i)
    Uses batched execution (A @ Z) for efficiency.

    Args:
        A: LinearOperator or array (n, n)
        key: PRNG key
        num_samples: Number of random probes
        distribution: 'rademacher' (default) or 'normal'

    Returns
    -------
        (estimate, std_error)
    """
    n = A.shape[-1]
    Z = _generate_probes(key, n, num_samples, distribution)  # (n, num_samples)

    # Batched matmul: (n, n) @ (n, samples) -> (n, samples)
    AZ = A @ Z

    # v^T A v = sum(v * Av) for each column
    samples = jnp.sum(Z * AZ, axis=0)  # (samples,)

    return jnp.mean(samples), jnp.std(samples, ddof=1) / jnp.sqrt(num_samples)

hutchinson_trace_and_diagonal(A: LinearOperatorLike, key: jax.Array, num_samples: int = 100, distribution: str = 'rademacher') -> dict[str, tuple[jax.Array, jax.Array]]

Jointly estimate trace and diagonal using shared probes.

Source code in linox/linalg/approx/hutchinson.py

def hutchinson_trace_and_diagonal(
    A: LinearOperatorLike,
    key: jax.Array,
    num_samples: int = 100,
    distribution: str = "rademacher",
) -> dict[str, tuple[jax.Array, jax.Array]]:
    """Jointly estimate trace and diagonal using shared probes."""
    n = A.shape[-1]
    Z = _generate_probes(key, n, num_samples, distribution)

    AZ = A @ Z

    # Diagonal samples: (n, samples)
    diag_samples = Z * AZ

    # Trace samples: sum over n -> (samples,)
    trace_samples = jnp.sum(diag_samples, axis=0)

    trace_mean = jnp.mean(trace_samples)
    trace_std = jnp.std(trace_samples, ddof=1) / jnp.sqrt(num_samples)

    diag_mean = jnp.mean(diag_samples, axis=1)
    diag_std = jnp.std(diag_samples, axis=1, ddof=1) / jnp.sqrt(num_samples)

    return {"trace": (trace_mean, trace_std), "diagonal": (diag_mean, diag_std)}

Stochastic Lanczos Quadrature (SLQ).

Approximates trace(f(A)) using stochastic probes and Lanczos tridiagonalization.

slq(A: LinearOperatorLike, func_scalar: callable, key: jax.Array, num_samples: int = 10, m: int = 20, distribution: str = 'rademacher') -> tuple[jax.Array, jax.Array]

Estimate trace(f(A)) using SLQ.

trace(f(A)) ≈ (1/M) * sum_i v_i^T f(A) v_i v_i^T f(A) v_i ≈ ||v_i||^2 * e_1^T f(T_m) e_1

Source code in linox/linalg/approx/slq.py

def slq(
    A: LinearOperatorLike,
    func_scalar: callable,  # f: scalar -> scalar (e.g. jnp.log)
    key: jax.Array,
    num_samples: int = 10,
    m: int = 20,  # Krylov iterations
    distribution: str = "rademacher",
) -> tuple[jax.Array, jax.Array]:
    """Estimate trace(f(A)) using SLQ.

    trace(f(A)) ≈ (1/M) * sum_i v_i^T f(A) v_i
    v_i^T f(A) v_i ≈ ||v_i||^2 * e_1^T f(T_m) e_1
    """
    n = A.shape[-1]

    # 1. Generate probes
    # Batched generation
    if distribution == "rademacher":
        keys = random.split(key, num_samples)
        V = jax.vmap(lambda k: 2 * random.bernoulli(k, shape=(n,)) - 1.0)(keys)
    elif distribution == "normal":
        keys = random.split(key, num_samples)
        V = jax.vmap(lambda k: random.normal(k, shape=(n,)))(keys)
    else:
        msg = f"Unknown distribution: {distribution}"
        raise ValueError(msg)

    # 2. Lanczos on each probe
    # Note: Lanczos is inherently sequential per vector, so we vmap the entire lanczos process.

    def process_probe(v):
        # f(A)v approx
        w = lanczos_matrix_function(A, v, func_scalar, m)
        # v^T f(A) v
        return jnp.dot(v, w)

    # vmap over samples
    estimates = jax.vmap(process_probe)(V)

    return jnp.mean(estimates), jnp.std(estimates, ddof=1) / jnp.sqrt(num_samples)

slq_logdet(A: LinearOperatorLike, key: jax.Array, num_samples: int = 10, m: int = 20) -> tuple[jax.Array, jax.Array]

Estimate log-determinant using SLQ.

logdet(A) = trace(log(A))

Source code in linox/linalg/approx/slq.py

def slq_logdet(
    A: LinearOperatorLike,
    key: jax.Array,
    num_samples: int = 10,
    m: int = 20,
) -> tuple[jax.Array, jax.Array]:
    """Estimate log-determinant using SLQ.

    logdet(A) = trace(log(A))
    """
    return slq(A, jnp.log, key, num_samples, m)

LSMR iterative solver for least-squares problems.

This module implements the LSMR (Least Squares Minimal Residual) algorithm for solving large-scale least-squares and linear systems in a matrix-free manner.

The implementation closely follows the matfree library (https://github.com/pnkraemer/matfree) by Nicholas Krämer et al., which itself follows the original LSMR algorithm by Fong and Saunders.

Key features: - Matrix-free: Only requires matrix-vector products - Handles over-determined, under-determined, and rank-deficient systems - JAX-compatible with automatic differentiation support - Iterative with configurable stopping criteria

References

.. [1] D. C.-L. Fong and M. A. Saunders, "LSMR: An iterative algorithm for sparse least-squares problems," SIAM Journal on Scientific Computing, vol. 33, no. 5, pp. 2950-2971, 2011.

.. [2] matfree: Matrix-free linear algebra in JAX https://github.com/pnkraemer/matfree

.. [3] A. Roy, N. Krämer, V. De Bortoli, and A. Doucet, "Gradients of Stochastic Trace Estimators via Differentiable Matrix-Free Linear Solvers," arXiv preprint, 2025.

lsmr_solve(A: LinearOperatorLike, b: ArrayLike, atol: float = 1e-06, btol: float = 1e-06, ctol: float = 1e-08, maxiter: int | None = None, damp: float = 0.0, x0: ArrayLike | None = None) -> tuple[jax.Array, dict]

Solve least-squares problem min ||Ax - b||_2 using LSMR algorithm.

LSMR is an iterative method for solving: - Linear systems Ax = b (when A is square and full rank) - Least-squares min ||Ax - b||_2 (when A is over-determined) - Minimum-norm solutions (when A is under-determined or rank-deficient)

The algorithm only requires matrix-vector products with A and A^T, making it suitable for large sparse or structured matrices.

This implementation follows matfree's approach, which matches SciPy's LSMR.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	Linear operator or matrix of shape (m, n). Should support both A @ v and A.T @ v operations.	required
b	ArrayLike	Right-hand side vector of shape (m,).	required
atol	float	Absolute tolerance for convergence. Default is 1e-6.	1e-06
btol	float	Relative tolerance for convergence. Default is 1e-6.	1e-06
ctol	float	Condition number tolerance. Default is 1e-8.	1e-08
maxiter	int	Maximum number of iterations. If None, uses min(m, n). Default is None.	None
damp	float	Damping parameter for regularization. Solves the problem min ||[A; damp*I] x - [b; 0]||_2 instead. Default is 0.0 (no damping).	0.0
x0	ArrayLike	Initial guess for the solution. If None, starts with zero vector. Default is None.	None

Returns:

Name	Type	Description
x	(Array, shape(n))	Solution vector.
info	dict	Dictionary containing: - 'istop': Stopping condition (0-9) - 'itn': Number of iterations performed - 'normr': Final residual norm ||r|| - 'normar': Final ||A^T r|| - 'normA': Estimate of ||A|| - 'condA': Estimate of cond(A) - 'normx': Norm of solution ||x||

Examples:

>>> import jax.numpy as jnp
>>> from linox import Matrix
>>> A = Matrix(jnp.eye(50))
>>> b = jnp.ones(50)
>>> x, info = lsmr_solve(A, b)
>>> assert jnp.allclose(x, b)

Notes

This implementation closely follows matfree's LSMR, which is based on Fong and Saunders (2011) and matches SciPy's implementation.

References

.. [1] D. C.-L. Fong and M. A. Saunders, "LSMR: An iterative algorithm for sparse least-squares problems," SIAM J. Sci. Compute., 2011. .. [2] https://github.com/pnkraemer/matfree

Source code in linox/linalg/approx/lsmr.py

def lsmr_solve(
    A: LinearOperatorLike,
    b: ArrayLike,
    atol: float = 1e-6,
    btol: float = 1e-6,
    ctol: float = 1e-8,
    maxiter: int | None = None,
    damp: float = 0.0,
    x0: ArrayLike | None = None,
) -> tuple[jax.Array, dict]:
    """Solve least-squares problem min ||Ax - b||_2 using LSMR algorithm.

    LSMR is an iterative method for solving:
    - Linear systems Ax = b (when A is square and full rank)
    - Least-squares min ||Ax - b||_2 (when A is over-determined)
    - Minimum-norm solutions (when A is under-determined or rank-deficient)

    The algorithm only requires matrix-vector products with A and A^T, making
    it suitable for large sparse or structured matrices.

    This implementation follows matfree's approach, which matches SciPy's LSMR.

    Parameters
    ----------
    A : LinearOperatorLike
        Linear operator or matrix of shape (m, n). Should support both
        `A @ v` and `A.T @ v` operations.
    b : ArrayLike
        Right-hand side vector of shape (m,).
    atol : float, optional
        Absolute tolerance for convergence. Default is 1e-6.
    btol : float, optional
        Relative tolerance for convergence. Default is 1e-6.
    ctol : float, optional
        Condition number tolerance. Default is 1e-8.
    maxiter : int, optional
        Maximum number of iterations. If None, uses min(m, n).
        Default is None.
    damp : float, optional
        Damping parameter for regularization. Solves the problem
        min ||[A; damp*I] x - [b; 0]||_2 instead. Default is 0.0 (no damping).
    x0 : ArrayLike, optional
        Initial guess for the solution. If None, starts with zero vector.
        Default is None.

    Returns
    -------
    x : jax.Array, shape (n,)
        Solution vector.
    info : dict
        Dictionary containing:
        - 'istop': Stopping condition (0-9)
        - 'itn': Number of iterations performed
        - 'normr': Final residual norm ||r||
        - 'normar': Final ||A^T r||
        - 'normA': Estimate of ||A||
        - 'condA': Estimate of cond(A)
        - 'normx': Norm of solution ||x||

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from linox import Matrix
    >>> A = Matrix(jnp.eye(50))
    >>> b = jnp.ones(50)
    >>> x, info = lsmr_solve(A, b)
    >>> assert jnp.allclose(x, b)

    Notes
    -----
    This implementation closely follows matfree's LSMR, which is based on
    Fong and Saunders (2011) and matches SciPy's implementation.

    References
    ----------
    .. [1] D. C.-L. Fong and M. A. Saunders, "LSMR: An iterative algorithm for
           sparse least-squares problems," SIAM J. Sci. Compute., 2011.
    .. [2] https://github.com/pnkraemer/matfree
    """
    b = jnp.asarray(b)
    m = b.shape[0]
    n = A.shape[1]

    if maxiter is None:
        maxiter = min(m, n)

    # Initialize x
    if x0 is None:
        # Match x shape to b's shape (handling (m, 1) case from base.py)
        x_shape = (n, *b.shape[1:]) if b.ndim > 1 else (n,)
        x = jnp.zeros(x_shape, dtype=b.dtype)
        Ax = jnp.zeros_like(b)
        u = b.copy()
        beta = jnp.linalg.norm(u)
    else:
        x = jnp.asarray(x0)
        Ax = A @ x
        u = b - Ax
        beta = jnp.linalg.norm(u)

    normb = jnp.linalg.norm(b)

    # Normalize u
    u /= jnp.where(beta > 0, beta, 1.0)

    # Apply A^T to u
    v = A.T @ u
    alpha = jnp.linalg.norm(v)
    v /= jnp.where(alpha > 0, alpha, 1.0)
    v = jnp.where(beta == 0, jnp.zeros_like(v), v)
    alpha = jnp.where(beta == 0, jnp.zeros_like(alpha), alpha)

    # Initialize variables for 1st iteration
    zetabar = alpha * beta
    alphabar = alpha
    rho = 1.0
    rhobar = 1.0
    cbar = 1.0
    sbar = 0.0

    h = v.copy()
    hbar = jnp.zeros_like(x)

    # Initialize variables for estimation of ||r||
    betadd = beta
    betad = 0.0
    rhodold = 1.0
    tautildeold = 0.0
    thetatilde = 0.0
    zeta = 0.0
    d = 0.0

    # Initialize variables for estimation of ||A|| and cond(A)
    normA2 = alpha * alpha
    maxrbar = 0.0
    minrbar = 1e10
    normA = jnp.sqrt(normA2)
    condA = 1.0
    normx = 0.0

    # Items for use in stopping rules
    normr = beta
    normar = alpha * beta

    # Iteration loop using lax.while_loop
    def cond_fun(carry):
        itn, istop, *_ = carry
        return (istop == 0) & (itn < maxiter)

    def body_fun(carry):
        (
            itn,
            istop,
            u,
            v,
            alpha,
            beta,
            alphabar,
            rhobar,
            rho,
            zeta,
            sbar,
            cbar,
            zetabar,
            hbar,
            h,
            x,
            betadd,
            thetatilde,
            rhodold,
            betad,
            tautildeold,
            d,
            normA2,
            maxrbar,
            minrbar,
            normar,
            normr,
            normA,
            condA,
            normx,
        ) = carry

        # Perform the next step of the bidiagonalization
        Av = A @ v
        u_new = Av - alpha * u
        beta_new = jnp.linalg.norm(u_new)

        # Stable update using jax.lax.cond to handle beta_new near 0
        def update_v_alpha(args):
            u_val, beta_val = args
            # Normalization
            u_normalized = u_val / beta_val

            # Update v
            v_new = A.T @ u_normalized - beta_val * v
            alpha_new = jnp.linalg.norm(v_new)

            # Normalize v if alpha > 0
            v_normalized = v_new / jnp.where(alpha_new > 0, alpha_new, 1.0)

            return u_normalized, v_normalized, alpha_new

        def no_update(args):
            u_val, _ = args
            # When beta is 0, we don't update v or alpha
            # u remains u_new (which is 0 vector if beta is 0, effectively)
            return u_val, v, alpha

        # Conditional execution
        u, v, alpha = lax.cond(beta_new > 0, update_v_alpha, no_update, (u_new, beta_new))
        beta = beta_new

        # Construct rotation Qhat_{k,2k+1}
        chat, shat, alphahat = _sym_ortho(alphabar, damp)

        # Use a plane rotation (Q_i) to turn B_i to R_i
        rhoold = rho
        c, s, rho = _sym_ortho(alphahat, beta)
        thetanew = s * alpha
        alphabar = c * alpha

        # Use a plane rotation (Qbar_i) to turn R_i^T to R_i^bar
        rhobarold = rhobar
        zetaold = zeta
        thetabar = sbar * rho
        rhotemp = cbar * rho
        cbar, sbar, rhobar = _sym_ortho(rhotemp, thetanew)
        zeta = cbar * zetabar
        zetabar = -sbar * zetabar

        # Update h, h_hat, x
        hbar = h - hbar * (thetabar * rho / (rhoold * rhobarold))
        x += (zeta / (rho * rhobar)) * hbar
        h = v - h * (thetanew / rho)

        # Estimate of ||r||
        # Apply rotation Qhat_{k,2k+1}
        betaacute = chat * betadd
        betacheck = -shat * betadd

        # Apply rotation Q_{k,k+1}
        betahat = c * betaacute
        betadd = -s * betaacute

        # Apply rotation Qtilde_{k-1}
        thetatildeold = thetatilde
        ctildeold, stildeold, rhotildeold = _sym_ortho(rhodold, thetabar)
        thetatilde = stildeold * rhobar
        rhodold = ctildeold * rhobar
        betad = -stildeold * betad + ctildeold * betahat

        tautildeold = (zetaold - thetatildeold * tautildeold) / rhotildeold
        taud = (zeta - thetatilde * tautildeold) / rhodold
        d += betacheck * betacheck
        normr = jnp.sqrt(d + (betad - taud) ** 2 + betadd * betadd)

        # Estimate ||A||
        normA2 += beta * beta
        normA = jnp.sqrt(normA2)
        normA2 += alpha * alpha

        # Estimate cond(A)
        maxrbar = jnp.maximum(maxrbar, rhobarold)
        minrbar = jnp.where(itn > 1, jnp.minimum(minrbar, rhobarold), minrbar)
        condA = jnp.maximum(maxrbar, rhotemp) / jnp.minimum(minrbar, rhotemp)

        # Compute norms for convergence testing
        normar = jnp.abs(zetabar)
        normx = jnp.linalg.norm(x)

        # Check whether we should stop
        itn += 1
        test1 = normr / normb
        z = normA * normr
        z_safe = jnp.where(z != 0, z, 1.0)
        test2 = jnp.where(z != 0, normar / z_safe, _LARGE_VALUE)
        test3 = 1.0 / condA
        t1 = test1 / (1 + normA * normx / normb)
        rtol = btol + atol * normA * normx / normb

        # Determine stopping condition (following matfree/scipy order)
        istop = 0
        istop = jnp.where(normar == 0, 9, istop)
        istop = jnp.where(normb == 0, 8, istop)
        istop = jnp.where(itn >= maxiter, 7, istop)
        istop = jnp.where(1 + test3 <= 1, 6, istop)
        istop = jnp.where(1 + test2 <= 1, 5, istop)
        istop = jnp.where(1 + t1 <= 1, 4, istop)
        istop = jnp.where(test3 <= ctol, 3, istop)
        istop = jnp.where(test2 <= atol, 2, istop)
        istop = jnp.where(test1 <= rtol, 1, istop)

        return (
            itn,
            istop,
            u,
            v,
            alpha,
            beta,
            alphabar,
            rhobar,
            rho,
            zeta,
            sbar,
            cbar,
            zetabar,
            hbar,
            h,
            x,
            betadd,
            thetatilde,
            rhodold,
            betad,
            tautildeold,
            d,
            normA2,
            maxrbar,
            minrbar,
            normar,
            normr,
            normA,
            condA,
            normx,
        )

    init_carry = (
        0,  # itn
        0,  # istop
        u,
        v,
        alpha,
        beta,
        alphabar,
        rhobar,
        rho,
        zeta,
        sbar,
        cbar,
        zetabar,
        hbar,
        h,
        x,
        betadd,
        thetatilde,
        rhodold,
        betad,
        tautildeold,
        d,
        normA2,
        maxrbar,
        minrbar,
        normar,
        normr,
        normA,
        condA,
        normx,
    )

    final_carry = lax.while_loop(cond_fun, body_fun, init_carry)

    # Extract final values
    (
        itn,
        istop,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        x,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        _,
        normar,
        normr,
        normA,
        condA,
        normx,
    ) = final_carry

    info = {
        "istop": istop,
        "itn": itn,
        "normr": normr,
        "normar": normar,
        "normA": normA,
        "condA": condA,
        "normx": normx,
    }

    return x, info

Matrix-free SVD via Lanczos bidiagonalization.

This module implements matrix-free singular value decomposition (SVD) using Lanczos bidiagonalization for large-scale problems where only a few singular values and vectors are needed.

The implementation is inspired by the matfree library (https://github.com/pnkraemer/matfree) by Nicholas Krämer et al.

Key algorithms: - Lanczos bidiagonalization: Reduces matrix to bidiagonal form - Partial SVD: Computes k largest singular values/vectors

References

.. [1] N. Krämer, M. Schober, and P. Hennig, "Gradients of functions of large matrices," arXiv preprint arXiv:2405.17277, 2024. https://arxiv.org/abs/2405.17277

.. [2] matfree: Matrix-free linear algebra in JAX https://github.com/pnkraemer/matfree

.. [3] G. H. Golub and C. F. Van Loan, "Matrix Computations," 4th ed., Johns Hopkins, 2013.

svd_partial(A: LinearOperatorLike, k: int, num_iters: int | None = None, u0: ArrayLike | None = None, reortho: bool = True) -> tuple[jax.Array, jax.Array, jax.Array]

Compute partial SVD using Lanczos bidiagonalization.

Computes the k largest singular values and corresponding singular vectors of a large matrix without forming the full matrix explicitly.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	Linear operator or matrix of shape (m, n).	required
k	int	Number of singular values/vectors to compute.	required
num_iters	int	Number of Lanczos iterations. Should be larger than k for good approximation. If None, uses min(2*k, min(m, n)). Default is None.	None
u0	ArrayLike	Initial vector of shape (m,) for bidiagonalization. If None, uses vector of ones. Default is None.	None
reortho	bool	Whether to use full reorthogonalization in bidiagonalization. This significantly improves numerical stability. Default is True.	True

Returns:

Name	Type	Description
U	(Array, shape(m, k))	Left singular vectors (columns).
S	(Array, shape(k))	Singular values in descending order.
Vt	(Array, shape(k, n))	Right singular vectors (rows).

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from linox import Matrix
>>> # Large matrix
>>> key = jax.random.PRNGKey(0)
>>> A_dense = jax.random.normal(key, (1000, 500))
>>> A = Matrix(A_dense)
>>> # Compute top 10 singular values/vectors
>>> U, S, Vt = svd_partial(A, k=10)
>>> print(f"Top 10 singular values: {S}")
>>> # Verify: A ≈ U @ diag(S) @ Vt
>>> A_approx = U @ jnp.diag(S) @ Vt
>>> error = jnp.linalg.norm(A_dense - A_approx)

Notes

This is a matrix-free alternative to jnp.linalg.svd for computing a few singular values/vectors of large sparse or structured matrices.

The algorithm: 1. Performs Lanczos bidiagonalization: A ≈ U_bi B V_bi^T 2. Computes SVD of small bidiagonal matrix B 3. Projects back to get singular vectors of A

References

Inspired by matfree library [1, 2].

Source code in linox/linalg/spectral.py

def svd_partial(
    A: LinearOperatorLike,
    k: int,
    num_iters: int | None = None,
    u0: ArrayLike | None = None,
    reortho: bool = True,
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """Compute partial SVD using Lanczos bidiagonalization.

    Computes the k largest singular values and corresponding singular vectors
    of a large matrix without forming the full matrix explicitly.

    Parameters
    ----------
    A : LinearOperatorLike
        Linear operator or matrix of shape (m, n).
    k : int
        Number of singular values/vectors to compute.
    num_iters : int, optional
        Number of Lanczos iterations. Should be larger than k for good
        approximation. If None, uses min(2*k, min(m, n)). Default is None.
    u0 : ArrayLike, optional
        Initial vector of shape (m,) for bidiagonalization. If None,
        uses vector of ones. Default is None.
    reortho : bool, optional
        Whether to use full reorthogonalization in bidiagonalization.
        This significantly improves numerical stability. Default is True.

    Returns
    -------
    U : jax.Array, shape (m, k)
        Left singular vectors (columns).
    S : jax.Array, shape (k,)
        Singular values in descending order.
    Vt : jax.Array, shape (k, n)
        Right singular vectors (rows).

    Examples
    --------
    >>> import jax
    >>> import jax.numpy as jnp
    >>> from linox import Matrix
    >>> # Large matrix
    >>> key = jax.random.PRNGKey(0)
    >>> A_dense = jax.random.normal(key, (1000, 500))
    >>> A = Matrix(A_dense)
    >>> # Compute top 10 singular values/vectors
    >>> U, S, Vt = svd_partial(A, k=10)
    >>> print(f"Top 10 singular values: {S}")
    >>> # Verify: A ≈ U @ diag(S) @ Vt
    >>> A_approx = U @ jnp.diag(S) @ Vt
    >>> error = jnp.linalg.norm(A_dense - A_approx)

    Notes
    -----
    This is a matrix-free alternative to jnp.linalg.svd for computing
    a few singular values/vectors of large sparse or structured matrices.

    The algorithm:
    1. Performs Lanczos bidiagonalization: A ≈ U_bi B V_bi^T
    2. Computes SVD of small bidiagonal matrix B
    3. Projects back to get singular vectors of A

    References
    ----------
    Inspired by matfree library [1, 2].
    """
    m = A.shape[0]
    n = A.shape[1]

    if num_iters is None:
        num_iters = min(2 * k, m, n)

    if u0 is None:
        u0 = jnp.ones(m)

    # Perform bidiagonalization
    U_bi, V_bi, alpha, beta = lanczos_bidiag(A, u0, num_iters, reortho=reortho)

    # Construct bidiagonal matrix
    B = jnp.diag(alpha)
    if beta.size > 0:
        B += jnp.diag(beta, k=1)

    # Compute SVD of small bidiagonal matrix
    U_small, S_small, Vt_small = jnp.linalg.svd(B, full_matrices=False)

    # Select top k singular values/vectors
    U_small = U_small[:, :k]
    S = S_small[:k]
    Vt_small = Vt_small[:k, :]

    # Project back to original space.
    #
    # The Golub-Kahan recurrence produces `A^T U_bi = V_bi B`, i.e.
    # `A ~= U_bi @ B.T @ V_bi.T`. Substituting `B = U_small S Vt_small` gives
    # `A ~= (U_bi @ Vt_small.T) S (U_small.T @ V_bi.T)`, so the left/right
    # factors of the small SVD attach to the *opposite* Krylov basis.
    U = U_bi @ Vt_small.T
    Vt = (V_bi @ U_small).T

    return U, S, Vt

lanczos_bidiag(A: LinearOperatorLike, u0: ArrayLike, num_iters: int, reortho: bool = True) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]

Lanczos bidiagonalization for SVD computation.

Reduces a matrix A to bidiagonal form: A ≈ U B V^T, where U and V are orthonormal and B is bidiagonal.

This is the foundation for computing partial SVD of large matrices where only a few singular values/vectors are needed.

Parameters:

Name	Type	Description	Default
A	LinearOperatorLike	Linear operator or matrix of shape (m, n). Should support both A @ v and A.T @ u operations.	required
u0	ArrayLike	Initial vector of shape (m,) for the bidiagonalization process. Will be normalized internally.	required
num_iters	int	Number of bidiagonalization iterations (size of bidiagonal matrix). Should be much smaller than min(m, n).	required
reortho	bool	Whether to use full reorthogonalization. This significantly improves numerical stability at the cost of O(num_iters^2) operations. Default is True (recommended).	True

Returns:

Name	Type	Description
U	(Array, shape(m, num_iters))	Left orthonormal basis vectors (columns).
V	(Array, shape(n, num_iters))	Right orthonormal basis vectors (columns).
alpha	(Array, shape(num_iters))	Diagonal elements of bidiagonal matrix B.
beta	(Array, shape(num_iters - 1))	Super-diagonal elements of bidiagonal matrix B.

Notes

The bidiagonal matrix B has the form: B = [[alpha[0], beta[0], 0, ...], [0, alpha[1], beta[1], ...], [0, 0, alpha[2], ...], [...]]

This is related to Golub-Kahan bidiagonalization and is used in algorithms like LSMR and partial SVD computation.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from linox import Matrix
>>> key = jax.random.PRNGKey(0)
>>> A = Matrix(jax.random.normal(key, (100, 50)))
>>> u0 = jnp.ones(100)
>>> U, V, alpha, beta = lanczos_bidiag(A, u0, num_iters=10)
>>> # U and V contain orthonormal vectors
>>> # B = diag(alpha) + diag(beta, 1) is bidiagonal

References

Inspired by matfree.decomp.bidiag [1, 2] and the Golub-Kahan process [3].

Source code in linox/linalg/spectral.py

def lanczos_bidiag(
    A: LinearOperatorLike,
    u0: ArrayLike,
    num_iters: int,
    reortho: bool = True,
) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
    """Lanczos bidiagonalization for SVD computation.

    Reduces a matrix A to bidiagonal form: A ≈ U B V^T, where U and V are
    orthonormal and B is bidiagonal.

    This is the foundation for computing partial SVD of large matrices where
    only a few singular values/vectors are needed.

    Parameters
    ----------
    A : LinearOperatorLike
        Linear operator or matrix of shape (m, n). Should support both
        `A @ v` and `A.T @ u` operations.
    u0 : ArrayLike
        Initial vector of shape (m,) for the bidiagonalization process.
        Will be normalized internally.
    num_iters : int
        Number of bidiagonalization iterations (size of bidiagonal matrix).
        Should be much smaller than min(m, n).
    reortho : bool, optional
        Whether to use full reorthogonalization. This significantly improves
        numerical stability at the cost of O(num_iters^2) operations.
        Default is True (recommended).

    Returns
    -------
    U : jax.Array, shape (m, num_iters)
        Left orthonormal basis vectors (columns).
    V : jax.Array, shape (n, num_iters)
        Right orthonormal basis vectors (columns).
    alpha : jax.Array, shape (num_iters,)
        Diagonal elements of bidiagonal matrix B.
    beta : jax.Array, shape (num_iters-1,)
        Super-diagonal elements of bidiagonal matrix B.

    Notes
    -----
    The bidiagonal matrix B has the form:
        B = [[alpha[0], beta[0],    0,       ...],
             [0,        alpha[1], beta[1],   ...],
             [0,        0,        alpha[2],  ...],
             [...]]

    This is related to Golub-Kahan bidiagonalization and is used in
    algorithms like LSMR and partial SVD computation.

    Examples
    --------
    >>> import jax
    >>> import jax.numpy as jnp
    >>> from linox import Matrix
    >>> key = jax.random.PRNGKey(0)
    >>> A = Matrix(jax.random.normal(key, (100, 50)))
    >>> u0 = jnp.ones(100)
    >>> U, V, alpha, beta = lanczos_bidiag(A, u0, num_iters=10)
    >>> # U and V contain orthonormal vectors
    >>> # B = diag(alpha) + diag(beta, 1) is bidiagonal

    References
    ----------
    Inspired by matfree.decomp.bidiag [1, 2] and the Golub-Kahan process [3].
    """
    u0 = jnp.asarray(u0)
    m = u0.shape[0]
    n = A.shape[1]

    # Normalize initial vector
    beta_0 = jnp.linalg.norm(u0)
    u0 /= beta_0

    # Pre-allocate arrays
    U = jnp.zeros((m, num_iters))
    V = jnp.zeros((n, num_iters))
    alpha = jnp.zeros(num_iters)
    beta = jnp.zeros(num_iters - 1) if num_iters > 1 else jnp.zeros(0)

    # Initialize first left vector
    U = U.at[:, 0].set(u0)

    def bidiag_step(k, carry):
        U_curr, V_curr, alpha_curr, beta_curr = carry

        # Get current left vector
        u = U_curr[:, k]

        # Compute right vector: v = A^T u
        v = A.T @ u

        # Orthogonalize against previous right vector if k > 0
        # In Golub-Kahan: v = A^T u - beta_{k-1} * v_{k-1}
        v_orth = lax.cond(
            k > 0,
            lambda: v - beta_curr[k - 1] * V_curr[:, k - 1],
            lambda: v,
        )

        # Full reorthogonalization against all previous V vectors
        if reortho:

            def reortho_v(v_in):
                # Compute V[:, :k].T @ v using a masked operation
                # Create a mask for the first k columns
                mask = jnp.arange(num_iters) < k
                V_masked = V_curr * mask[None, :]  # Broadcast mask over rows

                # Reorthogonalize twice for better numerical stability (matfree does this)
                coeffs = V_masked.T @ v_in
                v_out = v_in - V_masked @ coeffs
                coeffs = V_masked.T @ v_out
                v_out -= V_masked @ coeffs
                return v_out

            v_orth = lax.cond(k > 0, reortho_v, lambda x: x, v_orth)

        # Compute alpha_k and normalize v
        alpha_k = jnp.linalg.norm(v_orth)
        alpha_curr = alpha_curr.at[k].set(alpha_k)
        v_norm = v_orth / (alpha_k + 1e-16)

        # Store v
        V_curr = V_curr.at[:, k].set(v_norm)

        # Compute next left vector: u = A v - alpha_k * u
        # Only do this if not last iteration
        def compute_next_u():
            u_new = A @ v_norm - alpha_k * u

            # Full reorthogonalization against all previous U vectors
            if reortho:
                # Create a mask for the first k+1 columns
                mask = jnp.arange(num_iters) < (k + 1)
                U_masked = U_curr * mask[None, :]  # Broadcast mask over rows

                # Reorthogonalize twice for better numerical stability
                coeffs = U_masked.T @ u_new
                u_new -= U_masked @ coeffs
                coeffs = U_masked.T @ u_new
                u_new -= U_masked @ coeffs

            beta_k = jnp.linalg.norm(u_new)
            u_norm = u_new / (beta_k + 1e-16)
            return u_norm, beta_k

        u_next, beta_k = lax.cond(
            k < num_iters - 1,
            compute_next_u,
            lambda: (jnp.zeros(m), 0.0),
        )

        # Store u_next and beta if not last iteration
        U_next = lax.cond(
            k < num_iters - 1,
            lambda: U_curr.at[:, k + 1].set(u_next),
            lambda: U_curr,
        )

        beta_next = lax.cond(
            k < num_iters - 1,
            lambda: beta_curr.at[k].set(beta_k),
            lambda: beta_curr,
        )

        return (U_next, V_curr, alpha_curr, beta_next)

    # Run bidiagonalization iterations
    U, V, alpha, beta = lax.fori_loop(0, num_iters, bidiag_step, (U, V, alpha, beta))

    return U, V, alpha, beta