Numerical instability PoissonHmm

[amald08]

Hi, I've been trying to fit a PoissonHmm with some simulations made in NEST https://github.com/nest/nest-simulator, due to the length of the simulation and high spike counts I thought using the refactor would optimise the running time with respect to the other version of the SSM where everything works but it can take a while (not too long to get worried). However when using the jax-ssm I encounter numerical instability in the EM update step, the following assertion is raised no matter the number of iterations or even taking a small sample of the data:

assert np.isfinite(lp), "NaNs in marginal log probability"

I was wondering if there is a known limitation with a large number of spike counts or something that I am missing.