prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <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?
  
  </issue>
  <code>
  [start of jax/scipy/stats/multivariate_normal.py]
  1 # Copyright 2018 Google LLC
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     https://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 
  16 import numpy as np
  17 import scipy.stats as osp_stats
  18 
  19 from ... import lax
  20 from ...numpy.lax_numpy import _promote_dtypes_inexact, _constant_like, _wraps
  21 from ...numpy.lax_numpy import dot, subtract, einsum
  22 from ...numpy.linalg import det, inv
  23 
  24 
  25 @_wraps(osp_stats.multivariate_normal.logpdf, update_doc=False)
  26 def logpdf(x, mean, cov):
  27   x, mean, cov = _promote_dtypes_inexact(x, mean, cov)
  28   two = _constant_like(x, 2)
  29   dim = _constant_like(x, mean.shape[0])
  30   det_sig = det(cov).astype(cov.dtype)
  31   log_normalizer = lax.log(lax.mul(lax.pow(_constant_like(x, 2 * np.pi), dim),
  32     det_sig))
  33   x_shape = x.shape[:-1]
  34   if x_shape:
  35     x_2d = x.reshape((-1, mean.shape[0]))
  36     quadratic = einsum("ij,jk,ik->i", subtract(x_2d, mean), inv(cov), 
  37       subtract(x_2d, mean)).reshape(x_shape).astype(cov.dtype)
  38   else:
  39     quadratic = dot(dot(subtract(x, mean), inv(cov)), subtract(x, mean).T).astype(cov.dtype)
  40   return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)
  41 
  42 @_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)
  43 def pdf(x, mean, cov):
  44   return lax.exp(logpdf(x, mean, cov))
  45
  [end of jax/scipy/stats/multivariate_normal.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: 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):
  
verification_info: {"golden_diff": "diff --git a/jax/scipy/stats/multivariate_normal.py b/jax/scipy/stats/multivariate_normal.py\n--- a/jax/scipy/stats/multivariate_normal.py\n+++ b/jax/scipy/stats/multivariate_normal.py\n@@ -17,27 +17,29 @@\n import scipy.stats as osp_stats\n \n from ... import lax\n+from ...lax_linalg import cholesky, triangular_solve\n+from ... import numpy as jnp\n from ...numpy.lax_numpy import _promote_dtypes_inexact, _constant_like, _wraps\n-from ...numpy.lax_numpy import dot, subtract, einsum\n-from ...numpy.linalg import det, inv\n \n \n @_wraps(osp_stats.multivariate_normal.logpdf, update_doc=False)\n def logpdf(x, mean, cov):\n   x, mean, cov = _promote_dtypes_inexact(x, mean, cov)\n-  two = _constant_like(x, 2)\n-  dim = _constant_like(x, mean.shape[0])\n-  det_sig = det(cov).astype(cov.dtype)\n-  log_normalizer = lax.log(lax.mul(lax.pow(_constant_like(x, 2 * np.pi), dim),\n-    det_sig))\n-  x_shape = x.shape[:-1]\n-  if x_shape:\n-    x_2d = x.reshape((-1, mean.shape[0]))\n-    quadratic = einsum(\"ij,jk,ik->i\", subtract(x_2d, mean), inv(cov), \n-      subtract(x_2d, mean)).reshape(x_shape).astype(cov.dtype)\n+  if not mean.shape:\n+    return -1/2 * (x - mean) ** 2 / cov - 1/2 * (np.log(2*np.pi) + jnp.log(cov))\n   else:\n-    quadratic = dot(dot(subtract(x, mean), inv(cov)), subtract(x, mean).T).astype(cov.dtype)\n-  return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)\n+    n = mean.shape[-1]\n+    if not np.shape(cov):\n+      y = x - mean\n+      return (-1/2 * jnp.einsum('...i,...i->...', y, y) / cov\n+              - n/2 * (np.log(2*np.pi) + jnp.log(cov)))\n+    else:\n+      if cov.ndim < 2 or cov.shape[-2:] != (n, n):\n+        raise ValueError(\"multivariate_normal.logpdf got incompatible shapes\")\n+      L = cholesky(cov)\n+      y = triangular_solve(L, x - mean, lower=True, transpose_a=True)\n+      return (-1/2 * jnp.einsum('...i,...i->...', y, y) - n/2*np.log(2*np.pi)\n+              - jnp.log(L.diagonal()).sum())\n \n @_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)\n def pdf(x, mean, cov):\n", "issue": "Add multivariate normal pdf evalutation to jax.scipy\nIt 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: \r\n\r\n```\r\n@jit\r\ndef multi_gauss_logpdf(x, mean, cov):\r\n    \"\"\" Calculate the probability density of a\r\n        sample from the multivariate normal. \"\"\"\r\n    D = mean.shape[0]\r\n    (sign, logdet) = np.linalg.slogdet(cov)\r\n    p1 = D*np.log(2*np.pi) + logdet\r\n    p2 = (x-mean).T @ np.linalg.inv(cov) @ (x-mean)\r\n    return -1./2 * (p1 + p2)\r\n\r\nbatch_logpdf = vmap(multi_gauss_logpdf, in_axes=(0, None, None))\r\n```\r\n\r\nMy `lax`/primitive knowledge is still fairly limited but I will try to put together a pr. Any recommendations how to speed things up?\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport numpy as np\nimport scipy.stats as osp_stats\n\nfrom ... import lax\nfrom ...numpy.lax_numpy import _promote_dtypes_inexact, _constant_like, _wraps\nfrom ...numpy.lax_numpy import dot, subtract, einsum\nfrom ...numpy.linalg import det, inv\n\n\n@_wraps(osp_stats.multivariate_normal.logpdf, update_doc=False)\ndef logpdf(x, mean, cov):\n  x, mean, cov = _promote_dtypes_inexact(x, mean, cov)\n  two = _constant_like(x, 2)\n  dim = _constant_like(x, mean.shape[0])\n  det_sig = det(cov).astype(cov.dtype)\n  log_normalizer = lax.log(lax.mul(lax.pow(_constant_like(x, 2 * np.pi), dim),\n    det_sig))\n  x_shape = x.shape[:-1]\n  if x_shape:\n    x_2d = x.reshape((-1, mean.shape[0]))\n    quadratic = einsum(\"ij,jk,ik->i\", subtract(x_2d, mean), inv(cov), \n      subtract(x_2d, mean)).reshape(x_shape).astype(cov.dtype)\n  else:\n    quadratic = dot(dot(subtract(x, mean), inv(cov)), subtract(x, mean).T).astype(cov.dtype)\n  return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)\n\n@_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)\ndef pdf(x, mean, cov):\n  return lax.exp(logpdf(x, mean, cov))\n", "path": "jax/scipy/stats/multivariate_normal.py"}]}