probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
634 stars 70 forks source link

Addressing issue 347: time varying weights in LGSSM #356

Closed slinderman closed 5 months ago

slinderman commented 5 months ago

Talked to @mazzulab about issue with time-varying weights in LGSSM. Found that the sample code needs to pass through to lgssm_joint_sample in the inference code since the base SSM object doesn't support time-varying parameters in the sample code. Also dug into issues with fitting LGSSM and made some changes to his notebook to address those problems.