Stable psd solve - Githubissues

probml / dynamax

State Space Models library in JAX

https://probml.github.io/dynamax/

MIT License

634 stars 70 forks source link

Stable psd solve #335

Closed calebweinreb closed 10 months ago

calebweinreb commented 11 months ago

More stable version of psd_solve that addresses https://github.com/probml/dynamax/issues/334.

The input matrix is symmetrized
Only the diagonal is boosted
Uses cho_factor / cho_solve instead of linalg.solve

slinderman commented 10 months ago

Thanks! I believe @gileshd had originally implemented something similar, but we hit numerical stability issues in our tests. Yours may differ in that it explicitly symmetrizes as well as boosting the diagonal. I'll go ahead and merge!