使用tensordot无法重新创建具有“交错”输出索引的一行einsum函数？

Question

NumPy 的tensordot 和einsum 函数之间的异同有详细记录，并在本论坛中进行了广泛讨论（例如[1]、[2]、[3]、[4]、[5]） .但是，我遇到了一个使用einsum 进行矩阵乘法的实例，我发现使用tensordot 进行复制非常困难，如果不是不可能的话：如果我们的两个数组是，

>>> A = np.array([[0, 1], [1, 0]])
>>> B = np.arange(2 ** 4).reshape((2, 2, 2, 2))

是否存在一行tensordot 等价于以下内容？

>>> np.einsum("ab,ibjk->iajk", A, B)
array([[[[ 4,  5],
         [ 6,  7]],

        [[ 0,  1],
         [ 2,  3]]],


       [[[12, 13],
         [14, 15]],

        [[ 8,  9],
         [10, 11]]]]) 

根据我的发现，答案似乎是“否”。问题出现在输出维度iajk 的索引中。这里，数组A 的维度a 出现在数组B 的维度i 和j 之间。如果输出维度的索引改为aijk，np.tensordot(A, B, (1, 1)) 会正常工作。为了确定，我使用所有可能的轴进行了测试，

>>> output_einsum = np.einsum("ab,ibjk->iajk", A, B)
>>> axes_A = [-2, -1, 0, 1]
>>> axes_B = [-4, -3, -2, -1, 0, 1, 2, 3]
>>> for i in axes_A:
...     for j in axes_B:
...         output_tensordot = np.tensordot(A, B, axes=(i, j))
...         if np.allclose(ouput_einsum, output_tensordot):
...             print(i,j)
...

并发现没有任何允许的轴组合产生所需的结果。请注意，B 的维度将axes 参数的每个元素的长度限制为一。使用tensordot 不能在一行中重现具有交错输出尺寸的einsum 函数是否正确？如果是这样，是否存在多行解决方法？

Answer 1

正如我在之前的回答中所强调的，tensordot 是np.dot 的扩展，允许我们指定在乘积和中使用哪些维度。 dot 默认是 A 的最后一个，B 的倒数第二个。

这说明了dot 如何处理大于 2 的维度：

In [158]: np.dot(np.ones((2,3,4)),np.ones((5,4,7))).shape
Out[158]: (2, 3, 5, 7)

按照tensordot 的说法，B 的非收缩维度遵循A 的维度。因此，采用相同的数组，但移动轴，会产生相同的结果。

In [162]: np.tensordot(np.ones((2,4,3)),np.ones((5,7,4)),(1,2)).shape 
Out[162]: (2, 3, 5, 7)

在这些示例中，我选择了不同的维度，因此顺序更加明显。

tensordot 不提供对非收缩维度重新排序的方法。但之后您可以轻松地做到这一点。

您的示例周围有 2 个尺寸。这允许您指定任意轴组合，但需要使用allclose 来测试结果。

In [146]: >>> A = np.array([[0, 1], [1, 0]])
     ...: >>> B = np.arange(2 ** 4).reshape((2, 2, 2, 2))

在两个数组的第 2 轴上执行积和：

In [147]: C=np.tensordot(A,B,(1,1))
In [148]: C.shape
Out[148]: (2, 2, 2, 2)
In [149]: C
Out[149]: 
array([[[[ 4,  5],
         [ 6,  7]],

        [[12, 13],
         [14, 15]]],


       [[[ 0,  1],
         [ 2,  3]],

        [[ 8,  9],
         [10, 11]]]])

还有 einsum 及其默认结果排序 ('aijk')

In [150]: D= np.einsum('ab,ibjk',A,B)
In [151]: np.allclose(C,D)
Out[151]: True

tensordot 相当于这个dot：

In [152]: E = np.dot(A,B.reshape(2,2,4))
In [153]: E.shape
Out[153]: (2, 2, 4)
In [154]: np.allclose(C,E.reshape(2,2,2,2))
Out[154]: True
In [155]: np.allclose(E,np.einsum('ab,ibk',A,B.reshape(2,2,4)))
Out[155]: True