SYMBOL INDEX (525 symbols across 40 files)

FILE: benchmarks/gmres_fails_safely.py
  function tree_allclose (line 29) | def tree_allclose(x, y, *, rtol=1e-5, atol=1e-8):
  function make_problem (line 36) | def make_problem(mat_size: int, *, key):
  function benchmark_jax (line 44) | def benchmark_jax(mat_size: int, *, key):
  function benchmark_lx (line 63) | def benchmark_lx(mat_size: int, *, key):

FILE: benchmarks/lstsq_gradients.py
  function jax_solve (line 34) | def jax_solve(a):
  function lx_solve (line 39) | def lx_solve(a):

FILE: benchmarks/solver_speeds.py
  function tree_allclose (line 35) | def tree_allclose(x, y, *, rtol=1e-5, atol=1e-8):
  function base_wrapper (line 47) | def base_wrapper(a, b, solver):
  function jax_svd (line 61) | def jax_svd(a, b):
  function jax_gmres (line 66) | def jax_gmres(a, b):
  function jax_bicgstab (line 71) | def jax_bicgstab(a, b):
  function jax_cg (line 76) | def jax_cg(a, b):
  function jax_lu (line 81) | def jax_lu(matrix, vector):
  function jax_cholesky (line 85) | def jax_cholesky(matrix, vector):
  function jax_tridiagonal (line 89) | def jax_tridiagonal(matrix, vector):
  function create_problem (line 139) | def create_problem(solver, tags, size=3):
  function create_easy_iterative_problem (line 146) | def create_easy_iterative_problem(size, tags):
  function test_solvers (line 155) | def test_solvers(vmap_size, mat_size):

FILE: lineax/_misc.py
  function tree_where (line 23) | def tree_where(
  function resolve_rcond (line 30) | def resolve_rcond(rcond, n, m, dtype):
  function jacobian (line 41) | def jacobian(fn, in_size, out_size, holomorphic=False, has_aux=False, ja...
  function _asarray (line 58) | def _asarray(dtype, x):
  function _asarray_jvp (line 67) | def _asarray_jvp(dtype, x, tx):
  function default_floating_dtype (line 73) | def default_floating_dtype():
  function inexact_asarray (line 80) | def inexact_asarray(x):
  function complex_to_real_dtype (line 87) | def complex_to_real_dtype(dtype):
  function strip_weak_dtype (line 91) | def strip_weak_dtype(tree: PyTree) -> PyTree:
  function structure_equal (line 100) | def structure_equal(x, y) -> bool:

FILE: lineax/_norm.py
  function tree_dot (line 27) | def tree_dot(tree1: PyTree[ArrayLike], tree2: PyTree[ArrayLike]) -> Inex...
  function sum_squares (line 50) | def sum_squares(x: PyTree[ArrayLike]) -> Scalar:
  function two_norm (line 59) | def two_norm(x: PyTree[ArrayLike]) -> Scalar:
  function _two_norm (line 71) | def _two_norm(x: PyTree[ArrayLike]) -> Scalar:
  function _two_norm_jvp (line 86) | def _two_norm_jvp(x, tx):
  function rms_norm (line 104) | def rms_norm(x: PyTree[ArrayLike]) -> Scalar:
  function max_norm (line 123) | def max_norm(x: PyTree[ArrayLike]) -> Scalar:
  function _zero_grad_at_zero (line 143) | def _zero_grad_at_zero(x):
  function _zero_grad_at_zero_jvp (line 148) | def _zero_grad_at_zero_jvp(primals, tangents):

FILE: lineax/_operator.py
  function _frozenset (line 61) | def _frozenset(x: object | Iterable[object]) -> frozenset[object]:
  class AbstractLinearOperator (line 70) | class AbstractLinearOperator(eqx.Module):
    method __check_init__ (line 87) | def __check_init__(self):
    method mv (line 105) | def mv(
    method as_matrix (line 122) | def as_matrix(self) -> Inexact[Array, "a b"]:
    method transpose (line 137) | def transpose(self) -> "AbstractLinearOperator":
    method in_structure (line 150) | def in_structure(self) -> PyTree[jax.ShapeDtypeStruct]:
    method out_structure (line 161) | def out_structure(self) -> PyTree[jax.ShapeDtypeStruct]:
    method in_size (line 171) | def in_size(self) -> int:
    method out_size (line 183) | def out_size(self) -> int:
    method T (line 196) | def T(self) -> "AbstractLinearOperator":
    method __add__ (line 200) | def __add__(self, other) -> "AbstractLinearOperator":
    method __sub__ (line 205) | def __sub__(self, other) -> "AbstractLinearOperator":
    method __mul__ (line 210) | def __mul__(self, other) -> "AbstractLinearOperator":
    method __rmul__ (line 216) | def __rmul__(self, other) -> "AbstractLinearOperator":
    method __matmul__ (line 219) | def __matmul__(self, other) -> "AbstractLinearOperator":
    method __truediv__ (line 224) | def __truediv__(self, other) -> "AbstractLinearOperator":
    method __neg__ (line 230) | def __neg__(self) -> "AbstractLinearOperator":
  class MatrixLinearOperator (line 234) | class MatrixLinearOperator(AbstractLinearOperator):
    method __init__ (line 245) | def __init__(
    method mv (line 267) | def mv(self, vector):
    method as_matrix (line 273) | def as_matrix(self):
    method transpose (line 276) | def transpose(self):
    method in_structure (line 281) | def in_structure(self):
    method out_structure (line 285) | def out_structure(self):
  function _matmul (line 290) | def _matmul(matrix: ArrayLike, vector: ArrayLike) -> Array:
  function _tree_matmul (line 299) | def _tree_matmul(matrix: PyTree[ArrayLike], vector: PyTree[ArrayLike]) -...
  function _inexact_structure_impl2 (line 315) | def _inexact_structure_impl2(x):
  function _inexact_structure_impl (line 322) | def _inexact_structure_impl(x):
  function _inexact_structure (line 326) | def _inexact_structure(x: PyTree[jax.ShapeDtypeStruct]) -> PyTree[jax.Sh...
  class _Leaf (line 330) | class _Leaf:  # not a pytree
    method __init__ (line 331) | def __init__(self, value):
  class PyTreeLinearOperator (line 337) | class PyTreeLinearOperator(AbstractLinearOperator):
    method __init__ (line 373) | def __init__(
    method mv (line 426) | def mv(self, vector):
    method as_matrix (line 440) | def as_matrix(self):
    method transpose (line 459) | def transpose(self):
    method in_structure (line 481) | def in_structure(self):
    method out_structure (line 485) | def out_structure(self):
  class DiagonalLinearOperator (line 490) | class DiagonalLinearOperator(AbstractLinearOperator):
    method __init__ (line 502) | def __init__(self, diagonal: PyTree[ArrayLike]):
    method mv (line 509) | def mv(self, vector):
    method as_matrix (line 512) | def as_matrix(self):
    method transpose (line 515) | def transpose(self):
    method in_structure (line 518) | def in_structure(self):
    method out_structure (line 521) | def out_structure(self):
  class _NoAuxIn (line 525) | class _NoAuxIn(eqx.Module):
    method __call__ (line 529) | def __call__(self, x):
  class _Unwrap (line 533) | class _Unwrap(eqx.Module):
    method __call__ (line 536) | def __call__(self, x):
  class JacobianLinearOperator (line 541) | class JacobianLinearOperator(AbstractLinearOperator):
    method __init__ (line 575) | def __init__(
    method mv (line 622) | def mv(self, vector):
    method as_matrix (line 643) | def as_matrix(self):
    method transpose (line 646) | def transpose(self):
    method in_structure (line 657) | def in_structure(self):
    method out_structure (line 660) | def out_structure(self):
  class FunctionLinearOperator (line 666) | class FunctionLinearOperator(AbstractLinearOperator):
    method __init__ (line 679) | def __init__(
    method mv (line 707) | def mv(self, vector):
    method as_matrix (line 710) | def as_matrix(self):
    method transpose (line 713) | def transpose(self):
    method in_structure (line 727) | def in_structure(self):
    method out_structure (line 731) | def out_structure(self):
  class IdentityLinearOperator (line 738) | class IdentityLinearOperator(AbstractLinearOperator):
    method __init__ (line 746) | def __init__(
    method mv (line 769) | def mv(self, vector):
    method as_matrix (line 802) | def as_matrix(self):
    method transpose (line 812) | def transpose(self):
    method in_structure (line 815) | def in_structure(self):
    method out_structure (line 819) | def out_structure(self):
    method tags (line 824) | def tags(self):
  class TridiagonalLinearOperator (line 828) | class TridiagonalLinearOperator(AbstractLinearOperator):
    method __init__ (line 837) | def __init__(
    method mv (line 863) | def mv(self, vector):
    method as_matrix (line 869) | def as_matrix(self):
    method transpose (line 878) | def transpose(self):
    method in_structure (line 883) | def in_structure(self):
    method out_structure (line 887) | def out_structure(self):
  class TaggedLinearOperator (line 892) | class TaggedLinearOperator(AbstractLinearOperator):
    method __init__ (line 915) | def __init__(
    method mv (line 929) | def mv(self, vector):
    method as_matrix (line 932) | def as_matrix(self):
    method transpose (line 935) | def transpose(self):
    method in_structure (line 940) | def in_structure(self):
    method out_structure (line 943) | def out_structure(self):
  function _is_none (line 952) | def _is_none(x):
  class TangentLinearOperator (line 956) | class TangentLinearOperator(AbstractLinearOperator):
    method __check_init__ (line 964) | def __check_init__(self):
    method mv (line 967) | def mv(self, vector):
    method as_matrix (line 972) | def as_matrix(self):
    method transpose (line 977) | def transpose(self):
    method in_structure (line 984) | def in_structure(self):
    method out_structure (line 987) | def out_structure(self):
  class AddLinearOperator (line 991) | class AddLinearOperator(AbstractLinearOperator):
    method __check_init__ (line 1006) | def __check_init__(self):
    method mv (line 1012) | def mv(self, vector):
    method as_matrix (line 1020) | def as_matrix(self):
    method transpose (line 1023) | def transpose(self):
    method in_structure (line 1026) | def in_structure(self):
    method out_structure (line 1029) | def out_structure(self):
  class MulLinearOperator (line 1033) | class MulLinearOperator(AbstractLinearOperator):
    method mv (line 1048) | def mv(self, vector):
    method as_matrix (line 1051) | def as_matrix(self):
    method transpose (line 1054) | def transpose(self):
    method in_structure (line 1057) | def in_structure(self):
    method out_structure (line 1060) | def out_structure(self):
  class NegLinearOperator (line 1066) | class NegLinearOperator(AbstractLinearOperator):
    method mv (line 1079) | def mv(self, vector):
    method as_matrix (line 1082) | def as_matrix(self):
    method transpose (line 1085) | def transpose(self):
    method in_structure (line 1088) | def in_structure(self):
    method out_structure (line 1091) | def out_structure(self):
  class DivLinearOperator (line 1095) | class DivLinearOperator(AbstractLinearOperator):
    method mv (line 1110) | def mv(self, vector):
    method as_matrix (line 1114) | def as_matrix(self):
    method transpose (line 1117) | def transpose(self):
    method in_structure (line 1120) | def in_structure(self):
    method out_structure (line 1123) | def out_structure(self):
  class ComposedLinearOperator (line 1127) | class ComposedLinearOperator(AbstractLinearOperator):
    method __check_init__ (line 1145) | def __check_init__(self):
    method mv (line 1149) | def mv(self, vector):
    method as_matrix (line 1155) | def as_matrix(self):
    method transpose (line 1170) | def transpose(self):
    method in_structure (line 1173) | def in_structure(self):
    method out_structure (line 1176) | def out_structure(self):
  function _default_not_implemented (line 1197) | def _default_not_implemented(name: str, operator: AbstractLinearOperator...
  function linearise (line 1209) | def linearise(operator: AbstractLinearOperator) -> AbstractLinearOperator:
  function _ (line 1238) | def _(operator):
  function _ (line 1243) | def _(operator):
  function materialise (line 1266) | def materialise(operator: AbstractLinearOperator) -> AbstractLinearOpera...
  function _try_sparse_materialise (line 1320) | def _try_sparse_materialise(operator: AbstractLinearOperator) -> Abstrac...
  function _ (line 1344) | def _(operator):
  function _ (line 1351) | def _(operator):
  function _ (line 1356) | def _(operator):
  function _ (line 1372) | def _(operator):
  function diagonal (line 1397) | def diagonal(operator: AbstractLinearOperator) -> Shaped[Array, " size"]:
  function _leaf_from_keypath (line 1416) | def _leaf_from_keypath(pytree: PyTree, keypath: jtu.KeyPath) -> Array:
  function _ (line 1425) | def _(operator):
  function _ (line 1430) | def _(operator):
  function _ (line 1447) | def _(operator):
  function _ (line 1460) | def _(operator):
  function _ (line 1466) | def _(operator):
  function _ (line 1471) | def _(operator):
  function tridiagonal (line 1479) | def tridiagonal(
  function _ (line 1512) | def _(operator):
  function _ (line 1523) | def _(operator):
  function _ (line 1566) | def _(operator):
  function _ (line 1574) | def _(operator):
  function _ (line 1582) | def _(operator):
  function is_symmetric (line 1590) | def is_symmetric(operator: AbstractLinearOperator) -> bool:
  function _has_real_dtype (line 1607) | def _has_real_dtype(operator) -> bool:
  function _ (line 1625) | def _(operator):
  function _ (line 1639) | def _(operator):
  function _ (line 1644) | def _(operator):
  function _ (line 1649) | def _(operator):
  function is_diagonal (line 1657) | def is_diagonal(operator: AbstractLinearOperator) -> bool:
  function _ (line 1678) | def _(operator):
  function _ (line 1686) | def _(operator):
  function _ (line 1691) | def _(operator):
  function is_tridiagonal (line 1699) | def is_tridiagonal(operator: AbstractLinearOperator) -> bool:
  function _ (line 1720) | def _(operator):
  function _ (line 1727) | def _(operator):
  function has_unit_diagonal (line 1735) | def has_unit_diagonal(operator: AbstractLinearOperator) -> bool:
  function _ (line 1756) | def _(operator):
  function _ (line 1761) | def _(operator):
  function _ (line 1767) | def _(operator):
  function is_lower_triangular (line 1776) | def is_lower_triangular(operator: AbstractLinearOperator) -> bool:
  function _ (line 1797) | def _(operator):
  function _ (line 1803) | def _(operator):
  function _ (line 1808) | def _(operator):
  function is_upper_triangular (line 1816) | def is_upper_triangular(operator: AbstractLinearOperator) -> bool:
  function _ (line 1837) | def _(operator):
  function _ (line 1843) | def _(operator):
  function _ (line 1848) | def _(operator):
  function is_positive_semidefinite (line 1856) | def is_positive_semidefinite(operator: AbstractLinearOperator) -> bool:
  function _ (line 1877) | def _(operator):
  function _ (line 1882) | def _(operator):
  function _ (line 1888) | def _(operator):
  function is_negative_semidefinite (line 1897) | def is_negative_semidefinite(operator: AbstractLinearOperator) -> bool:
  function _ (line 1918) | def _(operator):
  function _ (line 1923) | def _(operator):
  function _ (line 1929) | def _(operator):
  function _ (line 1938) | def _(operator):
  function _ (line 1943) | def _(operator):
  function _ (line 1948) | def _(operator):
  function _ (line 1953) | def _(operator):
  function _ (line 1960) | def _(operator, transform=transform):
  function _ (line 1964) | def _(operator, transform=transform):
  function _ (line 1968) | def _(operator, transform=transform):
  function _ (line 1975) | def _(operator, transform=transform):
  function _ (line 1980) | def _(operator):
  function _ (line 1988) | def _(operator):
  function _ (line 1996) | def _(operator):
  function _ (line 2004) | def _(operator):
  function _ (line 2013) | def _(operator):
  function _ (line 2022) | def _(operator):
  function _ (line 2029) | def _(operator):
  function _ (line 2035) | def _(operator):
  function _ (line 2041) | def _(operator):
  function _ (line 2047) | def _(operator):
  function _ (line 2052) | def _(operator):
  function _ (line 2064) | def _(operator):
  function _ (line 2071) | def _(operator):
  function _ (line 2102) | def _(operator, check=check):
  function _ (line 2118) | def _(operator, check=check):
  function _ (line 2126) | def _(operator):
  class _ScalarSign (line 2130) | class _ScalarSign(enum.Enum):
  function _scalar_sign (line 2137) | def _scalar_sign(scalar) -> _ScalarSign:
  function _ (line 2154) | def _(operator):
  function _ (line 2166) | def _(operator):
  function _ (line 2180) | def _(operator):
  function _ (line 2190) | def _(operator):
  function _ (line 2201) | def _(operator):
  function _ (line 2206) | def _(operator):
  function _ (line 2222) | def _(operator, check=check, tag=tag):
  function _ (line 2237) | def _(operator, check=check):
  function _ (line 2242) | def _(operator):
  function _ (line 2254) | def _(operator, check=check):
  function _ (line 2260) | def _(operator):
  function _ (line 2267) | def _(operator):
  function _ (line 2278) | def _(operator):
  function _ (line 2283) | def _(operator):
  function conj (line 2296) | def conj(operator: AbstractLinearOperator) -> AbstractLinearOperator:
  function _ (line 2311) | def _(operator):
  function _ (line 2316) | def _(operator):
  function _ (line 2322) | def _(operator):
  function _ (line 2328) | def _(operator):
  function _ (line 2333) | def _(operator):
  function _ (line 2342) | def _(operator):
  function _ (line 2347) | def _(operator):
  function _ (line 2356) | def _(operator):
  function _ (line 2361) | def _(operator):
  function _ (line 2368) | def _(operator):
  function _ (line 2373) | def _(operator):
  function _ (line 2378) | def _(operator):
  function _ (line 2383) | def _(operator):
  function _ (line 2388) | def _(operator):

FILE: lineax/_solution.py
  class RESULTS (line 52) | class RESULTS(eqxi.Enumeration):
  class Solution (line 71) | class Solution(eqx.Module):

FILE: lineax/_solve.py
  function _to_shapedarray (line 66) | def _to_shapedarray(x):
  function _to_struct (line 73) | def _to_struct(x):
  function _assert_false (line 85) | def _assert_false(x):
  function _is_none (line 89) | def _is_none(x):
  function _sum (line 93) | def _sum(*args):
  function _linear_solve_impl (line 97) | def _linear_solve_impl(_, state, vector, options, solver, throw, *, chec...
  function _linear_solve_abstract_eval (line 133) | def _linear_solve_abstract_eval(operator, state, vector, options, solver...
  function _linear_solve_jvp (line 152) | def _linear_solve_jvp(primals, tangents):
  function _is_undefined (line 271) | def _is_undefined(x):
  function _assert_defined (line 275) | def _assert_defined(x):
  function _keep_undefined (line 279) | def _keep_undefined(v, ct):
  function _linear_solve_transpose (line 287) | def _linear_solve_transpose(inputs, cts_out):
  class AbstractLinearSolver (line 343) | class AbstractLinearSolver(eqx.Module, Generic[_SolverState]):
    method init (line 347) | def init(
    method compute (line 384) | def compute(
    method transpose (line 409) | def transpose(
    method conj (line 440) | def conj(
    method assume_full_rank (line 466) | def assume_full_rank(self) -> bool:
  function _lookup (line 496) | def _lookup(token) -> AbstractLinearSolver:
  class AutoLinearSolver (line 518) | class AutoLinearSolver(AbstractLinearSolver[_AutoLinearSolverState]):
    method _select_solver (line 555) | def _select_solver(self, operator: AbstractLinearOperator):
    method select_solver (line 602) | def select_solver(self, operator: AbstractLinearOperator) -> AbstractL...
    method init (line 615) | def init(self, operator, options) -> _AutoLinearSolverState:
    method compute (line 619) | def compute(
    method transpose (line 630) | def transpose(self, state: _AutoLinearSolverState, options: dict[str, ...
    method conj (line 637) | def conj(self, state: _AutoLinearSolverState, options: dict[str, Any]):
    method assume_full_rank (line 644) | def assume_full_rank(self):
  function linear_solve (line 657) | def linear_solve(
  function invert (line 809) | def invert(
  function stop_gradient_transpose (line 876) | def stop_gradient_transpose(ct, x):

FILE: lineax/_solver/bicgstab.py
  class BiCGStab (line 35) | class BiCGStab(AbstractLinearSolver[_BiCGStabState]):
    method __check_init__ (line 57) | def __check_init__(self):
    method init (line 70) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 78) | def compute(
    method transpose (line 207) | def transpose(self, state: _BiCGStabState, options: dict[str, Any]):
    method conj (line 214) | def conj(self, state: _BiCGStabState, options: dict[str, Any]):
    method assume_full_rank (line 221) | def assume_full_rank(self):

FILE: lineax/_solver/cg.py
  class CG (line 48) | class CG(AbstractLinearSolver[_CGState]):
    method __check_init__ (line 75) | def __check_init__(self):
    method init (line 88) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 114) | def compute(
    method transpose (line 229) | def transpose(
    method conj (line 239) | def conj(
    method assume_full_rank (line 249) | def assume_full_rank(self):
  function NormalCG (line 271) | def NormalCG(*args, **kwargs):

FILE: lineax/_solver/cholesky.py
  class Cholesky (line 34) | class Cholesky(AbstractLinearSolver[_CholeskyState]):
    method init (line 43) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 65) | def compute(
    method transpose (line 80) | def transpose(
    method conj (line 87) | def conj(
    method assume_full_rank (line 94) | def assume_full_rank(self):

FILE: lineax/_solver/diagonal.py
  class Diagonal (line 36) | class Diagonal(AbstractLinearSolver[_DiagonalState]):
    method init (line 48) | def init(
    method compute (line 66) | def compute(
    method transpose (line 85) | def transpose(self, state: _DiagonalState, options: dict[str, Any]):
    method conj (line 93) | def conj(self, state: _DiagonalState, options: dict[str, Any]):
    method assume_full_rank (line 104) | def assume_full_rank(self):

FILE: lineax/_solver/gmres.py
  class GMRES (line 39) | class GMRES(AbstractLinearSolver[_GMRESState]):
    method __check_init__ (line 65) | def __check_init__(self):
    method init (line 78) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 106) | def compute(
    method _gmres_compute (line 242) | def _gmres_compute(
    method _arnoldi_gram_schmidt (line 331) | def _arnoldi_gram_schmidt(
    method _normalise (line 401) | def _normalise(
    method transpose (line 415) | def transpose(self, state: _GMRESState, options: dict[str, Any]):
    method conj (line 422) | def conj(self, state: _GMRESState, options: dict[str, Any]):
    method assume_full_rank (line 429) | def assume_full_rank(self):

FILE: lineax/_solver/lsmr.py
  class LSMR (line 55) | class LSMR(AbstractLinearSolver[_LSMRState]):
    method __check_init__ (line 76) | def __check_init__(self):
    method init (line 91) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 94) | def compute(
    method _givens (line 361) | def _givens(self, a, b):
    method transpose (line 411) | def transpose(self, state: _LSMRState, options: dict[str, Any]):
    method conj (line 417) | def conj(self, state: _LSMRState, options: dict[str, Any]):
    method assume_full_rank (line 423) | def assume_full_rank(self):

FILE: lineax/_solver/lu.py
  class LU (line 37) | class LU(AbstractLinearSolver[_LUState]):
    method init (line 43) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 56) | def compute(
    method transpose (line 68) | def transpose(
    method conj (line 83) | def conj(
    method assume_full_rank (line 97) | def assume_full_rank(self):

FILE: lineax/_solver/misc.py
  function preconditioner_and_y0 (line 30) | def preconditioner_and_y0(
  function pack_structures (line 70) | def pack_structures(operator: AbstractLinearOperator) -> PackedStructures:
  function ravel_vector (line 79) | def ravel_vector(
  function unravel_solution (line 93) | def unravel_solution(
  function transpose_packed_structures (line 108) | def transpose_packed_structures(

FILE: lineax/_solver/normal.py
  function normal_preconditioner_and_y0 (line 31) | def normal_preconditioner_and_y0(options: dict[str, Any], tall: bool):
  class Normal (line 53) | class Normal(
    method init (line 105) | def init(self, operator, options):
    method compute (line 122) | def compute(
    method transpose (line 142) | def transpose(
    method conj (line 161) | def conj(
    method assume_full_rank (line 180) | def assume_full_rank(self):

FILE: lineax/_solver/qr.py
  class QR (line 37) | class QR(AbstractLinearSolver):
    method init (line 55) | def init(self, operator, options):
    method compute (line 67) | def compute(
    method transpose (line 96) | def transpose(self, state: _QRState, options: dict[str, Any]):
    method conj (line 107) | def conj(self, state: _QRState, options: dict[str, Any]):
    method assume_full_rank (line 117) | def assume_full_rank(self):

FILE: lineax/_solver/svd.py
  class SVD (line 38) | class SVD(AbstractLinearSolver[_SVDState]):
    method init (line 49) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 55) | def compute(
    method transpose (line 80) | def transpose(self, state: _SVDState, options: dict[str, Any]):
    method conj (line 88) | def conj(self, state: _SVDState, options: dict[str, Any]):
    method assume_full_rank (line 95) | def assume_full_rank(self):

FILE: lineax/_solver/triangular.py
  class Triangular (line 43) | class Triangular(AbstractLinearSolver[_TriangularState]):
    method init (line 49) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 68) | def compute(
    method transpose (line 87) | def transpose(self, state: _TriangularState, options: dict[str, Any]):
    method conj (line 101) | def conj(self, state: _TriangularState, options: dict[str, Any]):
    method assume_full_rank (line 114) | def assume_full_rank(self):

FILE: lineax/_solver/tridiagonal.py
  class Tridiagonal (line 36) | class Tridiagonal(AbstractLinearSolver[_TridiagonalState]):
    method init (line 41) | def init(self, operator: AbstractLinearOperator, options: dict[str, An...
    method compute (line 54) | def compute(
    method transpose (line 74) | def transpose(self, state: _TridiagonalState, options: dict[str, Any]):
    method conj (line 81) | def conj(self, state: _TridiagonalState, options: dict[str, Any]):
    method assume_full_rank (line 87) | def assume_full_rank(self):

FILE: lineax/_tags.py
  class _HasRepr (line 16) | class _HasRepr:
    method __init__ (line 17) | def __init__(self, string: str):
    method __repr__ (line 20) | def __repr__(self):
  function _ (line 47) | def _(tags: frozenset[object], tag=tag):
  function _ (line 53) | def _(tags: frozenset[object]):
  function _ (line 59) | def _(tags: frozenset[object]):
  function transpose_tags (line 64) | def transpose_tags(tags: frozenset[object]):

FILE: tests/conftest.py
  function getkey (line 26) | def getkey():

FILE: tests/helpers.py
  function _construct_matrix_impl (line 31) | def _construct_matrix_impl(
  function construct_matrix (line 75) | def construct_matrix(getkey, solver, tags, num=1, *, size=3, dtype=jnp.f...
  function construct_singular_matrix (line 86) | def construct_singular_matrix(getkey, solver, tags, num=1, dtype=jnp.flo...
  function construct_poisson_matrix (line 101) | def construct_poisson_matrix(size, dtype=jnp.float64):
  function _transpose (line 142) | def _transpose(operator, matrix):
  function _linearise (line 146) | def _linearise(operator, matrix):
  function _materialise (line 150) | def _materialise(operator, matrix):
  function params (line 157) | def params(only_pseudo):
  function tree_allclose (line 177) | def tree_allclose(x, y, *, rtol=1e-5, atol=1e-8):
  function has_tag (line 181) | def has_tag(tags, tag):
  function _operators_append (line 188) | def _operators_append(x):
  function make_matrix_operator (line 194) | def make_matrix_operator(getkey, matrix, tags):
  function make_trivial_pytree_operator (line 199) | def make_trivial_pytree_operator(getkey, matrix, tags):
  function make_function_operator (line 206) | def make_function_operator(getkey, matrix, tags):
  function make_jac_operator (line 214) | def make_jac_operator(getkey, matrix, tags):
  function make_jacfwd_operator (line 228) | def make_jacfwd_operator(getkey, matrix, tags):
  function make_jacrev_operator (line 242) | def make_jacrev_operator(getkey, matrix, tags):
  function make_trivial_diagonal_operator (line 278) | def make_trivial_diagonal_operator(getkey, matrix, tags):
  function make_identity_operator (line 285) | def make_identity_operator(getkey, matrix, tags):
  function make_tridiagonal_operator (line 291) | def make_tridiagonal_operator(getkey, matrix, tags):
  function make_add_operator (line 312) | def make_add_operator(getkey, matrix, tags):
  function make_mul_operator (line 322) | def make_mul_operator(getkey, matrix, tags):
  function make_composed_operator (line 328) | def make_composed_operator(getkey, matrix, tags):
  function finite_difference_jvp (line 342) | def finite_difference_jvp(fn, primals, tangents):
  function jvp_jvp_impl (line 355) | def jvp_jvp_impl(

FILE: tests/test_adjoint.py
  function test_adjoint (line 20) | def test_adjoint(make_operator, dtype, getkey):
  function test_functional_pytree_adjoint (line 58) | def test_functional_pytree_adjoint():
  function test_functional_pytree_adjoint_complex (line 68) | def test_functional_pytree_adjoint_complex():
  function test_preconditioner_adjoint (line 95) | def test_preconditioner_adjoint(solver):

FILE: tests/test_invert.py
  function _well_conditioned_matrix (line 24) | def _well_conditioned_matrix(getkey, size=3, dtype=jnp.float64):
  function _well_conditioned_psd_matrix (line 32) | def _well_conditioned_psd_matrix(getkey, size=3, dtype=jnp.float64):
  function test_mv (line 41) | def test_mv(getkey):
  function test_composition_identity (line 52) | def test_composition_identity(getkey):
  function test_double_inverse (line 63) | def test_double_inverse(getkey):
  function test_pseudoinverse_overdetermined (line 77) | def test_pseudoinverse_overdetermined(getkey):
  function test_pseudoinverse_underdetermined (line 88) | def test_pseudoinverse_underdetermined(getkey):
  function test_solver_cholesky (line 102) | def test_solver_cholesky(getkey):
  function test_solver_cg (line 113) | def test_solver_cg(getkey):
  function test_vmap (line 127) | def test_vmap(getkey):
  function test_grad_wrt_vector (line 141) | def test_grad_wrt_vector(getkey):
  function test_jvp_wrt_vector (line 156) | def test_jvp_wrt_vector(getkey):
  function test_grad_wrt_operator (line 172) | def test_grad_wrt_operator(getkey):
  function test_jvp_wrt_operator (line 190) | def test_jvp_wrt_operator(getkey):

FILE: tests/test_jvp.py
  function test_jvp (line 46) | def test_jvp(

FILE: tests/test_jvp_jvp1.py
  function _clear_cache (line 31) | def _clear_cache():
  function test_jvp_jvp (line 40) | def test_jvp_jvp(

FILE: tests/test_jvp_jvp2.py
  function _clear_cache (line 31) | def _clear_cache():
  function test_jvp_jvp (line 40) | def test_jvp_jvp(

FILE: tests/test_lsmr.py
  function test_ill_conditioned (line 13) | def test_ill_conditioned():
  function test_zero_rhs (line 20) | def test_zero_rhs():
  function test_damp_regularizes (line 34) | def test_damp_regularizes():
  function test_damp (line 47) | def test_damp():

FILE: tests/test_misc.py
  function test_inexact_asarray_no_copy (line 22) | def test_inexact_asarray_no_copy():
  function test_inexact_asarray_jvp (line 30) | def test_inexact_asarray_jvp():
  function test_zero_matrix (line 37) | def test_zero_matrix(dtype):

FILE: tests/test_norm.py
  function _square (line 23) | def _square(x):
  function _two_norm (line 27) | def _two_norm(x):
  function _rms_norm (line 31) | def _rms_norm(x):
  function _max_norm (line 35) | def _max_norm(x):
  function test_nonzero (line 39) | def test_nonzero():
  function test_zero (line 72) | def test_zero():
  function test_complex (line 85) | def test_complex():
  function test_size_zero (line 102) | def test_size_zero():

FILE: tests/test_operator.py
  function test_ops (line 36) | def test_ops(make_operator, getkey, dtype):
  function test_structures_vector (line 87) | def test_structures_vector(make_operator, getkey):
  function _setup (line 111) | def _setup(getkey, matrix, tag: object | frozenset[object] = frozenset()):
  function _assert_except_diag (line 131) | def _assert_except_diag(cond_fun, operators, flip_cond):
  function test_linearise (line 143) | def test_linearise(dtype, getkey):
  function test_materialise (line 163) | def test_materialise(dtype, getkey):
  function test_materialise_large (line 170) | def test_materialise_large(dtype, getkey):
  function test_diagonal (line 177) | def test_diagonal(dtype, getkey):
  function test_tridiagonal (line 194) | def test_tridiagonal(dtype, getkey):
  function test_is_symmetric (line 239) | def test_is_symmetric(dtype, getkey):
  function test_is_diagonal (line 250) | def test_is_diagonal(dtype, getkey):
  function test_is_diagonal_scalar (line 261) | def test_is_diagonal_scalar(dtype, getkey):
  function test_is_diagonal_tridiagonal (line 269) | def test_is_diagonal_tridiagonal(dtype, getkey):
  function test_has_unit_diagonal (line 277) | def test_has_unit_diagonal(dtype, getkey):
  function test_is_lower_triangular (line 289) | def test_is_lower_triangular(dtype, getkey):
  function test_is_upper_triangular (line 300) | def test_is_upper_triangular(dtype, getkey):
  function test_is_positive_semidefinite (line 311) | def test_is_positive_semidefinite(dtype, getkey):
  function test_is_negative_semidefinite (line 326) | def test_is_negative_semidefinite(dtype, getkey):
  function test_is_tridiagonal (line 341) | def test_is_tridiagonal(dtype, getkey):
  function test_tangent_as_matrix (line 354) | def test_tangent_as_matrix(dtype, getkey):
  function test_materialise_function_linear_operator (line 377) | def test_materialise_function_linear_operator(dtype, getkey):
  function test_pytree_transpose (line 400) | def test_pytree_transpose(dtype, getkey):
  function test_diagonal_tangent (line 420) | def test_diagonal_tangent():
  function test_identity_with_different_structures (line 432) | def test_identity_with_different_structures():
  function test_identity_with_different_structures_complex (line 461) | def test_identity_with_different_structures_complex():
  function test_zero_pytree_as_matrix (line 491) | def test_zero_pytree_as_matrix(dtype):
  function test_jacrev_operator (line 498) | def test_jacrev_operator():

FILE: tests/test_singular.py
  function test_small_singular (line 39) | def test_small_singular(make_operator, solver, tags, ops, getkey, dtype):
  function test_bicgstab_breakdown (line 57) | def test_bicgstab_breakdown(getkey, dtype):
  function test_gmres_stagnation_or_breakdown (line 76) | def test_gmres_stagnation_or_breakdown(getkey, dtype):
  function test_nonsquare_pytree_operator1 (line 113) | def test_nonsquare_pytree_operator1(solver):
  function test_nonsquare_pytree_operator2 (line 136) | def test_nonsquare_pytree_operator2(solver):
  function test_nonsquare_mat_vec (line 162) | def test_nonsquare_mat_vec(solver, full_rank, jvp, wide, dtype, getkey):
  function test_nonsquare_vec (line 213) | def test_nonsquare_vec(solver, full_rank, jvp, wide, dtype, getkey):
  function test_iterative_singular (line 256) | def test_iterative_singular(getkey, solver, tags, use_state, make_operat...

FILE: tests/test_solve.py
  function test_gmres_large_dense (line 24) | def test_gmres_large_dense(getkey):
  function test_nontrivial_pytree_operator (line 41) | def test_nontrivial_pytree_operator():
  function test_nontrivial_diagonal_operator (line 51) | def test_nontrivial_diagonal_operator():
  function test_mixed_dtypes (line 65) | def test_mixed_dtypes(solver):
  function test_mixed_dtypes_complex (line 78) | def test_mixed_dtypes_complex(solver):
  function test_mixed_dtypes_complex_real (line 91) | def test_mixed_dtypes_complex_real(solver):
  function test_mixed_dtypes_triangular (line 103) | def test_mixed_dtypes_triangular():
  function test_mixed_dtypes_complex_triangular (line 115) | def test_mixed_dtypes_complex_triangular():
  function test_mixed_dtypes_complex_real_triangular (line 127) | def test_mixed_dtypes_complex_real_triangular():
  function test_ad_closure_function_linear_operator (line 139) | def test_ad_closure_function_linear_operator(getkey):
  function test_grad_vmap_symbolic_cotangent (line 156) | def test_grad_vmap_symbolic_cotangent():
  function test_iterative_solver_max_steps_only (line 184) | def test_iterative_solver_max_steps_only(solver):
  function test_solver_init_not_differentiated (line 197) | def test_solver_init_not_differentiated(getkey):
  function test_nonfinite_input (line 249) | def test_nonfinite_input():

FILE: tests/test_transpose.py
  class TestTranspose (line 25) | class TestTranspose:
    method assert_transpose_fixture (line 27) | def assert_transpose_fixture(_):
    method test_transpose (line 43) | def test_transpose(
    method test_pytree_transpose (line 54) | def test_pytree_transpose(_, assert_transpose_fixture):  # pyright: ig...

FILE: tests/test_vmap.py
  function test_vmap (line 43) | def test_vmap(
  function test_grad_vmap_basic (line 93) | def test_grad_vmap_basic(getkey):
  function test_grad_vmap_advanced (line 110) | def test_grad_vmap_advanced(getkey):

FILE: tests/test_vmap_jvp.py
  function test_vmap_jvp (line 45) | def test_vmap_jvp(

FILE: tests/test_vmap_vmap.py
  function test_vmap_vmap (line 44) | def test_vmap_vmap(

FILE: tests/test_well_posed.py
  function test_small_wellposed (line 34) | def test_small_wellposed(make_operator, solver, tags, ops, getkey, dtype):
  function test_pytree_wellposed (line 57) | def test_pytree_wellposed(solver, getkey, dtype):