inf using multivariate_normal.logpdf with multiple batches

[sidravi1]

multivariate_normal.logpdf gives infs 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)

[jakevdp]

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:

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.