Solvers¶

linear_solve ¶

linear_solve(
    operator: AbstractLinearOperator,
    vector: PyTree[Any],
    solver: Any = None,
    *,
    options: dict[str, Any] | None = None,
    state: PyTree[Any] = sentinel,
    throw: bool = True,
) -> tuple[Solution, Any]

Solve operator @ x = vector, returning the solution and an updated state.

A wrapper over lineax.linear_solve for the stateful sparse API. It runs the solver's init or update to fold the operator into a state, solves, then tracks the solution against the state so a later release is ordered after it. Unlike lineax.linear_solve, it returns a (solution, state) tuple:

solution, state = splineax.linear_solve(operator, vector, solver, state=state)

With no state, a fresh one is built with solver.init. The default solver is AutoSparseLinearSolver, which picks a backend for the platform and precision.

AutoSparseLinearSolver ¶

Bases: AbstractLinearSolver[Any]

Selects a sparse direct solver based on the JAX platform, precision, and what is installed, and by default refines its solution with iterative refinement.

On CPU with x64 enabled, dispatches to Pardiso (Intel oneMKL Pardiso, factorization reuse) if the optional pardiso-mkl-jax dependency is installed, otherwise KLU (SuiteSparse, factorization reuse). Both are double precision only, hence the x64 requirement. On any other backend, or on CPU when x64 is disabled, it dispatches to Spsolve, which works in single or double precision and on any backend. It exposes the same stateful API as Pardiso and KLU (update, init_symbolic), so it can be substituted for either. When it dispatches to Spsolve, the reuse calls degrade to no-ops.

pardiso_mkl_jax does not support complex matrices (see Pardiso's docstring), so init falls back to KLU for a complex operator even when Pardiso was otherwise selected, keeping Auto able to solve anything KLU can. init_symbolic cannot make the same check, since a bare sparsity pattern carries no values to inspect, so it stays on Pardiso. Construct KLU() directly for symbolic-pattern reuse on a complex operator.

By default the chosen solver is wrapped in IterativeRefinement, which improves each solution until its relative residual is within tolerance or a step cap is spent. Pass an IterativeRefinementSettings to tune those, or iterative_refinement=False to solve with the chosen direct solver alone.

platform `class-attribute` `instance-attribute` ¶

platform: str | None = None

Platform to select for. If None, jax.default_backend() is used. Set to e.g. "cpu", "gpu", or "tpu" to override the choice explicitly. Pardiso/KLU are chosen only when this resolves to "cpu" and x64 is enabled, otherwise Spsolve is chosen.

iterative_refinement `class-attribute` `instance-attribute` ¶

iterative_refinement: bool | IterativeRefinementSettings = (
    True
)

Whether to refine the direct solve, and with what settings. True refines with the IterativeRefinementSettings defaults, False disables it, and an explicit IterativeRefinementSettings sets the tolerance and step cap.

select_solver ¶

select_solver(
    operator: AbstractLinearOperator,
) -> SparseLinearSolver[Any]

The exact solver AutoSparseLinearSolver will run, including any refinement.

Mirrors lineax.AutoLinearSolver.select_solver. With refinement on this is an IterativeRefinement wrapping the chosen direct solver. The operator is accepted for signature parity but selection depends only on the platform.

KLU ¶

Bases: AbstractLinearSolver[_KLUState]

Sparse direct solver wrapping the klujax (SuiteSparse KLU) library.

This solver keeps the operator in its native sparse (COO) storage rather than densifying it, and so is intended for use with the sparse operators in this package (BCOOLinearOperator and BCSRLinearOperator).

klujax is CPU and double-precision only: float32/complex64 inputs are upcast to float64/complex128. It does not enable JAX's x64 mode or force the CPU platform on import, so jax_enable_x64 must already be on before this solver runs. klujax raises a clear error otherwise.

This solver can only handle square nonsingular operators.

init_symbolic ¶

init_symbolic(
    sparsity: _Sparsity, options: dict[str, Any] = {}
) -> _KLUState

Analyze a sparsity pattern into a symbolic-only state, no values yet.

Accepts a BCOO, BCSR, BCOOLinearOperator, BCSRLinearOperator, SparseJacobianLinearOperator, SparseJacobianLinearOperatorColoring, or JacobianColoring. update then folds in an operator sharing the pattern and reuses this analysis. The pattern must be concrete here, not a traced value.

update ¶

update(
    state: _KLUState,
    operator: AbstractLinearOperator,
    options: dict[str, Any] = {},
) -> _KLUState

Fold a new operator into state, reusing the analysis where the pattern holds.

Repeated calls with the same operator object are a no-op. When the operator shares the state's sparsity tag, the symbolic analysis is reused and only the numeric factorization is redone. Otherwise the operator is analyzed from scratch.

Pardiso ¶

Pardiso()

Bases: AbstractLinearSolver[_PardisoState]

Sparse direct solver wrapping pardiso_mkl_jax (Intel oneMKL Pardiso).

This solver keeps the operator in its native sparse (CSR) storage rather than densifying it, and so is intended for use with the sparse operators in this package (BCOOLinearOperator and BCSRLinearOperator).

pardiso_mkl_jax is CPU, real-valued, and double-precision only: float32 inputs are upcast to float64, and complex operators raise TypeError (Pardiso's complex matrix types are not supported by pardiso_mkl_jax yet). It does not enable JAX's x64 mode or force the CPU platform on import, so jax_enable_x64 must already be on before this solver runs.

This solver can only handle square nonsingular operators.

Requires the optional pardiso-mkl-jax dependency (pip install splineax[pardiso]). Constructing Pardiso() raises ImportError if it is not installed. AutoSparseLinearSolver prefers Pardiso over KLU on CPU with x64 enabled, falling back to KLU automatically when pardiso-mkl-jax is missing.

Arguments:

Nothing.

init_symbolic ¶

init_symbolic(
    sparsity: _Sparsity, options: dict[str, Any] = {}
) -> _PardisoState

Record the pattern for reuse, deferring analysis to the first update.

Under Pardiso's default weighted matching the analysis depends on the matrix values, so a values-independent symbolic phase is not sound. This keeps only the shape and a pattern tag. The first update with real values runs analyze and factor.

update ¶

update(
    state: _PardisoState,
    operator: AbstractLinearOperator,
    options: dict[str, Any] = {},
) -> _PardisoState

Fold a new operator into state, reusing the analysis where the pattern holds.

Repeated calls with the same operator object are a no-op. When the operator shares the state's pattern and an analysis already exists, only the numeric factorization is redone. Otherwise the operator is analyzed from scratch.

Spsolve ¶

Bases: AbstractLinearSolver[_SpsolveState]

Sparse direct solver wrapping jax.experimental.sparse.linalg.spsolve.

This solver keeps the operator in its native sparse (CSR) storage rather than densifying it, and so is intended for use with the sparse operators in this package (BCOOLinearOperator and BCSRLinearOperator). Internally spsolve performs a sparse QR factorization (CUDA native; on CPU it falls back to scipy.sparse.linalg.spsolve).

It has no separate factorization phase, so the stateful reuse API (init_symbolic, update, and the state's release) is a set of no-ops here, for parity with KLU and Pardiso.

This solver can only handle square nonsingular operators.

init_symbolic ¶

init_symbolic(
    sparsity: _Sparsity, options: dict[str, Any] = {}
) -> _SpsolveState

No-op symbolic init, for parity with KLU.

Spsolve cannot pre-analyze a sparsity pattern, so this returns an empty state that update fills with the first real operator.

update ¶

update(
    state: _SpsolveState,
    operator: AbstractLinearOperator,
    options: dict[str, Any] = {},
) -> _SpsolveState

Rebuild the state from operator, since Spsolve reuses no factorization.

Repeated calls with the same operator object are a no-op.

IterativeRefinement ¶

Bases: AbstractLinearSolver[_IterativeRefinementState]

Wraps a stateful solver and refines each solve with iterative refinement.

The wrapped solver supplies the factorization and the per-step solves. compute runs the refinement loop (see iterative_refinement), reusing that factorization for both the initial solve and every correction. Every other method delegates to the wrapped solver, so IterativeRefinement exposes the same stateful API (init, init_symbolic, update, transpose, conj) and can stand in for the solver it wraps. Its state releases through the inner state, so state.release() still frees it.

The wrapped solver must be square and nonsingular, since refinement assumes the correction solve returns a genuine approximate inverse.

init_symbolic ¶

init_symbolic(
    sparsity: _Sparsity, options: dict[str, Any] = {}
) -> _IterativeRefinementState

Analyze a sparsity pattern, deferring to the wrapped solver's init_symbolic.

The resulting state has no operator yet, so update must fold one in before a solve. Raises AttributeError if the wrapped solver has no symbolic phase.

update ¶

update(
    state: _IterativeRefinementState,
    operator: AbstractLinearOperator,
    options: dict[str, Any] = {},
) -> _IterativeRefinementState

Fold a new operator into state through the wrapped solver.

Returns the same state object when the wrapped update did, so an update with an unchanged operator stays a no-op.

IterativeRefinementSettings ¶

Bases: Module

The tol and max_steps of an iterative refinement, without a solver bound yet.

A solver that offers refinement as an option, such as AutoSparseLinearSolver, takes one of these instead of separate tolerance and step-cap arguments, so the two settings travel together. See IterativeRefinement for what each does.

tol `class-attribute` `instance-attribute` ¶

tol: float = eqx.field(default=1e-10, static=True)

Target relative residual, ||b - A x|| <= tol * ||b||.

max_steps `class-attribute` `instance-attribute` ¶

max_steps: int = eqx.field(default=10, static=True)

Maximum correction steps before returning NaN.

ReorderingScheme ¶

Bases: IntEnum

Stateful solve transform¶

Threads a solver state through a function's lineax.linear_solve calls, so its solves reuse a factorization. See Transforming existing Lineax code.

stateful_solve_transform ¶

stateful_solve_transform(
    fn: Callable[..., _OutputT],
    *,
    filter_solver: _FilterSolver = StatefulSolver,
    return_final_state: bool | None = None,
    pass_through_custom_diff: bool = False,
) -> _WrappedFunction[_OutputT]

Thread a solver state through a function's lineax.linear_solve calls.

The wrapped function takes the original arguments plus a state keyword for an initial state, defaulting to None, in which case init runs at the first solve. A function whose own signature already has a state argument cannot be wrapped, since the keyword is taken.

Arguments:

fn: the function to transform. It calls lineax.linear_solve internally.
filter_solver: which solves to thread, as a solver class matched by isinstance or a boolean predicate. The default StatefulSolver threads only solvers that implement the stateful API, so a plain dense lineax.LU() passes through.
return_final_state: when true the wrapped function returns (output, final_state), when false it returns the output alone and releases the threaded state. The default is true when an initial state is passed at call time, false otherwise.
pass_through_custom_diff: by default a matched solve inside a custom_jvp or custom_vjp raises, since the state cannot cross the custom rule. Set this true to let such a solve run without threading, so it works but does not reuse a factorization.

Returns:

A function taking fn's arguments plus a state keyword for an initial state. It returns fn's output, or the output paired with the final state when a state is kept.

Sparsity tags¶

sparsity_pattern_tag ¶

sparsity_pattern_tag(
    pattern: _Sparsity | None = None,
) -> object

Create a tag marking an operator's structural sparsity pattern.

Attach the tag to operators through their tags argument. Two operators carrying equal tags are asserted to have exactly the same index arrays, in the same order, so a solver may reuse one operator's factorization for the other.

Given a concrete pattern, the tag is content-hashed, so independently tagged operators with the same indices get equal tags. With no argument, or a pattern whose indices are traced under jit, the tag instead carries a random id. Thread that one tag object onto every operator sharing the pattern to mark them as equal.

sparse_indices_sorted `module-attribute` ¶

sparse_indices_sorted = _HasRepr('sparse_indices_sorted')

One global assertion that an operator's indices are already row-major sorted, so Pardiso and Spsolve may skip the sort they would otherwise do in init.

BCOOLinearOperator and BCSRLinearOperator add this automatically when the matrix they wrap already carries indices_sorted.

Protocols¶

Solvers structurally satisfy the SparseLinearSolver protocol, which extends the solver-agnostic StatefulSolver. The stateful reuse API is described in Stateful solves.

SparseLinearSolver ¶

Bases: StatefulSolver[_StateT], Protocol[_StateT]

Structural type for the sparse stateful solvers in this package.

Extends the solver-agnostic StatefulSolver (init, update, compute, transpose, conj, assume_full_rank) with init_symbolic, which analyzes a known sparsity pattern into a reusable state before any values are available. KLU, Pardiso, Spsolve, and AutoSparseLinearSolver all satisfy it structurally.

init_symbolic ¶

init_symbolic(
    sparsity: _Sparsity, options: dict[str, Any] = {}
) -> _StateT

Analyze a sparsity pattern into a state, reused by a later update.

StatefulSolver ¶

Bases: Protocol[_StateT]

A solver that creates and updates a reusable state.

This is the part of the lineax AbstractLinearSolver interface we rely on, plus update. A solver satisfies it structurally, so no base class is needed. update folds new information about the operator into an existing state.

The states a solver produces from init, update, and a state's track should share one pytree structure, so a state can be carried through a scan or while_loop, whose carry has a fixed structure. The sparse init_symbolic state may differ, since it holds only a symbolic analysis. Such a state must be updated before it is carried through a loop, for example by unrolling the first iteration.

update ¶

update(
    state: _StateT,
    operator: AbstractLinearOperator,
    options: dict[str, Any] = {},
) -> _StateT

Fold a new operator into state, reusing prior work where possible.

TrackingState ¶

Bases: Protocol

A solver state that records solves depending on it and frees its own memory.

A state may own memory that must outlive every solve made with it. track marks a solution as a dependency, so a later release is ordered after that solve, and release frees that memory once the state is done. Both live on the state, so a caller releases without a reference to the solver. A state that owns nothing implements track as a no-op returning self, and release as a no-op.

track ¶

track(solution: PyTree[Array]) -> Self

Return a new state whose eventual release is ordered after solution.

release ¶

release() -> None

Free any memory this state owns, ordered after its tracked solves.

PerformanceWarning ¶

Bases: UserWarning

Raised when a sparse solver has to do work that a differently prepared input would have avoided.

Currently only used by Spsolve and Pardiso, when their init sorts an unsorted BCOO or BCSR operator before solving. Both need row-major sorted indices and will silently sort them for you, but doing so on every init is wasted work if the same operator is solved more than once. Passing an already-sorted matrix (for a BCOO, call .sort_indices() once yourself) avoids the warning and the repeated cost.

Solvers¶

linear_solve ¶

AutoSparseLinearSolver ¶

platform class-attribute instance-attribute ¶

iterative_refinement class-attribute instance-attribute ¶

select_solver ¶

KLU ¶

init_symbolic ¶

update ¶

Pardiso ¶

init_symbolic ¶

update ¶

Spsolve ¶

init_symbolic ¶

update ¶

IterativeRefinement ¶

init_symbolic ¶

update ¶

IterativeRefinementSettings ¶

tol class-attribute instance-attribute ¶

max_steps class-attribute instance-attribute ¶

ReorderingScheme ¶

Stateful solve transform¶

stateful_solve_transform ¶

Sparsity tags¶

sparsity_pattern_tag ¶

sparse_indices_sorted module-attribute ¶

Protocols¶

SparseLinearSolver ¶

init_symbolic ¶

StatefulSolver ¶

update ¶

TrackingState ¶

track ¶

release ¶

PerformanceWarning ¶

platform `class-attribute` `instance-attribute` ¶

iterative_refinement `class-attribute` `instance-attribute` ¶

tol `class-attribute` `instance-attribute` ¶

max_steps `class-attribute` `instance-attribute` ¶

sparse_indices_sorted `module-attribute` ¶