[BUG] Incorrect summary statistics in `LinearGaussianConjugateSSM.fit_blocked_gibbs`

There is a slight mistake in the Gibbs sampling step of the dynamics parameters of the LinearGaussianConjugateSSM model.

Briefly, defining the dynamics parameters as $W_{z} \equiv [F, B, b]^\intercal $, for this model the prior is given by:

Q, W_{z} \sim \mathcal{MN}\mathcal{W}^{-1} (M_{z0}, V_{z0}, \nu_{q0}, \psi_{q0})

and states are subsequently sampled according to

\begin{aligned}
z^{(i)}_t \sim \mathcal{N}(F z^{(i)}_{t-1} + B u^{(i)}_t + b, Q) \equiv \mathcal{N}(W_{z} x^{(i)}_{zt}, Q)\\
\end{aligned}

where we defined $x^{(i)}_{zt} = [z^{(i)}_{t-1}, u^{(i)}_t, 1]^\intercal$. Using the conjugacy of the matrix normal inverse-Wishart distribution, the dynamics parameters and variances, $Q, W_{z} \sim \mathcal{MN}\mathcal{W}^{-1} (M_{z}, V_{z}, \nu_{q}, \psi_{q})$, given the states are derived from the summary statistics:

\begin{aligned}
S_{z}^{xx} = V_{z0} + \sum_{i=1}^{m} \sum_{t=2}^T x^{(i)}_{zt} {x^{(i)}_{zt}}^\intercal\\
S_{z}^{yy} = M_{z0} V_{z0} M_{z0}^\intercal +  \sum_{i=1}^{m} \sum_{t=2}^T z^{(i)}_{t} {z^{(i)}_{t}}^\intercal\\
S_{z}^{yx} = M_{z0} V_{z0} +  \sum_{i=1}^{m}  \sum_{t=2}^Tz^{(i)}_{t} {x^{(i)}_{t}}^\intercal \\
S_{z}^{y|x} = S_{z}^{yx} - S_{z}^{yx} (S_{z}^{xx})^{-1} {S_{z}^{yx}}^\intercal
\end{aligned}

However, in the code the summary statistics are calculated based on $[z^{(i)}_{t}, u^{(i)}_t, 1]^\intercal $ instead of $x^{(i)}_{zt} = [z^{(i)}_{t-1}, u^{(i)}_t, 1]^\intercal $.
I had already corrected this mistake in PR #403. But if you want, I can make a separate PR with just this fix instead.

Kind regards,

Hylke

Start at dynamax/linear_gaussian_ssm/models.py around line 684 and inspect the summary-statistics calculation in LinearGaussianConjugateSSM.fit_blocked_gibbs. Compare the constructed regressor with the issue's x_zt definition and verify that the resulting statistics use the preceding latent state; PR #403 provides context for the already discussed correction.