Fix off-by-one bugs

[nkbranigan]

• Fixes [BUG] Incorrect summary statistics in LinearGaussianConjugateSSM.fit_blocked_gibbs #416 and related off-by-one bugs for time-varying inputs and parameters
  • Problem: Most of Dynamax uses the convention that latent states at $t$ depend on inputs and parameters at $t$, but in some parts the latent states at $t$ depended on inputs and parameters at $t-1$
• Fixes additional bugs revealed during testing:
  • forecast outputs were all one time step too early
  • Parallel posterior sampling omitted inputs in its backward pass (previously would have crashed)
  • UKF indexed the wrong dimension for a time-varying dynamics covariance (previously would have crashed)
  • fixed_lag_smoother_hmm.ipynb used old API

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[slinderman]

Thanks for tracking all of these off-by-one errors down! Do the changes to the filtering code to special-case the last time step with predict=False lead to longer compilation times?

[nkbranigan]

Thank you for your review Scott! You’ve caught an issue I missed. The filtering changes do lead to longer compilation. For lgssm_filter, this PR is slower on my machine by ~50 ms with an outer jit (as in model.fit_em), and by ~500 ms with no outer jit (as in model.filter). The good news is that I think there’s a simple fix that restores compile times to what they were before this PR: #469.