如何用NumPy对指定轴执行点积运算以匹配目标形状？

你需要将形状为(2,2)的矩阵L，与形状为(2,5,3)的数组W中每个沿最后一维的(2,5)切片做矩阵乘法，最终得到形状为(2,5,3)的结果。以下是几种高效的向量化实现方式：

方法1：使用np.einsum（最直观） #

einsum可通过下标直接指定维度运算关系，完全匹配你的需求：

import numpy as np

# 示例数据
L = np.random.rand(2, 2)
W = np.random.rand(2, 5, 3)

# 核心运算
W_corr = np.einsum('ab,bci->aci', L, W)
print(W_corr.shape)  # 输出 (2, 5, 3)

下标解释：

• ab对应L的(2,2)维度
• bci对应W的(2,5,3)维度
• aci指定输出维度为(2,5,3)，表示对L的第2个维度（b）和W的第1个维度（b）做求和运算，保留其余维度。

方法2：使用np.tensordot（指定轴匹配） #

你之前调用tensordot时未指定轴，导致维度不匹配。只需明确指定要相乘的轴即可：

W_corr = np.tensordot(L, W, axes=((1), (0)))
print(W_corr.shape)  # 输出 (2, 5, 3)

参数解释：

• axes=((1), (0))表示将L的第1个轴（长度2）与W的第0个轴（长度2）做点积求和，最终保留L的第0轴、W的第1和第2轴，正好得到目标形状。

方法3：使用np.matmul结合维度转置 #

通过调整W的维度顺序，让matmul批量处理每个切片：

# 将W转置为(3,2,5)，把要处理的切片放到batch维度
W_reshaped = W.transpose(2, 0, 1)
# L会自动广播到(3,2,2)，与W_reshaped做批量矩阵乘法
result_batch = L @ W_reshaped
# 再转置回目标形状(2,5,3)
W_corr = result_batch.transpose(1, 2, 0)
print(W_corr.shape)  # 输出 (2, 5, 3)

验证正确性 #

以上三种方法的结果和你循环实现的结果完全一致，可通过np.allclose验证：

# 你的原方法
W_corr_original = []
for i in range(W.shape[-1]):
    res_ = L.dot(W[:, :, i])
    W_corr_original.append(res_)
W_corr_original = np.moveaxis(np.array(W_corr_original), 0, -1)

# 验证相等
print(np.allclose(W_corr, W_corr_original))  # 输出 True

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