Skip to content

Stateful solves

A sparse direct solver generally does some expensive work (like reordering or constructing a factorization) that does not change between right-hand sides, or between matrices that share a structure. splineax lets you do that work once and reuse it, through an explicit solver state that you create, update, and thread through your solves.

The idea is to give the solver information about the operators it will solve, as soon as you have it them, and then let it decide what to recompute. You never ask it to factor or refactor by hand. You hand it an operator, and it reuses whatever it still can, depending on the numerical properties of the operator and solver implementation details.

If you have a function that already calls lineax.linear_solve and you would rather not thread the state by hand, there is a stateful_solve_transform function transformation that does the threading for you, described in Transforming existing Lineax code.

The stateful solve API

Every solver in this package exposes the same small API.

  • state = solver.init(operator) builds a state for an operator.
  • state = solver.update(state, operator) folds a new operator into an existing state, reusing prior work where the two operators allow it.
  • solution, result, stats, state = solver.update_and_compute(state, operator, vector, options) updates the state and solves in one call. A solver can use it to decide what to rebuild after it has seen the result of the solve. splineax.linear_solve calls it.
  • state.release() says you are done with the state, so any memory it holds may go. It is optional, and a no-op for solvers whose state holds nothing.
  • state = state.track(solution) records that a solution depends on the state, so a later release is ordered after that solve.

lineax.linear_solve does not return an updated state, so on its own these steps read as:

state = solver.update(state, operator)
solution = lineax.linear_solve(operator, vector, solver, state=state)
state = state.track(solution)

splineax.linear_solve does all three for you and returns a (solution, state) tuple:

import jax
import jax.numpy as jnp
import lineax as lx
from jax.experimental.sparse import BCOO

import splineax as splx

# KLU and Pardiso require 64-bit mode.
jax.config.update("jax_enable_x64", True)

dense = jnp.array(
    [
        [10.0, 2.0, 0.0, 0.0],
        [3.0, 14.0, 5.0, 0.0],
        [0.0, 6.0, 18.0, 9.0],
        [0.0, 0.0, 1.0, 12.0],
    ]
)
b1 = jnp.array([1.0, 2.0, 3.0, 4.0])
b2 = b1[::-1]

operator = splx.BCOOLinearOperator(BCOO.fromdense(dense))
solver = splx.KLU()

# The first call builds a fresh state. Thread it back in to reuse the factorization.
solution, state = splx.linear_solve(operator, b1, solver)
solution, state = splx.linear_solve(operator, b2, solver, state=state)

# Release it once you are done solving.
state.release()

With no state, splineax.linear_solve builds one with init. With a state, it calls update, so passing the same operator again will cost nothing, and passing a matrix that shares the structure will reuse the analysis. Either way it tracks the solution before returning, so the state you get back is safe to release after the loop.

The default solver is AutoSparseLinearSolver, which picks a backend for the platform and precision. Any splineax solver works in its place.

Reuse across changing values

Often you know the sparsity pattern before the values, or you solve a family of matrices that share a pattern. Analyze the pattern once with init_symbolic, then update folds in each matrix and reuses that analysis:

# Only the structure matters here, not the values.
sparsity = BCOO.fromdense(dense)

state = solver.init_symbolic(sparsity)
state = solver.update(state, operator)
solution = lx.linear_solve(operator, b1, solver=solver, state=state).value
state.release()

init_symbolic accepts a BCOO, BCSR, BCOOLinearOperator, BCSRLinearOperator, a sparsity-pattern tag, a tagged lineax.JacobianLinearOperator or lineax.FunctionLinearOperator, or an asdex.ColoredPattern. Only its sparsity pattern is read. A tag or a tagged operator gives its pattern without materialising anything, so solver.init_symbolic(tag) works before you know any point to evaluate a Jacobian at.

Shared patterns between operators

For update to reuse an analysis, it has to know the new operator has the same structure as the last one. You assert that with a tag from splineax.sparsity_pattern_tag, attached to both operators:

tag = splx.sparsity_pattern_tag(sparsity)
first = splx.BCOOLinearOperator(BCOO.fromdense(dense), tags=tag)
second = splx.BCOOLinearOperator(BCOO.fromdense(2.0 * dense), tags=tag)

state = solver.init(first)
# The shared tag lets update reuse the analysis.
state = solver.update(state, second)
solution = lx.linear_solve(second, b1, solver=solver, state=state).value
state.release()

Two operators carrying the same tag are asserted to have exactly the same index arrays, in the same order. Given a concrete pattern the tag is a content hash, so operators tagged separately still compare equal when their indices match. With no argument, or under jit where the indices are traced, sparsity_pattern_tag() instead returns a marker you thread onto every operator sharing the pattern. Jacobian operators at different points that carry one tag from splineax.sparsity_coloring_tag reuse a factorization in the same way (see Operators). Transposing a tagged operator will let some solvers reuse factorizations as well.

There is also a sparse_indices_sorted tag. Attaching it to an operator asserts its indices are already row-major sorted, so Pardiso and Spsolve skip the sort they would otherwise do.

Solves inside jax.jit

A factorization handle is an ordinary JAX value, not a native object tied to the Python side, so the whole lifecycle composes inside a jitted function. Build the state, solve, and release, all under one trace:

@jax.jit
def solve_under_jit(values, b):
    operator = splx.BCOOLinearOperator(
        BCOO((values, sparsity.indices), shape=sparsity.shape, indices_sorted=True)
    )
    solution, state = splx.linear_solve(operator, b, solver)
    state.release()
    return solution.value


x = solve_under_jit(sparsity.data, b1)

state.track ties the state and the solution obtained using it together: it creates an artificial data dependency of the state on the solution (via entangle-jax) so that any further operation that uses the state will have to be ordered after the solve. This is handled by splineax.linear_solve automatically.

CuDSS tracks its state the same way, but for a different reason. Each factorize renames the token's cache entry, so a solve that still holds the old token has to run before the update that factorizes new values. Otherwise the solve finds no entry and rebuilds the factorization from the token's arrays. CuDSS has no traced release, so jit skips a release inside it. The cache evicts old factorizations by itself, and a solve whose factorization was evicted rebuilds it. That is slower but still correct.

What each solver reuses

The API is the same across solvers, but what they reuse differs.

KLU keeps two handles, a symbolic analysis and a numeric factorization. init builds both. update on a matching pattern reuses the analysis and rebuilds the numeric factor for the new values, also when the new operator has the transposed pattern. transpose reuses both and solves the transposed system directly.

Pardiso keeps one factorization handle. Under its default weighted matching, an analysis that ignores the values is not sound, so init_symbolic defers the analysis. It records the pattern only, and the first update with real values runs analyze and factor. Later updates on the same pattern refactor while reusing that analysis. Like KLU, it reuses the handle for an operator with the transposed pattern.

CuDSS keeps one token that carries its analysis forward. init analyzes and factorizes, init_symbolic analyzes only, and update on a matching pattern reruns the numeric factor from the stored token. cuDSS renames the token's cache entry as it advances a phase rather than dropping the analysis, so update reuses the analysis without re-running it. With CuDSSReordering.COLAMD or CuDSSReordering.BTF_COLAMD, update also reuses the previous pivots through cuDSS's refactorization, and factorizes fresh when the reused pivots come out badly scaled, the same guard KLU uses. Under the other reorderings cuDSS has no cheaper refactorization, so update always factorizes. transpose reuses the factors for a symmetric matrix, and builds a genuine A^T factorization for a general one, since cuDSS has no transpose solve. For the same reason update analyzes an operator with the transposed pattern again.

Spsolve has no separate factorization phase, so the reuse API is a set of no-ops for parity. update rebuilds the state, release frees nothing, and track returns the state unchanged. Code written against the API runs unchanged on any backend, and AutoSparseLinearSolver forwards to whichever it picked.

Diagnostics in a loop

A solve returns the diagnostics of its factorization in solution.stats["diagnostics"] (see Factorization diagnostics). Inside stateful_solve_transform these stats are the way to read them, since the final state only comes back at the end. The example below searches for the smallest shift that makes matrix + shift * I positive definite. It makes the factorization again for every shift, and the loop condition reads the flag.

shifts = jnp.array([0.0, 1.0, 2.0, 4.0, 8.0])


def search():
    def not_definite(carry):
        index, definite = carry
        return ~definite & (index + 1 < len(shifts))

    def next_shift(carry):
        index = carry[0] + 1
        shifted = operator_at(shifts[index])
        solution = lx.linear_solve(shifted, b, splx.Pardiso())
        return index, solution.stats["diagnostics"].positive_definite

    index, _ = jax.lax.while_loop(
        not_definite, next_shift, (jnp.int32(0), jnp.bool_(False))
    )
    return shifts[index]


smallest = jax.jit(splx.stateful_solve_transform(search))()

The helper operator_at builds the symmetric operator of matrix + shift * I on one sparsity pattern, so the transform reuses a single analysis.

Backend-agnostic code

Type a routine against the splineax.SparseLinearSolver protocol and let the caller pick the solver:

def solve_many(
    solver: splx.SparseLinearSolver,
    operator: lx.AbstractLinearOperator,
    right_hand_sides: list[jax.Array],
):
    solution, state = splx.linear_solve(operator, right_hand_sides[0], solver)
    results = [solution.value]
    for b in right_hand_sides[1:]:
        solution, state = splx.linear_solve(operator, b, solver, state=state)
        results.append(solution.value)
    state.release()
    return results


solve_many(splx.AutoSparseLinearSolver(), operator, [b1, b2])

Profiling a stateful solve sequence

To debug what a solver actually reuses or rebuilds across a state-sequence, wrap your solves in splineax.create_solve_profile (or splineax.profile_solves for a jitted function), which records every analyze, factor, refactor, and solve step. See the solve profile guide for details.