使用tensordot进行批量矩阵乘法

Question

如果我想计算矩阵-矩阵乘积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会很简单，但似乎不是这样。

Answer 1

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

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

(2, 2, 3)