Repository: probml/dynamax

Issue:
No batch support for `LinearGaussianConjugateSSM.fit_blocked_gibbs`
Unlike suggested by the type hints of the `emissions: Float[Array, "nbatch num_timesteps emission_dim"]` argument in `LinearGaussianConjugateSSM.fit_blocked_gibbs`, this function can not fit multiple observations. 

A minimal reproducing example:
```Python
from itertools import count
from functools import partial
from jax import random, vmap
from dynamax.linear_gaussian_ssm import LinearGaussianConjugateSSM

state_dim = 2
emission_dim = 3
num_timesteps = 4
m_samples = 5
keys = map(random.PRNGKey, count())
m_keys = random.split(next(keys), num=m_samples)

model = LinearGaussianConjugateSSM(state_dim, emission_dim)
params, _ = model.initialize(next(keys))
_, y_obs = vmap(partial(model.sample, params, num_timesteps=num_timesteps))(m_keys)

model.fit_blocked_gibbs(next(keys), params, sample_size=6, emissions=y_obs)
```
Gives the error:

> ---------------------------------------------------------------------------
> TypeError                                 Traceback (most recent call last)
> Cell In[15], [line 16](vscode-notebook-cell:?execution_count=15&line=16)
>      [14](vscode-notebook-cell:?execution_count=15&line=14) m_keys = random.split(next(keys), num=m_samples)
>      [15](vscode-notebook-cell:?execution_count=15&line=15) _, y_obs = vmap(partial(model.sample, params, num_timesteps=num_timesteps))(m_keys)
> ---> [16](vscode-notebook-cell:?execution_count=15&line=16) model.fit_blocked_gibbs(next(keys), params, sample_size=6, emissions=y_obs)
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:692, in LinearGaussianConjugateSSM.fit_blocked_gibbs(self, key, initial_params, sample_size, emissions, inputs)
>     [690](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:690) for _ in progress_bar(range(sample_size)):
>     [691](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:691)     sample_of_params.append(current_params)
> --> [692](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:692)     current_params = one_sample(current_params, next(keys))
>     [694](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:694) return pytree_stack(sample_of_params)
> 
>     [... skipping hidden 13 frame]
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:681, in LinearGaussianConjugateSSM.fit_blocked_gibbs.<locals>.one_sample(_params, rng)
>     [679](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:679) rngs = jr.split(rng, 2)
>     [680](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:680) # Sample latent states
> --> [681](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:681) states = lgssm_posterior_sample(rngs[0], _params, emissions, inputs)
>     [682](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:682) # Sample parameters
>     [683](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/models.py:683) _stats = sufficient_stats_from_sample(states)
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:608, in lgssm_posterior_sample(key, params, emissions, inputs, jitter)
>     [605](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:605) inputs = jnp.zeros((num_timesteps, 0)) if inputs is None else inputs
>     [607](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:607) # Run the Kalman filter
> --> [608](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:608) filtered_posterior = lgssm_filter(params, emissions, inputs)
>     [609](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:609) ll, filtered_means, filtered_covs, *_ = filtered_posterior
>     [611](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:611) # Sample backward in time
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:380, in preprocess_args.<locals>.wrapper(*args, **kwargs)
>     [377](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:377) num_timesteps = len(emissions)
>     [378](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:378) full_params, inputs = preprocess_params_and_inputs(params, num_timesteps, inputs)
> --> [380](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:380) return f(full_params, emissions, inputs=inputs)
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:511, in lgssm_filter(params, emissions, inputs)
>     [509](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:509) # Run the Kalman filter
>     [510](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:510) carry = (0.0, params.initial.mean, params.initial.cov)
> --> [511](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:511) (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps))
>     [512](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:512) return PosteriorGSSMFiltered(marginal_loglik=ll, filtered_means=filtered_means, filtered_covariances=filtered_covs)
> 
>     [... skipping hidden 9 frame]
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:502, in lgssm_filter.<locals>._step(carry, t)
>     [499](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:499) ll += _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y)
>     [501](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:501) # Condition on this emission
> --> [502](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:502) filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y)
>     [504](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:504) # Predict the next state
>     [505](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:505) pred_mean, pred_cov = _predict(filtered_mean, filtered_cov, F, B, b, Q, u)
> 
> File ~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:302, in _condition_on(prior_mean, prior_cov, emission_matrix, input_weights, emission_bias, emission_cov, inpt, emission)
>     [299](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:299)     S = jnp.diag(emission_cov) + emission_matrix @ prior_cov @ emission_matrix.T
>     [301](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:301) residual = emission - input_weights @ inpt - emission_bias - emission_matrix @ prior_mean
> --> [302](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:302) mu_cond = prior_mean + K @ residual
>     [303](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:303) Sigma_cond = prior_cov - K @ S @ K.T
>     [304](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/workspace/dynamax/dynamax/linear_gaussian_ssm/inference.py:304) return mu_cond, symmetrize(Sigma_cond)
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1050, in _forward_operator_to_aval.<locals>.op(self, *args)
>    [1049](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1049) def op(self, *args):
> -> [1050](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1050)   return getattr(self.aval, f"_{name}")(self, *args)
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:573, in _defer_to_unrecognized_arg.<locals>.deferring_binary_op(self, other)
>     [571](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:571) args = (other, self) if swap else (self, other)
>     [572](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:572) if isinstance(other, _accepted_binop_types):
> --> [573](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:573)   return binary_op(*args)
>     [574](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:574) # Note: don't use isinstance here, because we don't want to raise for
>     [575](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:575) # subclasses, e.g. NamedTuple objects that may override operators.
>     [576](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:576) if type(other) in _rejected_binop_types:
> 
>     [... skipping hidden 14 frame]
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9161, in matmul(a, b, precision, preferred_element_type)
>    [9159](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9159) a = lax.squeeze(a, tuple(a_squeeze))
>    [9160](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9160) b = lax.squeeze(b, tuple(b_squeeze))
> -> [9161](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9161) out = lax.dot_general(
>    [9162](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9162)   a, b, (((ndim(a) - 1,), (ndim(b) - 1 - b_is_mat,)), (a_batch, b_batch)),
>    [9163](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9163)   precision=precision, preferred_element_type=preferred_element_type)
>    [9164](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9164) result = lax.transpose(out, perm)
>    [9165](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/numpy/lax_numpy.py:9165) return lax_internal._convert_element_type(result, preferred_element_type, output_weak_type)
> 
>     [... skipping hidden 7 frame]
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3342, in _dot_general_shape_rule(lhs, rhs, dimension_numbers, precision, preferred_element_type, out_type)
>    [3339](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3339) if not core.definitely_equal_shape(lhs_contracting_shape, rhs_contracting_shape):
>    [3340](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3340)   msg = ("dot_general requires contracting dimensions to have the same "
>    [3341](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3341)          "shape, got {} and {}.")
> -> [3342](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3342)   raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
>    [3344](https://file+.vscode-resource.vscode-cdn.net/home/hylke/workspace/dynamax/docs/notebooks/linear_gaussian_ssm/~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:3344) return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)
> 
> TypeError: dot_general requires contracting dimensions to have the same shape, got (3,) and (4,).

The function does work when the "nbatch" dimension is removed. For example, this works okay:
```python
# Remove the batch dimension.
model.fit_blocked_gibbs(next(keys), params, sample_size=6, emissions=y_obs[0])
```

Discussion:

Solution patch:
diff --git a/dynamax/linear_gaussian_ssm/models.py b/dynamax/linear_gaussian_ssm/models.py
index e949662..4e9ca04 100644
--- a/dynamax/linear_gaussian_ssm/models.py
+++ b/dynamax/linear_gaussian_ssm/models.py
@@ -7,12 +7,12 @@ import tensorflow_probability.substrates.jax.distributions as tfd
 
 from fastprogress.fastprogress import progress_bar
 from functools import partial
-from jax import jit, vmap
+from jax import jit, tree, vmap
 from jax.tree_util import tree_map
 from jaxtyping import Array, Float
 from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN
 from typing import Any, Optional, Tuple, Union, runtime_checkable
-from typing_extensions import Protocol 
+from typing_extensions import Protocol
 
 from dynamax.ssm import SSM
 from dynamax.linear_gaussian_ssm.inference import lgssm_joint_sample, lgssm_filter, lgssm_smoother, lgssm_posterior_sample
@@ -24,7 +24,7 @@ from dynamax.utils.bijectors import RealToPSDBijector
 from dynamax.utils.distributions import MatrixNormalInverseWishart as MNIW
 from dynamax.utils.distributions import NormalInverseWishart as NIW
 from dynamax.utils.distributions import mniw_posterior_update, niw_posterior_update
-from dynamax.utils.utils import pytree_stack, psd_solve
+from dynamax.utils.utils import ensure_array_has_batch_dim, pytree_stack, psd_solve
 
 @runtime_checkable
 class SuffStatsLGSSM(Protocol):
@@ -206,7 +206,7 @@ class LinearGaussianSSM(SSM):
                key: PRNGKeyT,
                num_timesteps: int,
                inputs: Optional[Float[Array, "num_timesteps input_dim"]] = None) \
-                -> Tuple[Float[Array, "num_timesteps state_dim"], 
+                -> Tuple[Float[Array, "num_timesteps state_dim"],
                          Float[Array, "num_timesteps emission_dim"]]:
         """Sample from the model.
         
@@ -357,7 +357,7 @@ class LinearGaussianSSM(SSM):
                 input_weights=params.emissions.input_weights,
                 cov=1e8 * jnp.ones(self.emission_dim)) # ignore dummy observatiosn
             )
-        
+
         dummy_emissions = jnp.zeros((num_forecast_timesteps, self.emission_dim))
         forecast_inputs = forecast_inputs if forecast_inputs is not None else \
             jnp.zeros((num_forecast_timesteps, 0))
@@ -367,7 +367,7 @@ class LinearGaussianSSM(SSM):
         H = params.emissions.weights
         b = params.emissions.bias
         R = params.emissions.cov if params.emissions.cov.ndim == 2 else jnp.diag(params.emissions.cov)
-        
+
         forecast_emissions = forecast_states.filtered_means @ H.T + b
         forecast_emissions_cov = H @ forecast_states.filtered_covariances @ H.T + R
         return forecast_states.filtered_means, \
@@ -662,18 +662,20 @@ class LinearGaussianConjugateSSM(LinearGaussianSSM):
         Returns:
             parameter object, where each field has `sample_size` copies as leading batch dimension.
         """
-        num_timesteps = len(emissions)
+        batch_emissions = ensure_array_has_batch_dim(emissions, self.emission_shape)
+        batch_inputs = ensure_array_has_batch_dim(inputs, self.inputs_shape)
 
-        if inputs is None:
-            inputs = jnp.zeros((num_timesteps, 0))
+        num_batches, num_timesteps = batch_emissions.shape[:2]
+
+        if batch_inputs is None:
+            batch_inputs = jnp.zeros((num_batches, num_timesteps, 0))
 
-        def sufficient_stats_from_sample(states):
+        def sufficient_stats_from_sample(y, inputs, states):
             """Convert samples of states to sufficient statistics."""
             inputs_joint = jnp.concatenate((inputs, jnp.ones((num_timesteps, 1))), axis=1)
             # Let xn[t] = x[t+1]          for t = 0...T-2
             x, xp, xn = states, states[:-1], states[1:]
             u, up = inputs_joint, inputs_joint[:-1]
-            y = emissions
 
             init_stats = (x[0], jnp.outer(x[0], x[0]), 1)
 
@@ -733,9 +735,13 @@ class LinearGaussianConjugateSSM(LinearGaussianSSM):
             """Sample a single set of states and compute their sufficient stats."""
             rngs = jr.split(rng, 2)
             # Sample latent states
-            states = lgssm_posterior_sample(rngs[0], _params, emissions, inputs)
+            batch_keys = jr.split(rngs[0], num=num_batches)
+            forward_backward_batched = vmap(partial(lgssm_posterior_sample, params=_params))
+            batch_states = forward_backward_batched(batch_keys, emissions=batch_emissions, inputs=batch_inputs)
+            _batch_stats = vmap(sufficient_stats_from_sample)(batch_emissions, batch_inputs, batch_states)
+            # Aggregate statistics from all observations.
+            _stats = tree.map(lambda x: jnp.sum(x, axis=0), _batch_stats)
             # Sample parameters
-            _stats = sufficient_stats_from_sample(states)
             return lgssm_params_sample(rngs[1], _stats)
 
 
diff --git a/dynamax/utils/distributions.py b/dynamax/utils/distributions.py
index ffa1d56..5b6330e 100644
--- a/dynamax/utils/distributions.py
+++ b/dynamax/utils/distributions.py
@@ -97,23 +97,12 @@ class InverseWishart(tfd.TransformedDistribution):
     def _variance(self):
         """Compute the marginal variance of each entry of the matrix.
         """
-        def _single_variance(df, scale):
-            """Compute the marginal variance of each entry of the matrix."""
-            assert scale.ndim == 2
-            assert df.shape == scale.shape
-            dim = scale.shape[-1]
-            diag = jnp.diag(scale)
-            rows = jnp.arange(dim)[:, None].repeat(3, axis=1)
-            cols = jnp.arange(dim)[None, :].repeat(3, axis=0)
-            numer = (df - dim + 1) * scale**2 + (df - dim - 1) * diag[rows] * diag[cols]
-            denom = (df - dim) * (df - dim - 1)**2 * (df - dim - 3)
-            return numer / denom
-
+        dim = self.scale.shape[-1]
         dfs, scales = jnp.broadcast_arrays(jnp.array(self.df)[..., None, None], self.scale)
-        if scales.ndim == 2:
-            return _single_variance(dfs, scales)
-        else:
-            return vmap(_single_variance)(dfs, scales)
+        Ψii_Ψjj = jnp.einsum('...ii,...jj->...ij', scales, scales)
+        numer = (dfs - dim + 1) * scales**2 + (dfs - dim - 1) * Ψii_Ψjj
+        denom = (dfs - dim) * (dfs - dim - 1)**2 * (dfs - dim - 3)
+        return numer / denom
 
 
 class NormalInverseWishart(tfd.JointDistributionSequential):


Test patch:
diff --git a/dynamax/linear_gaussian_ssm/models_test.py b/dynamax/linear_gaussian_ssm/models_test.py
index 50b5aff..bd15674 100644
--- a/dynamax/linear_gaussian_ssm/models_test.py
+++ b/dynamax/linear_gaussian_ssm/models_test.py
@@ -1,8 +1,12 @@
 """
 Tests for the linear Gaussian SSM models.
 """
+from functools import partial
+from itertools import count
 
 import pytest
+from jax import vmap
+import jax.numpy as jnp
 import jax.random as jr
 
 from dynamax.linear_gaussian_ssm import LinearGaussianSSM
@@ -29,3 +33,20 @@ def test_sample_and_fit(cls, kwargs, inputs):
     fitted_params, lps = model.fit_em(params, param_props, emissions, inputs=inputs, num_iters=3)
     assert monotonically_increasing(lps)
     fitted_params, lps = model.fit_sgd(params, param_props, emissions, inputs=inputs, num_epochs=3)
+
+def test_fit_blocked_gibbs_batched():
+    """
+    Test that the blocked Gibbs sampler works for multiple observations.
+    """
+    state_dim = 2
+    emission_dim = 3
+    num_timesteps = 4
+    m_samples = 5
+    keys = map(jr.PRNGKey, count())
+    m_keys = jr.split(next(keys), num=m_samples)
+
+    model = LinearGaussianConjugateSSM(state_dim, emission_dim)
+    params, _ = model.initialize(next(keys))
+    _, y_obs = vmap(partial(model.sample, params, num_timesteps=num_timesteps))(m_keys)
+
+    model.fit_blocked_gibbs(next(keys), params, sample_size=6, emissions=y_obs)
\ No newline at end of file
diff --git a/dynamax/utils/distributions_test.py b/dynamax/utils/distributions_test.py
index 836c687..552c9e6 100644
--- a/dynamax/utils/distributions_test.py
+++ b/dynamax/utils/distributions_test.py
@@ -6,6 +6,7 @@ import jax.random as jr
 
 from jax.tree_util import tree_map
 from jax.scipy.stats import norm
+import pytest
 from scipy.stats import invgamma
 from tensorflow_probability.substrates import jax as tfp
 
@@ -41,7 +42,8 @@ def test_inverse_wishart_log_prob(df=7.0, dim=3, scale_factor=3.0, n_samples=10)
     assert jnp.all(jnp.isfinite(lps))
 
 
-def test_inverse_wishart_sample(df=7.0, dim=3, scale_factor=3.0, n_samples=10000, num_std=6):
+@pytest.mark.parametrize("dim", [2, 3, 4])
+def test_inverse_wishart_sample(dim, df=7.0, scale_factor=3.0, n_samples=10000, num_std=6):
     """Test that the sample mean is within a (large) interval around the true mean.
     To determine the interval to 6 times the standard deviation of the Monte
     Carlo estimator.
@@ -57,6 +59,35 @@ def test_inverse_wishart_sample(df=7.0, dim=3, scale_factor=3.0, n_samples=10000
     mc_std = jnp.sqrt(iw.variance() / n_samples)
     assert jnp.allclose(samples.mean(axis=0), iw.mean(), atol=num_std * mc_std)
 
+def test_inverse_wishart_variance_vectorization():
+    """Test the vectorization of marginal variance calculation."""
+    p = 3
+    # Vectorize over two inverse-Wishart distributions.
+    Ψ1 = jnp.array([[ 3.915627  ,  0.05046278, -0.7466818 ],
+       [ 0.05046278,  5.368932  , -0.57225686],
+       [-0.7466818 , -0.57225686,  3.162733  ]], dtype=jnp.float32)
+    𝜈1 = 7.5  # >p + 3
+    Ψ2 = jnp.array([[ 4.3864717, -0.8869951, -0.9199044],
+       [-0.8869951,  5.07704  , -0.8494578],
+       [-0.9199044, -0.8494578,  4.1288185]], dtype=jnp.float32)
+    𝜈2 = 8.0  # >p + 3
+    assert all(jnp.linalg.eigvals(Ψ1) > 0)
+    assert all(jnp.linalg.eigvals(Ψ2) > 0)
+
+    # Make a (2, 1) batch shape to test vectorization over >1 leading axis.
+    𝜈 = jnp.stack([𝜈1, 𝜈2]) # Shape: (2,)
+    𝜈 = 𝜈[:, None]  # Shape: (2, 1)
+    Ψ = jnp.stack([Ψ1, Ψ2])  # Shape: (2, 3, 3)
+    Ψ = Ψ[:, None]  # Shape: (2, 1, 3, 3)
+    variance_composite = InverseWishart(df=𝜈, scale=Ψ).variance()
+    assert variance_composite.shape[-2:] == (p, p)
+
+    # Verify that vectorized and non-vectorized calculation gives the same variance.
+    variance1 = InverseWishart(df=𝜈1, scale=Ψ1).variance()
+    variance2 = InverseWishart(df=𝜈2, scale=Ψ2).variance()
+    jnp.allclose(variance_composite[0, 0], variance1)
+    jnp.allclose(variance_composite[1, 0], variance2)
+
 def test_inverse_wishart_sample_non_diagonal_scale(n_samples: int = 10_000, num_std=3):
     """Test sample mean of an inverse-Wishart distr. w/ non-diagonal scale matrix."""
     k = 2