EM yields NaN after some iterations

[atlaie]

Hi,

I'm currently using Dynamax to implement a Switching Linear Regression (LinearRegressionHMM in Dynamax) and I wanted to know a bit more about the fit_em function. Particularly, I'd like to train the SLR on some data that looks like this:

i.e., in part of its domain, it's constant. Here's a link with the .npz file containing this data (called "target" in that npz file), as well as the predictors ("predictors" in the npz) that I'm trying to use to call fit_em.

Also, I've tried adding some small noise so that the variance is non-null, but still, after some steps, the EM algorithm blows up:

Notice how after the 4th iteration (out of 100) it's just full of NaNs. When I try with SGD (fit_sgd) it works and converges to reasonable predictions.

Let me know if any further details are needed,

Thanks!