dynamax
dynamax copied to clipboard
Add general info form message passing algorithms for linear Gaussian SSMs
The current versions of information form message passing for LGSSMs use a hybrid representation of dynamics parameters and information form emission parameters. It is often helpful to have routines that operate directly on the block tridiagonal precision matrix and the precision-weighted mean. Examples of such implementations are:
- https://github.com/lindermanlab/ssm-jax/blob/main/ssm/distributions/mvn_block_tridiag.py#L11
- https://colab.research.google.com/drive/1kY39c4UQNeG1CpviAql-iZ0dxMY8o8mn?usp=sharing
- https://github.com/lindermanlab/ssm/blob/master/ssm/messages.py#L1043
- https://github.com/lindermanlab/ssm/blob/master/ssm/messages.py#L1080
(I've also written information form samplers and smoothers in JAX, along the lines of the last two links.)
We should add similar inference routines to the LGSSM info inference code.