Repository: probml/dynamax

Issue:
[Bug] `InverseWishart` gives incorrect samples
Hi,

I was playing around with the `InverseWishart` distribution. But when I make a Monte Carlo estimate of the average $\langle Σ \rangle$ with $Σ \sim \mathcal{W}^{-1}(\psi, \nu)$, it does not correspond to $\frac{\psi}{\nu - k - 1}$ (which is the [Inverse-Wishart mean](https://en.wikipedia.org/wiki/Inverse-Wishart_distribution) according to Wikipedia). Here is a minimal reproducing example:

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

key = jax.random.key(42)
n_samples = 10_000
k = 2
𝜈 = 3.5  # 𝜈 > k
Ψ = jnp.array([[20.712932, 25.124634],
       [25.124634, 32.814785]], dtype=jnp.float32)  # k x k
assert all(jnp.linalg.eigvals(Ψ) > 0)  # Is positive definite.

iw = InverseWishart(df=𝜈, scale=Ψ)
Σs = iw.sample(sample_shape=n_samples, seed=key)
actual_Σ_avg = jnp.mean(Σs, axis=0)
expected_Σ_avg = Ψ / (𝜈 - k - 1)
```
Possibly affects:
- `dynamax.utils.distributions.NormalInverseWishart`
- `dynamax.utils.distributions.MatrixNormalInverseWishart`
- `dynamax.linear_gaussian_ssm.LinearGaussianConjugateSSM`

Unrelated to this, the unit test `dynamax.utils.distributions_test.test_inverse_wishart_sample` works, but only when `dim=3` (and not for other values).

Kind regards,

Hylke

Hints from issue comments:
@hylkedonker: Fixed in PR #408

Solution patch:
diff --git a/dynamax/utils/distributions.py b/dynamax/utils/distributions.py
index 8bf8bfb..ffa1d56 100644
--- a/dynamax/utils/distributions.py
+++ b/dynamax/utils/distributions.py
@@ -42,8 +42,8 @@ class InverseWishart(tfd.TransformedDistribution):
         # Wishart distribution
         dim = scale.shape[-1]
         eye = jnp.broadcast_to(jnp.eye(dim), scale.shape)
-        cho_scale = jnp.linalg.cholesky(scale)
-        inv_scale_tril = solve_triangular(cho_scale, eye, lower=True)
+        inv_scale = psd_solve(A=scale, b=eye)
+        inv_scale_tril = jnp.linalg.cholesky(inv_scale)
 
         super().__init__(
             tfd.WishartTriL(df, scale_tril=inv_scale_tril),


Test patch:
diff --git a/dynamax/utils/distributions_test.py b/dynamax/utils/distributions_test.py
index ebad700..836c687 100644
--- a/dynamax/utils/distributions_test.py
+++ b/dynamax/utils/distributions_test.py
@@ -57,6 +57,29 @@ 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_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
+    𝜈 = 5.5  # 𝜈 > k
+    Ψ = jnp.array([[20.712932, 25.124634],
+        [25.124634, 32.814785]], dtype=jnp.float32)  # k x k
+    Ψ_diag = jnp.diagonal(Ψ)
+    assert all(jnp.linalg.eigvals(Ψ) > 0)  # Is positive definite.
+
+    iw = InverseWishart(df=𝜈, scale=Ψ)
+    Σs = iw.sample(sample_shape=n_samples, seed=jr.key(42))
+    actual_Σ_avg = jnp.mean(Σs, axis=0)
+
+    # Closed form expression of mean.
+    true_Σ_avg = Ψ / (𝜈 - k - 1)
+    # Closed form expression of variance.
+    numerator = (𝜈 - k + 1) * Ψ**2 + (𝜈 - k - 1) * jnp.outer(Ψ_diag, Ψ_diag)
+    denominator = (𝜈 - k) * (𝜈 - k - 1)**2 * (𝜈 - k - 3)
+    true_Σ_var = numerator / denominator
+
+    mc_std = jnp.sqrt(true_Σ_var / n_samples)
+    assert jnp.allclose(actual_Σ_avg, true_Σ_avg, atol=num_std * mc_std)
+
 
 def test_normal_inverse_wishart_mode(loc=0., mean_conc=1.0, df=7.0, dim=3, scale_factor=3.0):
     """