NAN returned in large data and large iteration number

Hi, I'm using LinearGaussianSSM with large time seires data: 73 emission dimension, 1000 length

I noticed that both fit_em() and posterior_predictive() will return me nan parameters or nan posterior predictions if I ran too many iterations, around 150 iterations. And if reduced the emission dimension or state dimension it can hold longer.

I felt like this is probably an overflow or underflow problem. I wonder if there is anything we can do in dynamax or jax to prevent it from returning nan.

A separate question: in dynamax.linear_gaussian_ssm.inference -> lgssm_posterior_sample() -> _step()

return state, state

And this is only used once and the first state, which should be exactly the same as second one, is discarded. I wonder if this is for later development or it's just an omit.

Start with dynamax.linear_gaussian_ssm.inference, especially lgssm_posterior_sample() and its _step(), and reproduce the reported fit_em() and posterior_predictive() behavior using the stated emission dimension, sequence length, and iteration count. Determine whether the large-data run produces NaNs and clarify the purpose of returning state twice in _step(); done means the numerical failure is addressed or characterized and the return-value question is resolved.