Repository: probml/dynamax

Issue:
[Bug] Variance of inverse Wishart distribution
As alluded to in #407, the `InverseWishart.variance` function gives an error when the shape of the scale matrix is not `(3,3)`. Here is a minimal example:

```python
import jax.numpy as jnp
from dynamax.utils.distributions import InverseWishart

k = 2
𝜈 = 5.5  # 𝜈 > k + 3
Ψ = jnp.array([[20.712932, 25.124634],
       [25.124634, 32.814785]], dtype=jnp.float32)  # k x k

InverseWishart(df=𝜈, scale=Ψ).variance()
```

Which gives the error:

> File ~/.local/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:1601, in Distribution.variance(self, name, **kwargs)
>    1599 with self._name_and_control_scope(name):
>    1600   try:
> -> 1601     return self._variance(**kwargs)
>    1602   except NotImplementedError:
>    1603     try:
> 
> File ~/workspace/dynamax/dynamax/utils/distributions.py:114, in InverseWishart._variance(self)
>     112 dfs, scales = jnp.broadcast_arrays(jnp.array(self.df)[..., None, None], self.scale)
>     113 if scales.ndim == 2:
> --> 114     return _single_variance(dfs, scales)
>     115 else:
>     116     return vmap(_single_variance)(dfs, scales)
> 
> File ~/workspace/dynamax/dynamax/utils/distributions.py:108, in InverseWishart._variance.<locals>._single_variance(df, scale)
>     106 rows = jnp.arange(dim)[:, None].repeat(3, axis=1)
>     107 cols = jnp.arange(dim)[None, :].repeat(3, axis=0)
> --> 108 numer = (df - dim + 1) * scale**2 + (df - dim - 1) * diag[rows] * diag[cols]
>     109 denom = (df - dim) * (df - dim - 1)**2 * (df - dim - 3)
>     110 return numer / denom
> 
> 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 args = (other, self) if swap else (self, other)
>     572 if isinstance(other, _accepted_binop_types):
> --> 573   return binary_op(*args)
>     574 # Note: don't use isinstance here, because we don't want to raise for
>     575 # subclasses, e.g. NamedTuple objects that may override operators.
>     576 if type(other) in _rejected_binop_types:
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/numpy/ufunc_api.py:179, in ufunc.__call__(self, out, where, *args)
>     177   raise NotImplementedError(f"where argument of {self}")
>     178 call = self.__static_props['call'] or self._call_vectorized
> --> 179 return call(*args)
> 
>     [... skipping hidden 13 frame]
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py:1256, in multiply(x, y)
>    1230 """Multiply two arrays element-wise.
>    1231 
>    1232 JAX implementation of :obj:`numpy.multiply`. This is a universal function,
>    (...)
>    1253   Array([ 0, 10, 20, 30], dtype=int32)
>    1254 """
>    1255 x, y = promote_args("multiply", x, y)
> -> 1256 return lax.mul(x, y) if x.dtype != bool else lax.bitwise_and(x, y)
> 
>     [... skipping hidden 8 frame]
> 
> File ~/.local/lib/python3.12/site-packages/jax/_src/lax/lax.py:128, in _try_broadcast_shapes(name, *shapes)
>     126       result_shape.append(non_1s[0])
>     127     else:
> --> 128       raise TypeError(f'{name} got incompatible shapes for broadcasting: '
>     129                       f'{", ".join(map(str, map(tuple, shapes)))}.')
>     130 return tuple(result_shape)
> 
> TypeError: mul got incompatible shapes for broadcasting: (2, 2), (2, 3).

I think the problem in the code is that the terms $\psi_{ii} \psi_{jj}$ are not correctly calculated on [distributions.py:L106](https://github.com/probml/dynamax/blob/b02d8922e248268789842423eae9e4ec3bec269e/dynamax/utils/distributions.py#L106).

Kind regards,

Hylke

Discussion:
@hylkedonker: Fixed in 82719e9502cc9fec6141fb82f5ea7c2b09628db1. 

Solution patch:
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/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