Solvers¶
AutoSparseLinearSolver ¶
Bases: AbstractSparseLinearSolver[SparseBasicState | SparseSymbolicState | SparseNumericState]
Selects a sparse direct solver based on the JAX platform and precision.
Dispatches to KLU (SuiteSparse direct solve with factorization reuse) only on CPU
with x64 enabled, since klujax is double precision only. 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. Exposes the same factorization API as KLU
(factorize, factorize_symbolic), so it can be substituted for KLU verbatim.
When it dispatches to Spsolve, these factorization calls degrade to no-ops.
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. KLU is chosen only when
this resolves to "cpu" and x64 is enabled, otherwise Spsolve is chosen.
select_solver ¶
Check which solver AutoSparseLinearSolver will dispatch to.
Mirrors lineax.AutoLinearSolver.select_solver. The operator is accepted for
signature parity but selection depends only on the platform.
KLU ¶
Bases: AbstractSparseLinearSolver[_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, and importing it enables JAX's x64 mode and forces
the CPU platform globally (this happens lazily, on the first solve).
This solver can only handle square nonsingular operators.
factorize ¶
factorize(
operator: AbstractLinearOperator,
options: dict[str, Any] = {},
) -> AbstractContextManager[SparseNumericState]
Pre-compute a full (symbolic + numeric) factorization for reuse.
Equivalent to self.init(operator, options).factorize().
factorize_symbolic ¶
factorize_symbolic(
sparsity: BCOO
| BCSR
| BCOOLinearOperator
| BCSRLinearOperator
| SparseJacobianLinearOperator
| SparseJacobianLinearOperatorColoring
| JacobianColoring,
)
Open a scope with a pre-computed KLU symbolic factorization.
Yields a _KLUSymbolicScope. Inside the block, call:
- .init(operator) to create a _KLUSymbolicState for lx.linear_solve
(uses solve_with_symbol: factors numerically on each call, symbolic reused).
- .init(operator).factorize() or equivalently .factorize(operator) to also
pre-compute the numeric factorization (uses solve_with_numeric).
The symbolic handle is freed when the with block exits, after all
registered solve-result dependencies have been consumed.
Args:
sparsity: Sparse matrix whose sparsity pattern to pre-analyze. Accepts
BCOO, BCSR, BCOOLinearOperator, BCSRLinearOperator,
SparseJacobianLinearOperator,
SparseJacobianLinearOperatorColoring, or JacobianColoring.
For the latter three, the indices are taken host-side from the
precomputed asdex sparsity pattern, without materialising the
Jacobian numerically. That host-side read means the coloring
must be concrete here, not a traced value inside a jitted
function.
Spsolve ¶
Bases: AbstractSparseLinearSolver[_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).
This solver can only handle square nonsingular operators.
ReorderingScheme ¶
Bases: IntEnum
Base class and protocols¶
All solvers subclass AbstractSparseLinearSolver and structurally satisfy the
SparseLinearSolver protocol. The factorization-reuse API is described below (see
Advanced usage).
AbstractSparseLinearSolver ¶
Bases: AbstractLinearSolver[_SolverState], Generic[_SolverState]
Abstract base for sparse direct solvers that support factorization reuse.
Extends the lineax AbstractLinearSolver interface with factorize and
factorize_symbolic. Concrete subclasses (KLU, Spsolve,
AutoSparseLinearSolver) are therefore usable both with lineax.linear_solve
(which requires an AbstractLinearSolver) and the factorization-reuse API. They
also structurally satisfy the SparseLinearSolver protocol.
factorize
abstractmethod
¶
factorize(
operator: AbstractLinearOperator,
options: dict[str, Any] = {},
) -> AbstractContextManager[SparseNumericState]
Pre-compute a full factorization for reuse across multiple solves.
factorize_symbolic
abstractmethod
¶
Pre-compute a symbolic factorization from a known sparsity pattern.
SparseLinearSolver ¶
Bases: Protocol
Structural type for sparse direct solvers that expose factorization reuse on
top of the lineax AbstractLinearSolver interface.
SparseSymbolicScope ¶
SparseBasicState ¶
Bases: Protocol
The state returned by SparseLinearSolver.init.
Can be turned into a numeric factorization for reuse across solves.
factorize ¶
Pre-compute a numeric factorization, yielding a reusable state.
SparseSymbolicState ¶
Bases: Protocol
A state that reuses a pre-computed symbolic factorization.
Returned by SparseSymbolicScope.init. Directly solvable, and can additionally
be turned into a numeric factorization.
factorize ¶
Pre-compute a numeric factorization, reusing the symbolic one.
SparseNumericState ¶
Bases: Protocol
A fully factorized sparse solver state, ready to pass to lineax.linear_solve.
Marker protocol: a terminal state with no further factorization step.