Repository: google/jax

Issue:
Add multivariate normal pdf evalutation to jax.scipy
It would be great to have a Multivariate gaussian pdf/logpdf implementation, similar to the univariate version in [jax.scipy.stats.norm](https://jax.readthedocs.io/en/latest/_modules/jax/scipy/stats/norm.html#logpdf). I am currently working with this hacky function: 

```
@jit
def multi_gauss_logpdf(x, mean, cov):
    """ Calculate the probability density of a
        sample from the multivariate normal. """
    D = mean.shape[0]
    (sign, logdet) = np.linalg.slogdet(cov)
    p1 = D*np.log(2*np.pi) + logdet
    p2 = (x-mean).T @ np.linalg.inv(cov) @ (x-mean)
    return -1./2 * (p1 + p2)

batch_logpdf = vmap(multi_gauss_logpdf, in_axes=(0, None, None))
```

My `lax`/primitive knowledge is still fairly limited but I will try to put together a pr. Any recommendations how to speed things up?


Discussion:
@RobertTLange are you actively working on a PR? If not, I can probably help with this one, since I've written similar code too many times in the past!

It's usually better (faster and more numerically stable) to do a Cholesky and a triangular solve, and you can compute the log det in terms of the diagonal elements of the Cholesky too. Here's [an example](https://github.com/mattjj/pybasicbayes/blob/61f65ad6c781288605ec5f7347efcc5dbd73c4fc/pybasicbayes/distributions/gaussian.py#L67-L75).
@mattjj thanks for coming back and sorry for the late response. Yes, it is on my todo-list. My focus is suffering a little. Should be able to get this done in the next days. Okaydokey than I have learned some new computational linear algebra tricks ;)

P.S.: I wrote a little JAX intro tutorial (https://roberttlange.github.io/posts/2020/03/blog-post-10/). Thanks for the great project and all the effort!
Awesome tutorial, and artwork!

Actually, I just noticed: a [multivariate normal logpdf function](https://github.com/google/jax/blob/6545cf3421e45aa06beb8ad2bb3c73d34ce7716a/jax/scipy/stats/multivariate_normal.py#L26-L40) seems to be checked in already from #268, but it's inefficient because of how it computes `inv` and `det` (it even computes `inv` twice!). Looks like I LGTM'd it though, probably thinking we could improve the efficiency and numerics in follow-up work.

Does that function work for you? if so, we can change this issue title to be about improving it.

Resolving patch:
diff --git a/jax/scipy/stats/multivariate_normal.py b/jax/scipy/stats/multivariate_normal.py
--- a/jax/scipy/stats/multivariate_normal.py
+++ b/jax/scipy/stats/multivariate_normal.py
@@ -17,27 +17,29 @@
 import scipy.stats as osp_stats
 
 from ... import lax
+from ...lax_linalg import cholesky, triangular_solve
+from ... import numpy as jnp
 from ...numpy.lax_numpy import _promote_dtypes_inexact, _constant_like, _wraps
-from ...numpy.lax_numpy import dot, subtract, einsum
-from ...numpy.linalg import det, inv
 
 
 @_wraps(osp_stats.multivariate_normal.logpdf, update_doc=False)
 def logpdf(x, mean, cov):
   x, mean, cov = _promote_dtypes_inexact(x, mean, cov)
-  two = _constant_like(x, 2)
-  dim = _constant_like(x, mean.shape[0])
-  det_sig = det(cov).astype(cov.dtype)
-  log_normalizer = lax.log(lax.mul(lax.pow(_constant_like(x, 2 * np.pi), dim),
-    det_sig))
-  x_shape = x.shape[:-1]
-  if x_shape:
-    x_2d = x.reshape((-1, mean.shape[0]))
-    quadratic = einsum("ij,jk,ik->i", subtract(x_2d, mean), inv(cov), 
-      subtract(x_2d, mean)).reshape(x_shape).astype(cov.dtype)
+  if not mean.shape:
+    return -1/2 * (x - mean) ** 2 / cov - 1/2 * (np.log(2*np.pi) + jnp.log(cov))
   else:
-    quadratic = dot(dot(subtract(x, mean), inv(cov)), subtract(x, mean).T).astype(cov.dtype)
-  return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)
+    n = mean.shape[-1]
+    if not np.shape(cov):
+      y = x - mean
+      return (-1/2 * jnp.einsum('...i,...i->...', y, y) / cov
+              - n/2 * (np.log(2*np.pi) + jnp.log(cov)))
+    else:
+      if cov.ndim < 2 or cov.shape[-2:] != (n, n):
+        raise ValueError("multivariate_normal.logpdf got incompatible shapes")
+      L = cholesky(cov)
+      y = 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())
 
 @_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)
 def pdf(x, mean, cov):


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -757,7 +757,7 @@ def wrapped_fun(*args):
   def _CheckAgainstNumpy(self, numpy_reference_op, lax_op, args_maker,
                          check_dtypes=False, tol=None):
     args = args_maker()
-    numpy_ans = numpy_reference_op(*args)
     lax_ans = lax_op(*args)
+    numpy_ans = numpy_reference_op(*args)
     self.assertAllClose(numpy_ans, lax_ans, check_dtypes=check_dtypes,
                         atol=tol, rtol=tol)
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
@@ -260,24 +260,6 @@ def args_maker():
                             tol=1e-6)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
-  # TODO: currently it ignores the argument "shapes" and only tests dim=4
-  @genNamedParametersNArgs(3, jtu.rand_default)
-  def testMultivariateNormalLogPdf(self, rng_factory, shapes, dtypes):
-    rng = rng_factory()
-    scipy_fun = osp_stats.multivariate_normal.logpdf
-    lax_fun = lsp_stats.multivariate_normal.logpdf
-    dim = 4
-    shapex = (dim,)
-
-    def args_maker():
-      x, mean, cov = map(rng, (shapex, shapex, (dim, dim)), dtypes)
-      cov = random_correlation.rvs(onp.arange(1, 1+dim) * 2 / (dim + 1))
-      return [x, mean, cov]
-
-    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
-      tol=1e-4)
-    self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
-
   @genNamedParametersNArgs(3, jtu.rand_default)
   def testNormLogPdf(self, rng_factory, shapes, dtypes):
     rng = rng_factory()
@@ -400,6 +382,67 @@ def testIssue972(self):
       lsp_stats.norm.cdf(onp.full((4,), onp.inf, onp.float32)),
       check_dtypes=False)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_x={}_mean={}_cov={}".format(
+          jtu.format_shape_dtype_string(x_shape, x_dtype),
+          jtu.format_shape_dtype_string(mean_shape, mean_dtype)
+          if mean_shape is not None else None,
+          jtu.format_shape_dtype_string(cov_shape, cov_dtype)
+          if cov_shape is not None else None),
+       "x_shape": x_shape, "x_dtype": x_dtype,
+       "mean_shape": mean_shape, "mean_dtype": mean_dtype,
+       "cov_shape": cov_shape, "cov_dtype": cov_dtype,
+       "rng_factory": rng_factory}
+      for x_shape, mean_shape, cov_shape in [
+          # # These test cases cover default values for mean/cov, but we don't
+          # # support those yet (and they seem not very valuable).
+          # [(), None, None],
+          # [(), (), None],
+          # [(2,), None, None],
+          # [(2,), (), None],
+          # [(2,), (2,), None],
+          # [(3, 2), (3, 2,), None],
+          # [(5, 3, 2), (5, 3, 2,), None],
+
+          [(), (), ()],
+          [(3,), (), ()],
+          [(3,), (3,), ()],
+          [(3,), (3,), (3, 3)],
+          [(3, 4), (4,), (4, 4)],
+
+          # # These test cases are where scipy flattens things, which has
+          # # different batch semantics than some might expect
+          # [(5, 3, 2), (5, 3, 2,), ()],
+          # [(5, 3, 2), (5, 3, 2,), (5, 3, 2, 2)],
+          # [(5, 3, 2), (3, 2,), (5, 3, 2, 2)],
+          # [(5, 3, 2), (3, 2,), (2, 2)],
+      ]
+      for x_dtype, mean_dtype, cov_dtype in CombosWithReplacement(float_dtypes, 3)
+      if (mean_shape is not None or mean_dtype == onp.float32)
+      and (cov_shape is not None or cov_dtype == onp.float32)
+      for rng_factory in [jtu.rand_default]))
+  def testMultivariateNormalLogpdf(self, x_shape, x_dtype, mean_shape,
+                                   mean_dtype, cov_shape, cov_dtype, rng_factory):
+    rng = rng_factory()
+    def args_maker():
+      args = [rng(x_shape, x_dtype)]
+      if mean_shape is not None:
+        args.append(5 * rng(mean_shape, mean_dtype))
+      if cov_shape is not None:
+        if cov_shape == ():
+          args.append(0.1 + rng(cov_shape, cov_dtype) ** 2)
+        else:
+          factor_shape = (*cov_shape[:-1], 2 * cov_shape[-1])
+          factor = rng(factor_shape, cov_dtype)
+          args.append(onp.matmul(factor, onp.swapaxes(factor, -1, -2)))
+      return args
+
+    self._CheckAgainstNumpy(osp_stats.multivariate_normal.logpdf,
+                            lsp_stats.multivariate_normal.logpdf,
+                            args_maker, check_dtypes=True, tol=1e-3)
+    self._CompileAndCheck(lsp_stats.multivariate_normal.logpdf, args_maker,
+                          check_dtypes=True)
+
 
 if __name__ == "__main__":
     absltest.main()