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])

Arguments:

• well_posed: whether to only handle well-posed systems or not, as discussed above.

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)

Arguments:

• rcond: the cutoff for handling zero entries on the diagonal. Defaults to machine precision times max(N, M), where (N, M) is the shape of the operator. (I.e. N is the output size and M is the input size.)

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)

Arguments:

• well_posed: if False, then singular operators are accepted, and the pseudoinverse solution is returned. If True then passing a singular operator will cause an error to be raised instead.
• rcond: the cutoff for handling zero entries on the diagonal. Defaults to machine precision times N, where N is the input (or output) size of the operator. Only used if well_posed=False

---

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: Callable[[PyTree], Shaped[Array, '']] = <function max_norm>, stabilise_every: Optional[int] = 10, max_steps: Optional[int] = None)

Arguments:

• rtol: Relative tolerance for terminating solve.
• atol: Absolute tolerance for terminating solve.
• norm: The norm to use when computing whether the error falls within the tolerance. Defaults to the max norm.
• stabilise_every: The conjugate gradient is an iterative method that produces candidate solutions \(x_1, x_2, \ldots\), and terminates once \(r_i = \| Ax_i - b \|\) is small enough. For computational efficiency, the values \(r_i\) are computed using other internal quantities, and not by directly evaluating the formula above. However, this computation of \(r_i\) is susceptible to drift due to limited floating-point precision. Every stabilise_every steps, then \(r_i\) is computed directly using the formula above, in order to stabilise the computation.
• max_steps: The maximum number of iterations to run the solver for. If more steps than this are required, then the solve is halted with a failure.

---

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: Callable[[PyTree], Shaped[Array, '']] = <function max_norm>, stabilise_every: Optional[int] = 10, max_steps: Optional[int] = None)

Arguments: - rtol: Relative tolerance for terminating solve. - atol: Absolute tolerance for terminating solve. - norm: The norm to use when computing whether the error falls within the tolerance. Defaults to the max norm. - stabilise_every: The conjugate gradient is an iterative method that produces candidate solutions \(x_1, x_2, \ldots\), and terminates once \(r_i = \| Ax_i - b \|\) is small enough. For computational efficiency, the values \(r_i\) are computed using other internal quantities, and not by directly evaluating the formula above. However, this computation of \(r_i\) is susceptible to drift due to limited floating-point precision. Every stabilise_every steps, then \(r_i\) is computed directly using the formula above, in order to stabilise the computation. - max_steps: The maximum number of iterations to run the solver for. If more steps than this are required, then the solve is halted with a failure.

---

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)

Arguments:

• rtol: Relative tolerance for terminating solve.
• atol: Absolute tolerance for terminating solve.
• norm: The norm to use when computing whether the error falls within the tolerance. Defaults to the max norm.
• max_steps: The maximum number of iterations to run the solver for. If more steps than this are required, then the solve is halted with a failure.

---

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)

Arguments:

• rtol: Relative tolerance for terminating solve.
• atol: Absolute tolerance for terminating solve.
• norm: The norm to use when computing whether the error falls within the tolerance. Defaults to the max norm.
• max_steps: The maximum number of iterations to run the solver for. If more steps than this are required, then the solve is halted with a failure.
• restart: Size of the Krylov subspace built between restarts. The returned solution is the projection of the true solution onto this subpsace, so this direclty bounds the accuracy of the algorithm. Default is 20.
• stagnation_iters: The maximum number of iterations for which the solver may not decrease. If more than stagnation_iters restarts are performed without sufficient decrease in the residual, the algorithm is halted.