使用 tensordot 批量矩阵乘法答案

【问题标题】：Using tensordot to batch matrix multiplies使用 tensordot 批量矩阵乘法
【发布时间】：2021-11-20 06:35:14
【问题描述】：

如果我想计算矩阵-矩阵乘积a*b，我会做a@b 或np.dot(a,b)

a = np.random.rand(2,2)
b = np.random.rand(*a.shape)
c = a@b
c.shape
>>> (2,2)

一般来说，我可以使用 tensordot 来做同样的事情：

c = np.tensordot(a,b,1)
c.shape
>>> 2,2

但是，如果我向 a 和 b 添加一个维度，我将不再得到我想要的结果（另一个 2,2,3 数组）：

a = np.random.rand(2,2,3)
b = np.random.rand(*a.shape)
c = np.tensordot(a,b,1)
c.shape
>>> ValueError: shape-mismatch for sum

我尝试的任何轴排列，即使是可怕的b,axes=((0,1),(0,1))，都会导致输出形状不正确，或者由于无法完成计算而导致错误。

有没有办法完成我想做的事情？我的印象是使用 tensordot 会很简单，但似乎不是......

【问题讨论】：

【解决方案1】：

这给了我 (2, 2, 3) 输出。但请注意，我以前从未做过 3d 点积，所以请检查这些输出值是否正确：

output = np.matmul(a,b, axes=[(0, 1),(0, 1),(0, 1)])
print(output.shape)

(2, 2, 3)

【讨论】：