如何用einsum计算Numpy张量的相关矩阵？

你可以通过np.einsum直接指定维度的求和规则，实现对每个样本独立计算相关矩阵X.T @ X，具体如下：

核心实现代码 #

import numpy as np

# 假设X是形状为(1075, 252, 487)的张量
X = np.random.randn(1075, 252, 487)

# 用einsum计算逐样本相关矩阵，结果形状为(1075, 487, 487)
corr_matrices = np.einsum('ndp,ndq->npq', X, X)

表达式含义解释 #

• ndp：对应输入张量X的三个维度：n（样本数，1075）、d（每个样本的行维度，252）、p（每个样本的列维度，487）
• ndq：第二个X的维度标注，用q区分最终矩阵的列维度（实际与p是同一维度，仅为了清晰表示矩阵结构）
• ->npq：指定输出张量的维度：保留样本维度n，同时对d维度求和（对应矩阵乘法中的行乘列累加），最终得到每个样本的p×q（即487×487）相关矩阵

结果验证 #

你可以抽取单个样本，手动计算X.T @ X并与einsum的结果对比，确认正确性：

# 取第一个样本手动计算相关矩阵
manual_corr = X[0].T @ X[0]
# 验证einsum结果与手动计算是否一致
print(np.allclose(corr_matrices[0], manual_corr))  # 输出True表示结果一致

等价写法参考 #

如果你习惯矩阵转置的思路，也可以用转置后矩阵乘法实现，但einsum的写法更直观地展示了维度间的运算逻辑：

# 转置每个样本后做矩阵乘法，结果与einsum完全一致
corr_matrices_alt = X.transpose(0, 2, 1) @ X
print(np.allclose(corr_matrices, corr_matrices_alt))  # 输出True

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