0.1.1 release

[jpbrodrick89]

This is mostly a bug fix, documentation and under-the-hood performance improvement release with one new feature—the lx.invert transformation which produces an operator representing the inverse of a matrix. Use of coloring rules should make using implicit solvers in diffrax for tridiagonal Jacobian/FunctionLinearOperators at least an order of magnitude faster.

Breaking Changes

• Extraction of diagonal/tridiagonals of now leverages the promise of a matrix being tagged as diagonal/tridiagonal more heavily. If you have previously used the tag for an operator that you just wanted lineax to TREAT as diagonal/tridiagonal you may now get incorrect results. In most cases the right fix probably be to first manually extract (tri)diagonal and construct the (Tri)DiagonalOperator explicitly, please raise an issue if you need any further assistance.
• lineax.linear_solve now stop-gradient's automatically (update state docstring to be clearer; simplify its use by stop-gradient'ing it automatically #213), it is unlikely this will break any existing use-cases but may make manual stop-gradienting unecessary
• Removed AuxLinearOperator (Removed AuxLinearOperator. #203)

Features

• Add invert helper function to wrap lineax.linear_solve in FunctionLinearOperator. Materialising an inverse is now as simple as lx.invert(op).as_matrix(). (Add invert helper function to wrap linear_solve in FunctionLinearOperator #206)

Compatibility

• lineax v0.1.1 now requires JAX >= 0.10.0 which provides a lowering to LAPACK/cuSolver's` ormqr for more efficient QR solve adopted in Use ormqr for more efficient QR solve #219.

Bugfixes

• Fix derived tag check rules for composite operators (e.g. Composed/Neg/Mul/AddLinearOperator) (fix derived tag check rules #192)
• Linearisation of functions custom_vjp's are now supported by lineax.linearize(JacobianLinearOperator(f, x, jac="bwd")) by using jax.linear_transpose under the hood. (fix linearise for JacobianLinearOperator with jac=bwd and use linear_transpose in mv #191)
• Complex positive/negative semi-definite matrices no longer register as symmetric (fix is_symmetric for complex PSD/NSD matrices #200)
• lineax.LSMR no longer fails when initial residual is exactly zero. (HUGE thanks to @f0uriest for spotting this tricky and hard-to-spot bug Add shortcut to LSMR in event that initial residual is exactly zero #202)
• Differentiating through linear_solve's no longer differentiates through solver.init this means using solver's with no or incorrect jvp rule is now possible (prevent linear_solve calling jvp of solver.init #212)

Performance

• Coloring rules now used to massively speed up diagonal/tridiagonal extraction of tagged Jacobian/FunctionLinearOperators (Support "diagonal" primitives with no/slow JVP batch rule #164, Streamline tridiagonalisation of JacobianLinearOperator/FunctionLinearOperator using coloring method #165)
• Normal and iterative solvers now apply lineax.linearise under the hood to avoid multiple sequential AD passes (apply linearise in init for Normal and iterative solvers #198)
• Furthermore, lineax.Normal(lineax.Cholesky()) now materialises the inner operator before constructing the Gram matrix (materialise Normal(Cholesky()) inner solver #207)
• ComposedLinearOperator.as_matrix no longer materialises each matrix first but instead batches mv of the first operator over the second matrix (use mv in ComposedLinearOperator.as_matrix to avoid unecessary materialisation #196)
• JAX's ormqr now used for more efficient QR solves (Use ormqr for more efficient QR solve #219)

Documentation

• The lineax.LSMR iterative least square solver is now properly documented (added LSMR to docs #204) after @f0uriest's Add shortcut to LSMR in event that initial residual is exactly zero #202 bug-fixes make it more robust.

Other repo infra PR's not affecting Python package include #214, #216 and #218.

New Contributors

• @patrick-kidger-bot 🤖 made their first contribution in bump pymdown-extensions; pin pygments #216

Full Changelog: v0.1.0...v0.1.1

[jpbrodrick89]

Open to redrafts on the release notes, we can merge after diffrax 751.

[patrick-kidger]

🦭 of approval!

Congratulations on your first release for Lineax! 🎉