Repository: google/jax

Issue:
multivariate_normal.pdf does not work for multi-dimensional input

```python
size = 256
x, y = jnp.mgrid[-1:1:(2/size), -1:1:(2/size)]
pos = jnp.dstack((x, y))
mean = jnp.array([0.5, -0.2])
cov = jnp.array([[2.0, 0.3], [0.3, 0.5]])
gs = jax.scipy.stats.multivariate_normal.pdf(pos, mean=mean, cov=cov)
```
Expected: `gs` holds a tensor of shape `[256, 256]` with pdf values (as in SciPy).
Result:
```
/usr/local/lib/python3.7/dist-packages/jax/_src/lax/linalg.py in triangular_solve_shape_rule(a, b, left_side, **unused_kwargs)
    582     msg = ("triangular_solve requires both arguments to have the same number "
    583            "of dimensions and equal batch dimensions, got {} and {}.")
--> 584     raise TypeError(msg.format(a.shape, b.shape))
    585   common_dim = -2 if left_side else -1
    586   if a.shape[-1] != b.shape[common_dim]:

TypeError: triangular_solve requires both arguments to have the same number of dimensions and equal batch dimensions, got (2, 2) and (256, 256, 2).
```
JAX version: 0.2.21
SciPy version: 1.4.1
See also this notebook: https://colab.research.google.com/drive/1YOW5Q2HubHQbYtrdy1B_xnT-_qVeLKOh?usp=sharing



Discussion:
Thanks for the report - I'll take a look
It looks like some of these cases were explicitly left out in #2481; we should be able to add them.

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
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import partial
 
 import numpy as np
 import scipy.stats as osp_stats
@@ -42,7 +43,10 @@ def logpdf(x, mean, cov, allow_singular=None):
       if cov.ndim < 2 or cov.shape[-2:] != (n, n):
         raise ValueError("multivariate_normal.logpdf got incompatible shapes")
       L = lax.linalg.cholesky(cov)
-      y = lax.linalg.triangular_solve(L, x - mean, lower=True, transpose_a=True)
+      y = jnp.vectorize(
+        partial(lax.linalg.triangular_solve, lower=True, transpose_a=True),
+        signature="(n,n),(n)->(n)"
+      )(L, x - mean)
       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))
 


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
@@ -528,13 +528,7 @@ def testIssue972(self):
           [(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)],
+          [(2, 3, 4), (4,), (4, 4)],
       ]
       for x_dtype, mean_dtype, cov_dtype in itertools.combinations_with_replacement(jtu.dtypes.floating, 3)
       if (mean_shape is not None or mean_dtype == np.float32)
@@ -561,6 +555,56 @@ 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": "_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}
+      for x_shape, mean_shape, cov_shape in [
+          # These test cases are where scipy flattens things, which has
+          # different batch semantics than some might expect, so we manually
+          # vectorize scipy's outputs for the sake of testing.
+          [(5, 3, 2), (5, 3, 2), (5, 3, 2, 2)],
+          [(2,), (5, 3, 2), (5, 3, 2, 2)],
+          [(5, 3, 2), (2,), (5, 3, 2, 2)],
+          [(5, 3, 2), (5, 3, 2,), (2, 2)],
+          [(1, 3, 2), (3, 2,), (5, 1, 2, 2)],
+          [(5, 3, 2), (1, 2,), (2, 2)],
+      ]
+      for x_dtype, mean_dtype, cov_dtype in itertools.combinations_with_replacement(jtu.dtypes.floating, 3)
+      if (mean_shape is not None or mean_dtype == np.float32)
+      and (cov_shape is not None or cov_dtype == np.float32)))
+  def testMultivariateNormalLogpdfBroadcasted(self, x_shape, x_dtype, mean_shape,
+                                              mean_dtype, cov_shape, cov_dtype):
+    rng = jtu.rand_default(self.rng())
+    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(np.matmul(factor, np.swapaxes(factor, -1, -2)))
+      return args
+
+    osp_fun = np.vectorize(osp_stats.multivariate_normal.logpdf,
+                           signature="(n),(n),(n,n)->()")
+
+    self._CheckAgainstNumpy(osp_fun, lsp_stats.multivariate_normal.logpdf,
+                            args_maker, tol=1e-3)
+    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}