Scipy multivariate_normal能否对均值向量向量化？求内置解决方案

我在马尔可夫链采样器中使用multivariate_normal，需要利用向量化优化，同时计算多个不同均值向量的PDF对数。以下是最简示例：
第一段代码针对单个数据点和单个均值向量可正常运行；但第二段代码尝试计算多个均值向量时，出现如下错误：

ValueError: Array 'mean' must be a vector of length 600.

请问有没有内置解决方案，还是必须使用for循环？

代码示例 #

import numpy as np
from scipy.stats import multivariate_normal

np.random.seed(42)
cov = np.diag(np.ones(6))

# 单个均值、单个数据点，运行正常
mu = np.random.normal(0, 1, 6)
x = np.random.normal(0, 1, 6)
print(multivariate_normal.logpdf(x, mean=mu, cov=cov))

# 多个均值、多个数据点，报错
mu = np.random.normal(0, 1, (100, 6))
x = np.random.normal(0, 1, (100, 6))
print(multivariate_normal.logpdf(x, mean=mu, cov=cov))

编辑补充：我找到单个数据点的解决方案——利用正态分布中均值与数据可互换的特性，以下代码能得到正确结果：

mu = np.random.normal(0, 1, (100, 6))
x = np.random.normal(0, 1, 6)
print(multivariate_normal.logpdf(mu, mean=x, cov=cov))

不需要用for循环，scipy.stats.multivariate_normal.logpdf支持向量化输入，只需注意输入维度的匹配规则：

• 当计算N个数据点对应N个均值向量的对数PDF时，可借助scipy的广播机制实现向量化计算，避免循环损耗：

  import numpy as np
  from scipy.stats import multivariate_normal
  
  np.random.seed(42)
  D = 6
  N = 100
  cov = np.diag(np.ones(D))
  
  mu = np.random.normal(0, 1, (N, D))
  x = np.random.normal(0, 1, (N, D))
  
  # 方法1：手动向量化计算（适合对角协方差场景）
  inv_cov = np.diag(1.0 / np.diag(cov))
  diff = x - mu
  log_det_cov = np.sum(np.log(np.diag(cov)))
  log_pdf = -0.5 * (D * np.log(2 * np.pi) + log_det_cov + np.einsum('ni,ij,nj->n', diff, inv_cov, diff))
  print(log_pdf)
  
  # 方法2：利用scipy广播特性调整维度
  # 将x和mu转为(N, 1, D)形状，让函数正确识别多组输入
  log_pdf_vec = multivariate_normal.logpdf(x[:, None, :], mean=mu[:, None, :], cov=cov).flatten()
  print(log_pdf_vec)
  

• 报错的核心原因是multivariate_normal.logpdf默认要求mean为单个向量，传入多组均值时，需通过维度调整让函数触发广播逻辑，从而实现批量计算。

内容的提问来源于stack exchange，提问作者Aleksejs Fomins