Repository: google/jax

Issue:
multivariate_normal unexpected variances
Hello!

multivariate_normal does not produce the variances that I would expect. I checked the output of the multivariate_normal function in Numpy, which does provide the behaviour I expected.

In the code, I perform a sanity check that the covariance matrix is indeed positive definite (symmetric, positive eigenvalues), get 10.000 samples and calculate the variances.

```
import jax.numpy as np
import numpy as onp
from jax import random

N = 10000

cov_list = [[ 0.19,  0., -0.13, 0.], 
            [ 0., 0.19, 0., -0.13], 
            [-0.13, 0., 0.19, 0.], 
            [ 0., -0.13, 0., 0.19]]

cov_np = np.array(cov_list)
cov_onp = onp.array(cov_list)

mean_np = np.zeros(4)
mean_onp = onp.zeros(4)

positive_eig = (onp.linalg.eigvals(cov_onp) > 0.).all()
symmetric = ((cov_onp - cov_onp.T) == 0.).all()
print('Covariance matrix is symmetric:', symmetric, '\nAll eigenvalues are positive:', positive_eig)

out_onp = onp.random.multivariate_normal(mean_onp, cov_onp, N)
var = out_onp.var(axis=0)
print('\nNumpy: Variance of samples \n', var)

key = random.PRNGKey(0)
out_np = random.multivariate_normal(key, mean=mean_np, cov=cov_np, shape=(N,))
var = out_np.var(axis=0)
print('\nJax: Variance of samples \n', var)
```

The resulting output is as follows:

```
Covariance matrix is symmetric: True 
All eigenvalues are positive: True

Numpy: Variance of samples 
 [0.19001905 0.18978024 0.18641562 0.18503367]

Jax: Variance of samples 
 [0.28731692 0.28114378 0.09892643 0.10218105]
```

I get this result on two different machines (Ubuntu 16 and 18).

**jax version used**

jax 0.1.53
jaxlib 0.1.36

Any idea why this happens?

Cheers
Christian


Discussion:
Thanks for catching this. [There's a missing transpose in `random.multivariate_normal`](https://github.com/google/jax/blob/96677d9c6f549b643fd0cdf5664da2f22f45fbaf/jax/random.py#L523) (I think because I assumed the result would be lower-triangular, like in `numpy.linalg`, but imported the `cholesky` from `scipy.linalg`.) More concerning is we didn't have a good enough test on `multivariate_normal` to catch this.

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -39,8 +39,8 @@
 from jax.lib import cuda_prng
 from jax import core
 from jax import abstract_arrays
+from jax.numpy.linalg import cholesky
 from jax.scipy.special import logit
-from jax.scipy.linalg import cholesky
 from jax.interpreters import batching
 from jax.interpreters import xla
 


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -409,6 +409,30 @@ def testMultivariateNormal(self, dim, dtype):
       # eigenvectors follow a standard normal distribution.
       self._CheckKolmogorovSmirnovCDF(whitened.ravel(), scipy.stats.norm().cdf)
 
+  def testMultivariateNormalCovariance(self):
+    # test code based on https://github.com/google/jax/issues/1869
+    N = 100000
+    cov = np.array([[ 0.19,  0.00, -0.13,  0.00],
+                   [  0.00,  0.29,  0.00, -0.23],
+                   [ -0.13,  0.00,  0.39,  0.00],
+                   [  0.00, -0.23,  0.00,  0.49]])
+    mean = np.zeros(4)
+
+    out_onp = onp.random.RandomState(0).multivariate_normal(mean, cov, N)
+
+    key = random.PRNGKey(0)
+    out_jnp = random.multivariate_normal(key, mean=mean, cov=cov, shape=(N,))
+
+    var_onp = out_onp.var(axis=0)
+    var_jnp = out_jnp.var(axis=0)
+    self.assertAllClose(var_onp, var_jnp, rtol=1e-2, atol=1e-2,
+                        check_dtypes=False)
+
+    var_onp = onp.cov(out_onp, rowvar=False)
+    var_jnp = onp.cov(out_jnp, rowvar=False)
+    self.assertAllClose(var_onp, var_jnp, rtol=1e-2, atol=1e-2,
+                        check_dtypes=False)
+
   def testIssue222(self):
     x = random.randint(random.PRNGKey(10003), (), 0, 0)
     assert x == 0