不同维度NumPy数组的矩阵乘法：保持原形状的实现方法

我们有一个形状为(3, 2, 4)的NumPy数组A，包含6个1×4的向量，定义如下：

import numpy as np
A = np.array([
    [[1,2,3,4], [5,6,7,8]],
    [[7,8,9,10], [11,12,13,14]],
    [[15,16,17,18], [19,20,21,22]]
])

另有一个形状为(4, 4)的矩阵B：

B = np.array([
    [1,2,3,4],
    [5,6,7,8],
    [9,10,11,12],
    [13,14,15,16]
])

需求是让矩阵B与A中的每个1×4向量执行矩阵乘法，且结果保持和A相同的(3, 2, 4)形状，预期结果如下：

result = np.array([
    [[90, 100, 110, 120], [202, 228, 254, 280]],
    [[258, 292, 326, 360], [370, 420, 470, 520]],
    [[482, 548, 614, 680], [594, 676, 758, 840]]
])

解决方案 #

方法1：使用@运算符（推荐） #

NumPy的@运算符支持广播机制，直接对A和B执行矩阵乘法即可。A的最后一维（4）与B的第一维（4）匹配，广播会自动遍历A的前两维，输出结果保持原形状：

result = A @ B

验证形状：

print(result.shape)  # 输出 (3, 2, 4)

方法2：使用np.matmul函数 #

np.matmul是专门用于矩阵乘法的函数，功能与@运算符完全等价，同样支持广播：

result = np.matmul(A, B)

方法3：使用np.dot函数 #

np.dot也能实现该需求，不过在高维数组场景下，matmul和@的语义更明确，更推荐使用前两种方法：

result = np.dot(A, B)

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