Repository: google/jax

Issue:
`inf` using multivariate_normal.logpdf with multiple batches
`multivariate_normal.logpdf` gives `inf`s when there are multiple batches:

```
>>> from jax import scipy
>>> import jax.numpy as np
>>> from jax.scipy.stats import multivariate_normal
>>> multivariate_normal.logpdf(np.array([[1. ,2.],
...                                      [2., 3.]]),
...                            np.array([[1., 2.],
...                                      [2., 3.]]),
...                            np.array([[[1., 0.2],
...                                       [0.2, 1.]],
...                                      [[1., 0.3],
...                                       [0.3, 1.]]]))
DeviceArray([inf, inf], dtype=float32)
```









Discussion:
Thanks for the report!

I beleive the issue is that the diagonal of the decomposed covariance is not properly handled in the block case. It looks like changing this line: https://github.com/google/jax/blob/bb0750f31a35e2661834e8b532781785fe63496e/jax/_src/scipy/stats/multivariate_normal.py#L46-L47
to this:
```python
return (-1/2 * jnp.einsum('...i,...i->...', y, y) - n/2*np.log(2*np.pi)
        - jnp.log(L.diagonal(axis1=-1, axis2=-2)).sum(-1))
```
will make things work correctly.

I can prepare a fix with tests sometime tomorrow if nobody gets to it before that.

Resolving patch:
diff --git a/jax/_src/scipy/stats/multivariate_normal.py b/jax/_src/scipy/stats/multivariate_normal.py
--- a/jax/_src/scipy/stats/multivariate_normal.py
+++ b/jax/_src/scipy/stats/multivariate_normal.py
@@ -44,7 +44,7 @@ def logpdf(x, mean, cov, allow_singular=None):
       L = lax.linalg.cholesky(cov)
       y = lax.linalg.triangular_solve(L, x - mean, lower=True, transpose_a=True)
       return (-1/2 * jnp.einsum('...i,...i->...', y, y) - n/2*np.log(2*np.pi)
-              - jnp.log(L.diagonal()).sum())
+              - jnp.log(L.diagonal(axis1=-1, axis2=-2)).sum(-1))
 
 @_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)
 def pdf(x, mean, cov):


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -20,6 +20,7 @@
 import numpy as np
 import scipy.stats as osp_stats
 
+from jax import api
 from jax import test_util as jtu
 from jax.scipy import stats as lsp_stats
 from jax.scipy.special import expit
@@ -470,6 +471,24 @@ def args_maker():
     self._CompileAndCheck(lsp_stats.multivariate_normal.logpdf, args_maker,
                           rtol=1e-4, atol=1e-4)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_ndim={}_nbatch={}_dtype={}".format(ndim, nbatch, dtype.__name__),
+       "ndim": ndim, "nbatch": nbatch, "dtype": dtype}
+      for ndim in [2, 3]
+      for nbatch in [1, 3, 5]
+      for dtype in jtu.dtypes.floating))
+  def testMultivariateNormalLogpdfBatch(self, ndim, nbatch, dtype):
+    # Regression test for #5570
+    rng = jtu.rand_default(self.rng())
+    x = rng((nbatch, ndim), dtype)
+    mean = 5 * rng((nbatch, ndim), dtype)
+    factor = rng((nbatch, ndim, 2 * ndim), dtype)
+    cov = factor @ factor.transpose(0, 2, 1)
+
+    result1 = lsp_stats.multivariate_normal.logpdf(x, mean, cov)
+    result2 = api.vmap(lsp_stats.multivariate_normal.logpdf)(x, mean, cov)
+    self.assertArraysEqual(result1, result2)
+
 
 if __name__ == "__main__":
     absltest.main(testLoader=jtu.JaxTestLoader())