沿numpy张量点的最后一个轴收缩

Question

我对张量代数不是很熟悉，并且在理解如何使numpy.tensordot我的要求时遇到了麻烦。

我正在使用的示例很简单：给定一个具有形状(2,2,3)的张量a和另一个具有形状(2,1,3) b ，我想要一个具有形状(2,1)的结果张量c 。 该张量将是以下等效的python代码的结果：

n = a.shape[2]
c = np.zeros((2,n))
for k in range(n):
    c += a[:,:,k]*b[:,:,k]

文档说可选参数axes ：

如果为int N，则按顺序对a的最后N个轴和b的前N个轴求和。 相应轴的尺寸必须匹配。

但是我不明白这里需要哪些“轴”（此外，当轴是一个元组或一个元组的元组时，它会变得更加混乱）。 例子对我也不是很清楚。

Answer 1

tensordot工作方式，由于沿第一个轴的对齐要求，因此在这里（至少不能直接）工作。 您可以使用np.einsum来解决您的情况-

c = np.einsum('ijk,ilk->ij',a,b)

或者，使用np.matmul / @-operator np.matmul （Python 3.x）-

np.matmul(a,b.swapaxes(1,2))[...,0] # or (a @ b.swapaxes(1,2))[...,0]