Lineax provides high-performance, JIT-compilable, and differentiable linear solvers for systems of the form Ax=b, built on JAX and designed to integrate seamlessly with Equinox. It enables defining custom linear operators and choosing various direct or iterative solvers. The current version is 0.1.0, indicating an early stage of development, with a release cadence that tends to follow updates in the broader JAX and Equinox ecosystem.
pip install lineaxVerified import paths — ran on the pinned version, not inferred.
This quickstart demonstrates how to define a simple `lx.Matrix` operator and use `lx.solve` to find the solution to a linear system Ax=b. This example leverages JAX arrays and lineax's core functionality.
Always wrap your core computation logic (especially loops or repeated operations) with `jax.jit` for optimal performance. Remember that `jax.jit` requires pure functions and static shapes.
Embrace the functional programming paradigm: all operations on JAX arrays return new arrays. Use JAX primitives like `jax.vmap`, `jax.lax.scan`, `jax.lax.fori_loop`, or `jax.lax.cond` for control flow and transformations.
Before solving, verify the properties of your operator (e.g., condition number for `lx.Matrix`). For ill-conditioned systems, consider using iterative solvers (e.g., GMRES) or implementing appropriate preconditioning techniques.
Leverage `jaxtyping` for type annotations in your functions to catch dimension and type errors early during development. This improves code readability and robustness, especially in complex JAX applications.
Ensure your matrix is non-singular by checking its determinant (`jnp.linalg.det`) or condition number. If the system is inherently singular, consider alternative problem formulations. For ill-conditioned but non-singular systems, iterative solvers or preconditioning may be more robust.
Replace Python control flow with JAX control flow primitives like `jax.lax.cond`, `jax.lax.while_loop`, or `jax.lax.fori_loop`. Alternatively, ensure that any control flow dependent on JAX array values is moved outside the `jax.jit` decorated function.
Carefully review the `in_structure()` and `out_structure()` methods of your `AbstractLinearOperator` and ensure that the input vector for `lx.solve` has a shape compatible with `operator.out_structure()`.