DOCSEARCHbyAlgolia
Submit Docs

lineax

repository·main·Indexed 20 days ago

https://github.com/patrick-kidger/lineax

A JAX and Equinox library for linear solves and linear least squares. It supports PyTree-valued matrices, general linear operators (such as Jacobians), and provides numerically stable gradients. The library includes various operator types (Matrix, Diagonal, Tridiagonal, Function, etc.), a variety of solvers (iterative, structure-exploiting, and least squares), and an AutoLinearSolver for automatic dispatch based on operator tags.

Tokens
15.4K
Snippets
43
Records
72
Agent score
69%

What's inside lineax

  1. Compare Lineax with JAX core linear algebra

    main

    Lineax provides several advantages over standard jax.numpy and jax.scipy linear algebra operations:

    • New Solvers: Includes solvers like lineax.QR that are not in core JAX and can be faster than equivalents like jax.numpy.linalg.lstsq.
    • New Operators: Provides specialized operators such as lineax.JacobianLinearOperator.
    • Consistent API: Unifies the fragmented API found across jax.numpy, jax.scipy, and jax.scipy.sparse into a single consistent interface.
    • Numerical Stability: Offers more stable gradients, avoiding NaNs that can occur in some existing JAX implementations.
    • Performance: Provides faster compile and run times in specific scenarios.

  2. Use Structure-exploiting solvers for efficiency

    main

    If your linear operator has a known special structure, using a structure-exploiting solver can significantly improve performance. Note that these solvers will throw an error if the operator does not possess the required structure.

    Available Structure-exploiting Solvers:

    • lineax.Cholesky: Requires symmetric positive-definite structure.
    • lineax.Diagonal: For diagonal operators.
    • lineax.Triangular: For triangular operators.
    • lineax.Tridiagonal: For tridiagonal operators.
    • lineax.CG: Requires positive or negative definiteness.

  3. How linear operators work in Lineax

    main

    In Lineax, a linear operator is an abstraction used to represent a matrix $A$ in a linear system $Ax = b$. Instead of always materializing a full matrix, you can use different operator types depending on the structure of your problem to optimize for memory and computation. Many solvers (like lineax.CG) only require matrix-vector products, allowing you to use operators that never explicitly store the full matrix.

    Common operator types include:

    • lineax.MatrixLinearOperator: Stores the full matrix $A$ directly.
    • lineax.DiagonalLinearOperator: Stores only the diagonal of $A$ for efficiency.
    • lineax.TridiagonalLinearOperator: Optimized for tridiagonal matrices.
    • lineax.FunctionLinearOperator: Represents a linear operator via a function $F(x) = Ax$, avoiding matrix materialization entirely.
    • lineax.IdentityLinearOperator: Represents the identity matrix.
    • lineax.PyTreeLinearOperator: Handles operators acting on JAX PyTrees.
    • lineax.JacobianLinearOperator: Represents the Jacobian of a function.

  4. Use Iterative solvers for matrix-free operations

    main

    Iterative solvers are ideal when you do not want to (or cannot) instantiate a full matrix. They only require matrix-vector products, making them compatible with lineax.JacobianLinearOperator or lineax.FunctionLinearOperator.

    Available Iterative Solvers:

    • lineax.CG: Conjugate Gradient (requires specific structure like positive/negative definiteness).
    • lineax.BiCGStab: Biconjugate Gradient Stabilized.
    • lineax.GMRES: Generalized Minimal Residual method.
    • lineax.LSMR: Least Squares Minimal Residual (also functions as a least squares solver).

    Warning: lineax.BiCGStab and lineax.GMRES may fail to converge on certain problems, typically those that are not sparse.

  5. How tags work in Lineax

    main

    Lineax uses "tags" to mark linear operators with specific mathematical properties (e.g., being symmetric or positive semidefinite). These tags are an optional tool used to dispatch to more efficient solvers, which can improve both runtime and compile time.

    Warning: Tags are not verified against the actual underlying data. Lineax only checks the tag itself, not the matrix values. If you apply a tag to an operator that does not actually possess that property, lx.linear_solve may return an incorrect result without warning. This is analogous to the assume_a="pos" flag in scipy.linalg.solve.

    # Some rank-2 JAX array.
matrix = ...
# Some rank-1 JAX array.
vector = ...

# Declare that this matrix is positive semidefinite.
operator = lx.MatrixLinearOperator(matrix, lx.positive_semidefinite_tag)

# This tag is used to dispatch to a maximally-efficient linear solver.
# In this case, a Cholesky solver is used:
solution = lx.linear_solve(operator, vector)

# Whether operators are tagged can be checked:
assert lx.is_positive_semidefinite(operator)

  6. Choose a Least Squares solver for ill-posed problems

    main

    When dealing with ill-posed linear problems, use solvers designed for least squares. Available options include:

    • lineax.QR: QR decomposition-based solver.
    • lineax.SVD: Singular Value Decomposition-based solver.
    • lineax.Normal: Solves the normal equations.
    • lineax.LSMR: An iterative least squares method.
    • lineax.Diagonal: Can support ill-posed problems if initialized with well_posed=False.

  7. Represent a lower or upper triangular matrix

    main

    To represent a triangular matrix in Lineax, create a full matrix where the desired triangular part (lower or upper) contains your values and the opposite part contains zeros. You then wrap this matrix in a MatrixLinearOperator using the appropriate tag.

    This approach is the most efficient way to handle triangular matrices within JAX's ndarray-based programming model.

    # Example: Creating a lower triangular operator
operator = lx.MatrixLinearOperator(matrix, lx.lower_triangular_tag)

  8. Solve multiple systems of equations (AX = B)

    main

    Lineax solvers are designed to target single systems of linear equations ($Ax = b$). To solve multiple systems simultaneously (e.g., $AX = B$), use jax.vmap or equinox.filter_vmap to vectorize the solver.

    When using vmap, ensure you set the in_axes correctly so that the operator (the $A$ matrix) is treated as a constant (shared across the batch) while the right-hand side (the $B$ matrix) is mapped over the batch dimension.

    # Using equinox.filter_vmap
multi_linear_solve = eqx.filter_vmap(lx.linear_solve, in_axes=(None, 1))

# Or using standard jax.vmap
multi_linear_solve = jax.vmap(lx.linear_solve, in_axes=(None, 1))

  9. Use AutoLinearSolver for automatic solver dispatch

    main

    If you are unsure which specific solver to use, lineax.AutoLinearSolver is the recommended starting point. It automatically selects an efficient solver based on the structure and tags declared on your linear operator. To use it, ensure your linear operator has appropriate tags (see the tags documentation).

    # Example usage pattern
solver = lineax.AutoLinearSolver(operator)
result = solver.solve(rhs)

  11. Use tags to signal matrix structure to the solver

    main

    If you are using a general operator like lx.MatrixLinearOperator but know that the underlying matrix has specific properties (e.g., it is positive semi-definite), you can use tags to inform the solver.

    Lineax does not verify the mathematical properties of the matrix; tags act as a manual override to allow the AutoLinearSolver to select more efficient algorithms that exploit the signaled structure.

    import jax.numpy as jnp
import jax.random as jr
import lineax as lx

matrix = jr.normal(jr.PRNGKey(0), (4, 4))
# Signal that the matrix is positive semi-definite using a tag
operator = lx.MatrixLinearOperator(matrix.T @ matrix, lx.positive_semidefinite_tag)

# The solver will now select an algorithm optimized for PSD matrices
solution = lx.linear_solve(operator, vector)

  12. Compare Lineax to jax.numpy.linalg.lstsq

    main

    Lineax provides advantages over jax.numpy.linalg.lstsq in several areas:

    1. Speed (Forward Pass): In rectangular cases, using lx.QR() is significantly faster than the SVD-based approach used by the standard JAX implementation.
    2. Speed (Gradients): Lineax uses a more efficient autodifferentiation implementation, making gradient computations faster even when both libraries use the same underlying algorithm (like SVD).
    3. Correctness (Gradients): Lineax avoids certain NaN gradient issues present in core JAX's lstsq implementation.