Solvers¶
linear_solve ¶
linear_solve(
operator: AbstractLinearOperator,
vector: PyTree[Any],
solver: Any = None,
*,
options: dict[str, Any] | None = None,
state: Any = None,
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:
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. When a
state is passed, the returned state has the same type, so it can be threaded straight
back into the next call.
A solver that does not implement the stateful API (the StatefulSolver protocol),
such as a plain dense lineax.LU(), has no reusable state to thread. For such a
solver this behaves like lineax.linear_solve, returning the incoming state
unchanged alongside the solution so the (solution, state) return shape stays stable.
AutoSparseLinearSolver ¶
Bases: AbstractLinearSolver[TrackingSolverState]
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 a CUDA GPU, dispatches to CuDSS (factorization reuse, any precision)
if its optional dependency is installed, with no 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, KLU, and CuDSS (update, init_symbolic), so it can be substituted
for any of them. 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. CuDSS supports complex directly, so it needs no equivalent fallback.
By default the chosen solver is combined with an iterative solver (see iterative),
which refines each solution until its relative residual is within tolerance or a step
cap is spent. Pass an IterativeRefinementSettings to tune those, a HybridSettings
to pick another iterative solver, or iterative=False to solve with the chosen direct
solver alone.
platform
class-attribute
instance-attribute
¶
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. CuDSS is chosen only when it
resolves to "gpu" and a CUDA (not ROCm) device is visible. Otherwise Spsolve is
chosen.
iterative
class-attribute
instance-attribute
¶
Whether to combine the selected direct sparse solver with an iterative solver, either
to deal with limited precision (IterativeRefinement) or to allow increased
factorization reuse (HybridDirectIterative). True (the default) refines with the
IterativeRefinementSettings defaults, False uses the direct solver alone, and an
explicit IterativeRefinementSettings or HybridSettings allows choosing a specific
iterative solver and how it should be used.
select_solver ¶
The exact solver AutoSparseLinearSolver will run, including any iterative solver.
Mirrors lineax.AutoLinearSolver.select_solver. With an iterative solver this is a
HybridDirectIterative around 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 ¶
Analyze a sparsity pattern into a symbolic-only state, no values yet.
Accepts a BCOO, BCSR, BCOOLinearOperator, BCSRLinearOperator, a
sparsity-pattern tag, a tagged lineax.JacobianLinearOperator or
lineax.FunctionLinearOperator, or an asdex.ColoredPattern. update
then folds in an operator sharing the pattern and reuses this analysis. The
pattern must be concrete here, not a traced value.
update_and_compute ¶
update_and_compute(
state: _KLUState,
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[
PyTree[Array], RESULTS, dict[str, Any], _KLUState
]
Update state for operator, then solve. See update_then_compute.
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.
transposes_cheaply ¶
Whether transpose reuses the factorization. It always does, through tsolve.
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.
The matrix type follows the operator's lineax tags. An operator tagged
lineax.symmetric_tag that carries a sparsity_pattern_tag is factored as real
symmetric indefinite, from its upper triangle, and the state reports the inertia of the
matrix in state.diagnostics and in stats["diagnostics"]. Any other operator is
factored as a general matrix and reports nothing.
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 ¶
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_and_compute ¶
update_and_compute(
state: _PardisoState,
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[
PyTree[Array], RESULTS, dict[str, Any], _PardisoState
]
Update state for operator, then solve. See update_then_compute.
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.
transposes_cheaply ¶
Whether transpose reuses the factorization. It always does, natively.
CuDSS ¶
CuDSS(
reordering: CuDSSReordering = CuDSSReordering.DEFAULT,
memory: CuDSSMemory = CuDSSMemory.DEVICE,
device_id: int = 0,
diagnostics: bool = False,
matrix_type: CuDSSMatrixType | None = None,
)
Bases: AbstractLinearSolver[_CuDSSState]
Sparse direct solver wrapping NVIDIA's cuDSS library.
Unlike KLU/Pardiso (CPU-only) or Spsolve (no factorization reuse), this solver
runs on a CUDA GPU and reuses factorizations through the stateful solve API
(init, init_symbolic, update). It 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).
Supports float32, float64, complex64, and complex128 directly, with no
upcasting, unlike KLU/Pardiso.
The matrix type follows the operator's lineax tags: general (LU), symmetric (LDL^T), or
symmetric positive semidefinite (Cholesky). Pass matrix_type to force one instead,
which also applies to init_symbolic on a bare sparsity pattern that has no tags. With diagnostics=True, a real symmetric
operator reports the inertia of its matrix in state.diagnostics and in
stats["diagnostics"]. Cholesky reports nothing usable, so a positive semidefinite
operator is then factored as LDL^T like any other symmetric one. Reading them costs one
extra cuDSS query per factorization, so they are off by default and the diagnostics
are empty. AutoSparseLinearSolver picks CuDSS() without them, so construct
CuDSS(diagnostics=True) to get them.
This solver can only handle square nonsingular operators, and only runs on a CUDA GPU (not ROCm, not CPU, not TPU): an error is raised at trace time otherwise.
Every factorization lives in a size-bounded cache rather than behind an explicit
handle, so a release traced inside jax.jit is skipped: the cache evicts old
factorizations on its own, transparently (and correctly, if more slowly) rebuilding
one that is still referenced but was evicted. See Stateful solves
for details.
update on a shared sparsity pattern keeps the analysis and redoes only the numeric
phase. Under CuDSSReordering.COLAMD and CuDSSReordering.BTF_COLAMD it refactorizes
with the previous pivots, and factorizes fresh if the refactorized pivots come out
badly scaled, the same way KLU does. Under every other reordering cuDSS has no
cheaper refactorization, so update always factorizes.
Requires the optional cuDSS dependency, pip install splineax[cudss], which needs
CUDA 13, Python >=3.12, and x86_64 Linux. Constructing CuDSS() raises ImportError
if it isn't installed. AutoSparseLinearSolver prefers CuDSS on a CUDA GPU when it
is.
A plain, stateless solve (lx.linear_solve(op, b, solver=CuDSS()) with no state=)
re-runs the analysis on every call, minting a fresh cache entry each time. That is
correct but wasteful: thread a state through splineax.linear_solve for anything
solved more than once, exactly as recommended for KLU.
init_symbolic ¶
Analyze a sparsity pattern into an analyzed-only state, no values folded in.
Accepts a BCOO, BCSR, BCOOLinearOperator, BCSRLinearOperator, a
sparsity-pattern tag, a tagged lineax.JacobianLinearOperator or
lineax.FunctionLinearOperator, or an asdex.ColoredPattern. A later update
folds in an operator sharing the pattern and reuses this analysis. The pattern
must be concrete here, not a traced value.
update_and_compute ¶
update_and_compute(
state: _CuDSSState,
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[
PyTree[Array], RESULTS, dict[str, Any], _CuDSSState
]
Update state for operator, then solve. See update_then_compute.
update ¶
update(
state: _CuDSSState,
operator: AbstractLinearOperator,
options: dict[str, Any] = {},
) -> _CuDSSState
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 stored token's analysis is reused and only the numeric factorization is redone. Otherwise the operator is analyzed afresh.
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 ¶
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.
update_and_compute ¶
update_and_compute(
state: _SpsolveState,
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[
PyTree[Array], RESULTS, dict[str, Any], _SpsolveState
]
Update state for operator, then solve. See update_then_compute.
HybridDirectIterative ¶
Bases: AbstractLinearSolver[HybridState[Any]]
A direct solver and an iterative solver that share one factorization.
The direct solver factors the matrix, and the iterative solver uses that
factorization as its preconditioner. When a solve gets a new operator and reuse_direct
is on, the factorization of the previous operator is tried first. The iterative solve
runs with a low step cap (max_steps_stale). If its true residual reaches the
tolerance, the old factorization stays. Otherwise the factorization is made again for
the new operator, and the iterative solve runs with the high step cap (max_steps).
The factorization is only reused when the operators share a sparsity pattern tag (see
operator_pattern_tag), and only through update_and_compute, which linear_solve
uses. Plain update always makes a new factorization. Reuse covers solves that are
not differentiated. A derivative solve updates the factorization first.
A solution that never reaches the tolerance is returned as NaN with the result
max_steps_reached. Reuse never causes this, because the check after a new
factorization is the same as without reuse.
The wrapped solver must be square and nonsingular.
direct
instance-attribute
¶
The direct solver. Its state type depends on the solver, so it is not narrowed.
iterative
class-attribute
instance-attribute
¶
The iterative solver, with the options that are specific to it.
reuse_direct
class-attribute
instance-attribute
¶
Whether to try the factorization of the previous operator first. False makes a
new factorization for each new operator.
tol
class-attribute
instance-attribute
¶
The target relative residual, ||b - A x|| <= tol * ||b||.
init_symbolic ¶
Analyze a sparsity pattern through the direct solver.
The state has no operator, so update must give it one before a solve. This raises
AttributeError if the direct solver has no symbolic phase.
update ¶
update(
state: HybridState[Any],
operator: AbstractLinearOperator,
options: dict[str, Any] = {},
) -> HybridState[Any]
Fold a new operator into state through the direct solver.
This makes a new factorization, so the returned state is not stale. It returns the same state object when the direct solver did, so an unchanged operator is a no-op.
update_and_compute ¶
update_and_compute(
state: HybridState[Any],
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[
PyTree[Array], RESULTS, dict[str, Any], HybridState[Any]
]
Fold operator into state and solve, keeping the factorization if it suffices.
When reuse applies, the factorization stays and the iterative solve decides whether
it is good enough. Otherwise this is update followed by compute.
compute ¶
compute(
state: HybridState[Any],
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[PyTree[Array], RESULTS, dict[str, Any]]
Solve against state.
A stale factorization that fails to reach the tolerance is made again for this
solve only. Use update_and_compute to keep the new factorization.
RichardsonOptions ¶
Bases: Module
Options for iterative refinement with the factorization as the preconditioner.
Each step costs one back-substitution and one matrix-vector product. It converges when the preconditioned operator is close enough to the identity, so a stale factorization needs few steps only when the operator has changed a little.
max_steps
class-attribute
instance-attribute
¶
The step cap when the factorization was made for the current operator.
max_steps_stale
class-attribute
instance-attribute
¶
The step cap when the factorization was made for an earlier operator.
max_restarts
class-attribute
instance-attribute
¶
Passes after the first. Richardson checks the true residual itself, so zero.
GMRESOptions ¶
Bases: Module
Options for GMRES with the factorization as the preconditioner.
GMRES finds the best correction in a Krylov space, so it can converge when the factorization is too stale for Richardson iteration.
restart
class-attribute
instance-attribute
¶
The size of the Krylov space built before a restart.
stagnation_iters
class-attribute
instance-attribute
¶
How many restarts without a decrease of the residual end the solve.
max_steps
class-attribute
instance-attribute
¶
The step cap when the factorization was made for the current operator.
max_steps_stale
class-attribute
instance-attribute
¶
The step cap when the factorization was made for an earlier operator.
max_restarts
class-attribute
instance-attribute
¶
Passes after the first. GMRES stops on a preconditioned residual, so the true residual can end above the target and a pass from the current solution fixes it.
BiCGStabOptions ¶
Bases: Module
Options for BiCGStab with the factorization as the preconditioner.
BiCGStab handles a nonsymmetric operator with short recurrences. Each step needs two matrix-vector products and two solves.
max_steps
class-attribute
instance-attribute
¶
The step cap when the factorization was made for the current operator.
max_steps_stale
class-attribute
instance-attribute
¶
The step cap when the factorization was made for an earlier operator.
max_restarts
class-attribute
instance-attribute
¶
Passes after the first. BiCGStab stops on a preconditioned residual, so the true residual can end above the target and a pass from the current solution fixes it.
CGOptions ¶
Bases: Module
Options for conjugate gradients with the factorization as the preconditioner.
The operator must be symmetric and positive definite, and carry
lineax.positive_semidefinite_tag. The factorization of an earlier matrix of that
kind is a positive definite preconditioner. Each step needs one matrix-vector product
and one solve, with short recurrences and no growing Krylov space.
max_steps
class-attribute
instance-attribute
¶
The step cap when the factorization was made for the current operator.
max_steps_stale
class-attribute
instance-attribute
¶
The step cap when the factorization was made for an earlier operator.
max_restarts
class-attribute
instance-attribute
¶
Passes after the first. CG stops on a preconditioned residual, so the true residual can end above the target and a pass from the current solution fixes it.
ReuseOptions ¶
Bases: Module
Options for when HybridDirectIterative keeps a factorization for a new operator.
max_reuses
class-attribute
instance-attribute
¶
How many solves in a row may use the same factorization before it is made again.
None sets no limit.
slow_fraction
class-attribute
instance-attribute
¶
The fraction of max_steps_stale that a solve with an old factorization may use
before the next solve makes a new one. None never makes one for this reason.
HybridSettings ¶
Bases: Module
The options of a HybridDirectIterative, without a direct solver bound yet.
HybridState ¶
Bases: Module, Generic[_StateT]
The state of a HybridDirectIterative solve.
It wraps the state of the direct solver, together with the operator being solved and
the bookkeeping for the reuse of a factorization. A state straight from
init_symbolic has no operator and is not solvable until update gives it one.
inner_state
instance-attribute
¶
The state of the direct solver. Its factorization may be stale.
operator
instance-attribute
¶
The operator that this state represents. None for a symbolic-only state.
stale
instance-attribute
¶
Whether the factorization was made for an earlier operator.
reuses
instance-attribute
¶
How many solves in a row have used the current factorization while it was stale.
refactor_next
instance-attribute
¶
Whether the next solve makes a new factorization before solving.
transposed
class-attribute
instance-attribute
¶
Whether this state solves the transposed system. A transposed state always has a
current factorization, since transpose makes a new one where it was stale.
diagnostics
property
¶
The diagnostics of the inner state.
They describe the inner factorization, which is stale while stale is true, so
they belong to an earlier operator until the next refactorization.
track ¶
Order a later release after solution, through the inner state.
This does nothing for an inner state that owns no memory.
Richardson ¶
Bases: AbstractLinearSolver[_RichardsonState]
Preconditioned Richardson iteration, which is iterative refinement.
Starting from y0, each step adds the preconditioner applied to the residual,
y = y + M (b - A y). With the inverse of a factorization as the preconditioner,
every step is one back-substitution and one matrix-vector product.
It takes the same options as the Krylov solvers in lineax. A preconditioner
operator defaults to the identity, and y0 defaults to zeros. The iteration stops
when ||b - A y|| <= rtol * ||b|| + atol, with the relative part floored at
machine precision (see _CONVERGENCE_FLOOR_ULPS), or after max_steps steps.
IterativeRefinement ¶
IterativeRefinement(
solver: SparseLinearSolver[Any],
tol: float = 1e-10,
max_steps: int = 10,
) -> HybridDirectIterative
Wrap a stateful solver so that each solve is refined with iterative refinement.
The wrapped solver supplies the factorization. Each solve starts from one direct
solve and then corrects it with the residual until ||b - A x|| <= tol * ||b||, or
returns NaN after max_steps corrections. Every new operator is factored, so this is
a HybridDirectIterative with RichardsonOptions and reuse_direct=False.
Arguments:
solver: the stateful solver to wrap. It must be square and nonsingular.tol: the target relative residual. Defaults to1e-10.max_steps: the maximum number of correction steps. Defaults to10.
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. See
IterativeRefinement for what each does.
ReorderingScheme ¶
Bases: IntEnum
CuDSSReordering ¶
Bases: IntEnum
Fill-reducing reordering passed to analyze. Distinct from ReorderingScheme
(Spsolve's cuSOLVER reordering, an unrelated enum despite the similar name).
CuDSSMatrixType ¶
Bases: IntEnum
The factorization cuDSS runs, which CuDSS(matrix_type=...) can force.
CuDSSMemory ¶
Diagnostics¶
What a solver reports about its factorization, in state.diagnostics and in
solution.stats["diagnostics"]. See
Factorization diagnostics.
SolverDiagnostics ¶
Bases: Module
What a solver reports about the factorization it holds.
Every field is None when the solver does not report it, and an array otherwise. A
None field is part of the pytree structure, so the diagnostics of one state keep the
same structure inside a jax.lax.while_loop, jax.lax.cond or jax.lax.scan. A solver
that reports nothing holds SolverDiagnostics().
The diagnostics of a state are state.diagnostics. A solve also returns them as
solution.stats["diagnostics"], which is the way to read them inside
splineax.stateful_solve_transform, where the final state is only returned at the end.
positive_definite
class-attribute
instance-attribute
¶
Whether the matrix is positive definite, which means every pivot is positive and none was perturbed. A singular matrix whose zero pivot was replaced by a small number is not reported as definite.
positive
class-attribute
instance-attribute
¶
The number of positive eigenvalues of the factored matrix.
negative
class-attribute
instance-attribute
¶
The number of negative eigenvalues of the factored matrix.
perturbed_pivots
class-attribute
instance-attribute
¶
The number of pivots the solver replaced by a small number, which is nonzero for a singular or nearly singular matrix.
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,
) -> _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 callslineax.linear_solveinternally.filter_solver: which solves to thread, as a solver class matched byisinstanceor a boolean predicate. The defaultStatefulSolverthreads only solvers that implement the stateful API, so a plain denselineax.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 initialstateis passed at call time, false otherwise. 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: PatternSource
| _ContentPatternTag
| TaggedOperator
| None = None,
) -> PatternTag
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.
A content-hashed tag also lets a solver turn a tagged lineax.JacobianLinearOperator
or lineax.FunctionLinearOperator into a BCOO. The Jacobian coloring this needs is
computed on first use and cached on the tag. A pattern that already holds a coloring
(an asdex.ColoredPattern) keeps it. Use splineax.sparsity_coloring_tag to
compute the coloring up front.
Given a tag, or a lineax operator that carries one, that tag is returned.
sparse_indices_sorted
module-attribute
¶
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 ¶
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, and update_and_compute
does that and solves in one call, so a solver can choose how much of the state to
rebuild after seeing the result of the solve.
The states a solver produces from init, update/update_and_compute, 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.
update_and_compute ¶
update_and_compute(
state: _StateT,
operator: AbstractLinearOperator,
vector: PyTree[Array],
options: dict[str, Any],
) -> tuple[PyTree[Array], RESULTS, dict[str, Any], _StateT]
Fold operator into state and solve against vector.
Returns the solution, the result code, the solver statistics and the updated
state. A solver may defer or skip the work in update when the solve shows that
the existing state was good enough.
TrackingSolverState ¶
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.
A state also reports what its solver knows about the factorization it holds, as
diagnostics. A solver that reports nothing holds an all-None SolverDiagnostics.
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.