[Python] Use Scipy multivariate_normal with 3D covariance matrix

I have a covariance matrix with shape (n_samples, n, n).

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal

x, y = np.mgrid[-1:1:.1, -1:1:.1]
pos = np.dstack((x, y))

Sigma = np.array([
[[2.0, 0.3], [0.3, 0.5]],
[[2.0, 0.3], [0.3, 0.5]]
])


mvn = multivariate_normal.logpdf(
x=pos,
mean=[0.5, -0.2],
cov=Sigma
)


This results in the following error:

ValueError: Array 'cov' must be at most two-dimensional, but cov.ndim = 3
​

I don't want to use a "for loop" as the number of samples can be large. This blog shows to calculate the multivariate normal logpdf over multiple sets of parameters.

How do I change the code that it deals correctly with the following dimensions of the parameters

logpdfs = custom_multiple_logpdfs(x=pos, mu=[0.5, -0.2], cov=Sigma)

Continue reading...