如何在NumPy中实现三维数组内维度的点积向量化运算

问题场景 #

现有两个三维数组：

• X：形状为(10, 2, 3)，可理解为包含10个2×3的矩阵
• Y：形状为(8, 1, 3)，可理解为包含8个1×3的矩阵

需要计算两者内维度的点积，得到形状为(10, 8, 2)的结果数组Z，已有正确的循环实现，需替换为向量化方案提升运算效率。

原循环实现代码 #

import numpy as np

X = np.arange(0, 10 * 2 * 3).reshape(10, 2, 3)
Y = np.arange(0, 8 * 1 * 3).reshape(8, 1, 3)

# 初始化结果数组
Z = np.empty((10, 8, 2))
for i, x_i in enumerate(X):
    for j, y_j in enumerate(Y):
        z = x_i @ y_j.T
        # 展平后赋值
        Z[i, j, :] = z[:, 0]

向量化解决方案 #

这里提供两种高效的向量化实现方式：

方法1：使用np.einsum（直观匹配计算逻辑）

einsum可以清晰描述维度间的运算关系，直接对应需求中的计算逻辑：

Z_vectorized = np.einsum('ijk,lmk->ilm', X, Y)

• 维度解释：ijk对应X的(10,2,3)，lmk对应Y的(8,1,3)，指定最后一维（k）做点积，输出维度为ilm即(10,8,2)，完全匹配预期结果形状。

方法2：使用np.matmul结合维度扩展

通过扩展维度触发广播机制，实现批量矩阵乘法：

# 扩展维度后做批量矩阵乘法，再去除多余的尾维度
Z_vectorized = (X[:, None] @ Y[:, :, None].transpose(0,2,1)).squeeze(-1)

验证结果一致性 #

通过以下代码可验证向量化结果与原循环结果完全一致：

print(np.allclose(Z, Z_vectorized))  # 输出True则表示结果一致

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