Your original approach might be right, I’ve just never seen it coded that way. It seems to me like it’s not propagating uncertainty about each state into the future correctly. A priori, I expected something like this, using a scan to construct latent states. That’s what the pymc-statespace package does (but it runs a special algorithm to try to compute the likelihood faster. It’s unclear if that’s still true following some of the recent upgrades to scan likelihood inference)