Solvers

If you're not sure what to use, then pick lineax.AutoLinearSolver and it will automatically dispatch to an efficient solver depending on what structure your linear operator is declared to exhibit. (See the tags page.)

lineax.AbstractLinearSolver

lineax.AbstractLinearSolver

Abstract base class for all linear solvers.

init(self, operator: AbstractLinearOperator, options: dict[str, Any]) -> ~_SolverState abstractmethod

Do any initial computation on just the operator.

For example, an LU solver would compute the LU decomposition of the operator (and this does not require knowing the vector yet).

It is common to need to solve the linear system Ax=b multiple times in succession, with the same operator A and multiple vectors b. This method improves efficiency by making it possible to re-use the computation performed on just the operator.

Example

operator = lx.MatrixLinearOperator(...)
vector1 = ...
vector2 = ...
solver = lx.LU()
state = solver.init(operator, options={})
solution1 = lx.linear_solve(operator, vector1, solver, state=state)
solution2 = lx.linear_solve(operator, vector2, solver, state=state)

Arguments:

• operator: a linear operator.
• options: a dictionary of any extra options that the solver may wish to accept.

Returns:

A PyTree of arbitrary Python objects.

compute(self, state: ~_SolverState, vector: PyTree[Array], options: dict[str, Any]) -> tuple[PyTree[Array], RESULTS, dict[str, Any]] abstractmethod

Solves a linear system.

Arguments:

• state: as returned from lineax.AbstractLinearSolver.init.
• vector: the vector to solve against.
• options: a dictionary of any extra options that the solver may wish to accept. For example, lineax.CG accepts a preconditioner option.

Returns:

A 3-tuple of:

• The solution to the linear system.
• An integer indicating the success or failure of the solve. This is an integer which may be converted to a human-readable error message via lx.RESULTS[...].
• A dictionary of an extra statistics about the solve, e.g. the number of steps taken.

allow_dependent_columns(self, operator: AbstractLinearOperator) -> bool abstractmethod

Does this method ever produce non-NaN outputs for operators with linearly dependent columns? (Even if only sometimes.)

If True then a more expensive backward pass is needed, to account for the extra generality.

If you do not need to autodifferentiate through a custom linear solver then you simply define this method as

class MyLinearSolver(AbstractLinearsolver):
    def allow_dependent_columns(self, operator):
        raise NotImplementedError

Arguments:

• operator: a linear operator.

Returns:

Either True or False.

allow_dependent_rows(self, operator: AbstractLinearOperator) -> bool abstractmethod

Does this method ever produce non-NaN outputs for operators with linearly dependent rows? (Even if only sometimes)

If True then a more expensive backward pass is needed, to account for the extra generality.

If you do not need to autodifferentiate through a custom linear solver then you simply define this method as

class MyLinearSolver(AbstractLinearsolver):
    def allow_dependent_rows(self, operator):
        raise NotImplementedError

Arguments:

• operator: a linear operator.

Returns:

Either True or False.

transpose(self, state: ~_SolverState, options: dict[str, Any]) -> tuple[~_SolverState, dict[str, Any]] abstractmethod

Transposes the result of lineax.AbstractLinearSolver.init.

That is, it should be the case that

state_transpose, _ = solver.transpose(solver.init(operator, options), options)
state_transpose2 = solver.init(operator.T, options)

must be identical to each other.

It is relatively common (in particular when differentiating through a linear solve) to need to solve both Ax = b and A^T x = b. This method makes it possible to avoid computing both solver.init(operator) and solver.init(operator.T) if one can be cheaply computed from the other.

Arguments:

• state: as returned from solver.init.
• options: any extra options that were passed to solve.init.

Returns:

A 2-tuple of:

• The state of the transposed operator.
• The options for the transposed operator.

lineax.AutoLinearSolver (AbstractLinearSolver)

Automatically determines a good linear solver based on the structure of the operator.

• If well_posed=True:
  • If the operator is diagonal, then use lineax.Diagonal.
  • If the operator is tridiagonal, then use lineax.Tridiagonal.
  • If the operator is triangular, then use lineax.Triangular.
  • If the matrix is positive or negative definite, then use lineax.Cholesky.
  • Else use lineax.LU.

This is a good choice if you want to be certain that an error is raised for ill-posed systems.

• If well_posed=False:
  • If the operator is diagonal, then use lineax.Diagonal.
  • Else use lineax.SVD.

This is a good choice if you want to be certain that you can handle ill-posed systems.

• If well_posed=None:
  • If the operator is non-square, then use lineax.QR.
  • If the operator is diagonal, then use lineax.Diagonal.
  • If the operator is tridiagonal, then use lineax.Tridiagonal.
  • If the operator is triangular, then use lineax.Triangular.
  • If the matrix is positive or negative definite, then use lineax.Cholesky.
  • Else, use lineax.LU.

This is a good choice if your primary concern is computational efficiency. It will handle ill-posed systems as long as it is not computationally expensive to do so.

__init__(self, well_posed: Optional[bool])

select_solver(self, operator: AbstractLinearOperator) -> AbstractLinearSolver

Check which solver that lineax.AutoLinearSolver will dispatch to.

Arguments:

• operator: a linear operator.

Returns:

The linear solver that will be used.

---

lineax.LU (AbstractLinearSolver)

LU solver for linear systems.

This solver can only handle square nonsingular operators.

Least squares solvers

These are capable of solving ill-posed linear problems.

lineax.QR (AbstractLinearSolver)

QR solver for linear systems.

This solver can handle non-square operators.

This is usually the preferred solver when dealing with non-square operators.

Info

Note that whilst this does handle non-square operators, it still can only handle full-rank operators.

This is because JAX does not currently support a rank-revealing/pivoted QR decomposition, see issue #12897.

For such use cases, switch to lineax.SVD instead.

---

lineax.SVD (AbstractLinearSolver)

SVD solver for linear systems.

This solver can handle any operator, even nonsquare or singular ones. In these cases it will return the pseudoinverse solution to the linear system.

Equivalent to scipy.linalg.lstsq.

__init__(self, rcond: Optional[float] = None)

Info

In addition to these, lineax.Diagonal(well_posed=False) and lineax.NormalCG (below) also support ill-posed problems.

Structure-exploiting solvers

These require special structure in the operator. (And will throw an error if passed an operator without that structure.) In return, they are able to solve the linear problem much more efficiently.

lineax.Cholesky (AbstractLinearSolver)

Cholesky solver for linear systems. This is generally the preferred solver for positive or negative definite systems.

Equivalent to scipy.linalg.solve(..., assume_a="pos").

The operator must be square, nonsingular, and either positive or negative definite.

---

lineax.Diagonal (AbstractLinearSolver)

Diagonal solver for linear systems.

Requires that the operator be diagonal. Then \(Ax = b\), with \(A = diag[a]\), is solved simply by doing an elementwise division \(x = b / a\).

This solver can handle singular operators (i.e. diagonal entries with value 0).

__init__(self, well_posed: bool = False, rcond: Optional[float] = None)

---

lineax.Triangular (AbstractLinearSolver)

Triangular solver for linear systems.

The operator should either be lower triangular or upper triangular.

---

lineax.Tridiagonal (AbstractLinearSolver)

Tridiagonal solver for linear systems, using the Thomas algorithm.

Info

In addition to these, lineax.CG also requires special structure (positive or negative definiteness).

Iterative solvers

These solvers use only matrix-vector products, and do not require instantiating the whole matrix. This makes them good when used alongside e.g. lineax.JacobianLinearOperator or lineax.FunctionLinearOperator, which only provide matrix-vector products.

Warning

Note that lineax.BiCGStab and lineax.GMRES may fail to converge on some (typically non-sparse) problems.

lineax.CG (AbstractLinearSolver)

Conjugate gradient solver for linear systems.

The operator should be positive or negative definite.

Equivalent to scipy.sparse.linalg.cg.

This supports the following options (as passed to lx.linear_solve(..., options=...)).

• preconditioner: A positive definite lineax.AbstractLinearOperator to be used as preconditioner. Defaults to lineax.IdentityLinearOperator.
• y0: The initial estimate of the solution to the linear system. Defaults to all zeros.

Info

__init__(self, rtol: float, atol: float, norm: collections.abc.Callable[[PyTree], Shaped[Array, '']] = <function max_norm>, stabilise_every: Optional[int] = 10, max_steps: Optional[int] = None)

---

lineax.NormalCG (AbstractLinearSolver)

Conjugate gradient applied to the normal equations:

A^T A = A^T b

of a system of linear equations. Note that this squares the condition number, so it is not recommended. This is a fast but potentially inaccurate method, especially in 32 bit floating point precision.

This can handle nonsquare operators provided they are full-rank.

This supports the following options (as passed to lx.linear_solve(..., options=...)).

• preconditioner: A positive definite lineax.AbstractLinearOperator to be used as preconditioner. Defaults to lineax.IdentityLinearOperator.
• y0: The initial estimate of the solution to the linear system. Defaults to all zeros.

Info

__init__(self, rtol: float, atol: float, norm: collections.abc.Callable[[PyTree], Shaped[Array, '']] = <function max_norm>, stabilise_every: Optional[int] = 10, max_steps: Optional[int] = None)

---

lineax.BiCGStab (AbstractLinearSolver)

Biconjugate gradient stabilised method for linear systems.

The operator should be square.

Equivalent to jax.scipy.sparse.linalg.bicgstab.

This supports the following options (as passed to lx.linear_solve(..., options=...)).

• preconditioner: A positive definite lineax.AbstractLinearOperator to be used as a preconditioner. Defaults to lineax.IdentityLinearOperator.
• y0: The initial estimate of the solution to the linear system. Defaults to all zeros.

__init__(self, rtol: float, atol: float, norm: Callable = <function max_norm>, max_steps: Optional[int] = None)

---

lineax.GMRES (AbstractLinearSolver)

GMRES solver for linear systems.

The operator should be square.

Similar to jax.scipy.sparse.linalg.gmres.

This supports the following options (as passed to lx.linear_solve(..., options=...)).

• preconditioner: A positive definite lineax.AbstractLinearOperator to be used as preconditioner. Defaults to lineax.IdentityLinearOperator.
• y0: The initial estimate of the solution to the linear system. Defaults to all zeros.

__init__(self, rtol: float, atol: float, norm: Callable = <function max_norm>, max_steps: Optional[int] = None, restart: int = 20, stagnation_iters: int = 20)