Numpy文档中的Tensordot解释

Question

我不明白tensordot是如何工作的，我正在阅读官方文档，但我完全不明白那里发生了什么。

a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
c = np.tensordot(a,b, axes=([1,0],[0,1]))
c.shape
(5, 2)

为什么形状是(5，2)？到底发生了什么？

我也读过this article，但答案把我搞糊涂了。

 In [7]: A = np.random.randint(2, size=(2, 6, 5))
   ...:  B = np.random.randint(2, size=(3, 2, 4))
   ...: 

In [9]: np.tensordot(A, B, axes=((0),(1))).shape
Out[9]: (6, 5, 3, 4)

A : (2, 6, 5) -> reduction of axis=0
B : (3, 2, 4) -> reduction of axis=1

Output : `(2, 6, 5)`, `(3, 2, 4)` ===(2 gone)==> `(6,5)` + `(3,4)` => `(6,5,3,4)`

为什么形状是(6, 5, 3, 4)

Answer 1

In [196]: a = np.arange(60.).reshape(3,4,5)
     ...: b = np.arange(24.).reshape(4,3,2)
     ...: c = np.tensordot(a,b, axes=([1,0],[0,1]))
In [197]: c
Out[197]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

我发现einsum的等价物更容易“阅读”：

In [198]: np.einsum('ijk,jil->kl',a,b)
Out[198]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

tensordot通过转置和重塑输入来将问题简化为简单的dot

In [204]: a1 = a.transpose(2,1,0).reshape(5,12)
In [205]: b1 = b.reshape(12,2)
In [206]: np.dot(a1,b1)        # or a1@b1
Out[206]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

tensordot可以对结果进行进一步的操作，但这里不需要这样做。

在我得到正确的a1/b1之前，我必须尝试几件事。例如，a.transpose(2,0,1).reshape(5,12)会生成正确的形状，但值不同。

还有另一个版本：

In [210]: (a.transpose(1,0,2)[:,:,:,None]*b[:,:,None,:]).sum((0,1))
Out[210]: 
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])