# 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)
```

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`

.

- If the operator is diagonal, then use

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`

.

- If the operator is diagonal, then use

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`

.

- If the operator is non-square, then use

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.