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

¤

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