`psd_solve` leads to numerical errors

probml / dynamax

State Space Models library in JAX

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

MIT License

634 stars 70 forks source link

Closed calebweinreb closed 10 months ago

calebweinreb commented 11 months ago

The current implementation of psd_solve leads to numerical instability in some cases and could be improved:

Currently 1e-6 is added uniformly across the whole matrix. It would be more stable to just boost the diagonal.
Currently jnp.linalg.solve is used which doesn't take advantage of the PSD structure.