如何用shape为(2,2)和(2,2,3)的矩阵相乘得到shape为(2,3)的结果

问题描述 #

现有两个NumPy矩阵：

• a 形状为 (2, 2)
• b 形状为 (2, 2, 3)
  需要计算得到形状为 (2, 3) 的结果矩阵 c。

计算逻辑说明 #

从示例推导可知，计算规则为：
取 a 的第i个行向量（形状为 (2,)），与 b 的第i个二维子矩阵（形状为 (2, 3)）做矩阵乘法，得到的长度为3的行向量就是 c 的第i行。
对应示例的计算过程：

c[0] = [0.5, 0.5] × [[1,2,1],[1,3,1]] = [1. , 2.5, 1. ]
c[1] = [0.6, 0.4] × [[2,1,2],[3,1,3]] = [2.4 , 1. , 2.4 ]

实现代码 #

import numpy as np

a = np.array([[0.5, 0.5],
              [0.6, 0.4]])

b = np.array([[[1, 2, 1],
               [1, 3, 1]],

              [[2, 1, 2],
               [3, 1, 3]]])

# 方案1：用einsum实现，语义直观无额外维度转换
c = np.einsum('ij,ijk->ik', a, b)

# 方案2：用广播机制实现，更易理解
c = (a[..., np.newaxis] * b).sum(axis=1)

print(c)

运行输出与预期一致：

[[1.  2.5 1. ]
 [2.4 1.  2.4]]

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