Issue with LinearGaussiamSMM with time-varying emissions weights

[spyectr]

I am working with the LinearGaussianSSM with time varying emission weights based on the example in Online linear regression using Kalman filtering. I found problems with 1) setting model parameters and sampling data which generates samples with wrong emission dimensions; 2) performing inference with EM, when we freeze emission weights (Trainable==False) to given time-varying values, the model still learns them but with the wrong dimension; 3) fit_sgd gives error messages. I illustrated these issues in this notebook.

[slinderman]

Hopefully we've fixed these issues through our discussion! See PR https://github.com/probml/dynamax/pull/356.

Short story:

• had to pass lgssm sample code to underlying inference code
• fit_em doesn't work with trainable=False
• fit_sgd does work, but there were a few bugs in the notebook you shared.