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

问题描述

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

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

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

>>> np.einsum("ab,ibjk->iajk",A,B)
array([[[[ 4,5],[ 6,7]],[[ 0,[ 2,3]]],[[[12,13],[14,15]],[[ 8,9],[10,11]]]]) 

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

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

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

解决方法

正如我在之前的回答中强调的那样，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,5,7)

按照tensordot 的说法，B 的非收缩尺寸遵循A 的尺寸。所以采用相同的数组，但移动轴，产生相同的结果。

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

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

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

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

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

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

In [147]: C=np.tensordot(A,B,1))
In [148]: C.shape
Out[148]: (2,2)
In [149]: C
Out[149]: 
array([[[[ 4,5],[ 6,7]],[[12,13],[14,15]]],[[[ 0,[ 2,3]],[[ 8,9],[10,11]]]])

以及带有默认结果排序 ('aijk') 的 einsum

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,4))
In [153]: E.shape
Out[153]: (2,4)
In [154]: np.allclose(C,E.reshape(2,2))
Out[154]: True
In [155]: np.allclose(E,np.einsum('ab,ibk',4)))
Out[155]: True