Use the Sherman-Morison formula and the currently private jax.lax.linalg.cholesky_update function in inverses.py.
What value does this add?
It may be possible to increase the speed of the methods in inverse.py by utilising rank-1 updates. The primary drawback of such an update, via the Sherman-Morison formula, is that the numerical stability is unknown. However, when our matrix is symmetric, we are able to obtain better empirical stability and efficiency, by performing the update on the Cholesky factorisation, see https://timvieira.github.io/blog/post/2021/03/25/fast-rank-one-updates-to-matrix-inverse/.
JAX has a dedicated lowering for the Cholesky update, jax.lax.linalg.cholesky_update, but it doesn't appear to be part of the public API yet.
Is there an alternative you've considered?
No response
Additional context
Previous tests with rank-1 updates of the inverse of symmetric matrices yielded poor stability; However, these tests updated the matrix directly, and not the Cholesky factorisation.
What's the new feature?
Use the Sherman-Morison formula and the currently private
jax.lax.linalg.cholesky_updatefunction ininverses.py.What value does this add?
It may be possible to increase the speed of the methods in
inverse.pyby utilising rank-1 updates. The primary drawback of such an update, via the Sherman-Morison formula, is that the numerical stability is unknown. However, when our matrix is symmetric, we are able to obtain better empirical stability and efficiency, by performing the update on the Cholesky factorisation, see https://timvieira.github.io/blog/post/2021/03/25/fast-rank-one-updates-to-matrix-inverse/.JAX has a dedicated lowering for the Cholesky update,
jax.lax.linalg.cholesky_update, but it doesn't appear to be part of the public API yet.Is there an alternative you've considered?
No response
Additional context
Previous tests with rank-1 updates of the inverse of symmetric matrices yielded poor stability; However, these tests updated the matrix directly, and not the Cholesky factorisation.