numpy.tensordot 函数是如何逐步工作的？

Question

我是 numpy 的新手，所以我在可视化numpy.tensordot()函数的工作时numpy.tensordot()一些问题。 根据tensordot的文档，轴在参数中传递，其中轴=0 或 1 表示正常矩阵乘法，而轴=2 表示收缩。

有人可以解释一下乘法将如何处理给定的例子吗？

示例 1： a=[1,1] b=[2,2] for axes=0,1为什么它会在 axes=2 时引发错误？
示例 2： a=[[1,1],[1,1]] b=[[2,2],[2,2]] for axes=0,1,2

Answer 1

编辑：此答案的最初重点是axes是元组的情况，为每个参数指定一个或多个轴。 这种用法允许我们对传统的dot进行变体，特别是对于大于 2d 的数组（我在链接问题中的回答也是， https://stackoverflow.com/a/41870980/901925 ）。 作为标量的轴是一种特殊情况，它被翻译成元组版本。 所以它的核心仍然是一个dot积。

轴作为元组

In [235]: a=[1,1]; b=[2,2]

a和b是列表； tensordot将它们变成数组。

In [236]: np.tensordot(a,b,(0,0))
Out[236]: array(4)

由于它们都是一维数组，我们将轴值指定为 0。

如果我们尝试指定 1：

In [237]: np.tensordot(a,b,(0,1))
---------------------------------------------------------------------------
   1282     else:
   1283         for k in range(na):
-> 1284             if as_[axes_a[k]] != bs[axes_b[k]]:
   1285                 equal = False
   1286                 break

IndexError: tuple index out of range

它正在检查a的轴 0 的大小是否与b的轴 1 的大小匹配。 但由于b是 1d，它无法检查。

In [239]: np.array(a).shape[0]
Out[239]: 2
In [240]: np.array(b).shape[1]
IndexError: tuple index out of range

你的第二个例子是二维数组：

In [242]: a=np.array([[1,1],[1,1]]); b=np.array([[2,2],[2,2]])

指定的最后一个轴a和第一b （第二到最后），产生传统的矩阵（点）产物：

In [243]: np.tensordot(a,b,(1,0))
Out[243]: 
array([[4, 4],
       [4, 4]])
In [244]: a.dot(b)
Out[244]: 
array([[4, 4],
       [4, 4]])

更好的诊断值：

In [250]: a=np.array([[1,2],[3,4]]); b=np.array([[2,3],[2,1]])
In [251]: np.tensordot(a,b,(1,0))
Out[251]: 
array([[ 6,  5],
       [14, 13]])
In [252]: np.dot(a,b)
Out[252]: 
array([[ 6,  5],
       [14, 13]])

In [253]: np.tensordot(a,b,(0,1))
Out[253]: 
array([[11,  5],
       [16,  8]])
In [254]: np.dot(b,a)      # same numbers, different layout
Out[254]: 
array([[11, 16],
       [ 5,  8]])
In [255]: np.dot(b,a).T
Out[255]: 
array([[11,  5],
       [16,  8]])

另一个配对：

In [256]: np.tensordot(a,b,(0,0))
In [257]: np.dot(a.T,b)

(0,1,2) 轴是完全错误的。 轴参数应该是 2 个数字或 2 个元组，对应于 2 个参数。

tensordot的基本处理是对输入进行转置和重塑，以便它可以将结果传递给np.dot用于常规（a 的最后一个，b 的最后一个）矩阵乘积。

轴作为标量

如果我对tensordot代码的阅读是正确的，则axes参数将转换为两个列表：

def foo(axes):
    try:
        iter(axes)
    except Exception:
        axes_a = list(range(-axes, 0))
        axes_b = list(range(0, axes))
    else:
        axes_a, axes_b = axes
    try:
        na = len(axes_a)
        axes_a = list(axes_a)
    except TypeError:
        axes_a = [axes_a]
        na = 1
    try:
        nb = len(axes_b)
        axes_b = list(axes_b)
    except TypeError:
        axes_b = [axes_b]
        nb = 1

    return axes_a, axes_b

对于标量值 0,1,2，结果为：

In [281]: foo(0)
Out[281]: ([], [])
In [282]: foo(1)
Out[282]: ([-1], [0])
In [283]: foo(2)
Out[283]: ([-2, -1], [0, 1])

axes=1与在元组中指定相同：

In [284]: foo((-1,0))
Out[284]: ([-1], [0])

对于 2：

In [285]: foo(((-2,-1),(0,1)))
Out[285]: ([-2, -1], [0, 1])

在我的最新示例中， axes=2与在 2 个数组的所有轴上指定一个dot相同：

In [287]: np.tensordot(a,b,axes=2)
Out[287]: array(18)
In [288]: np.tensordot(a,b,axes=((0,1),(0,1)))
Out[288]: array(18)

这与在数组的扁平化 1d 视图上做dot相同：

In [289]: np.dot(a.ravel(), b.ravel())
Out[289]: 18

我已经演示了这些数组的传统点积， axes=1情况。

axes=0与axes=((),()) ，两个数组没有求和轴：

In [292]: foo(((),()))
Out[292]: ([], [])

np.tensordot(a,b,((),()))与np.tensordot(a,b,axes=0)

当输入数组为 1d 时， foo(2)转换中的-2会给您带来问题。 axes=1是一维数组的“收缩”。 换句话说，不要太字面理解文档中的描述。 他们只是试图描述代码的动作； 它们不是正式的规范。

等价物

我认为einsum的轴规格更清晰、更强大。 这是 0,1,2 的等价物

In [295]: np.einsum('ij,kl',a,b)
Out[295]: 
array([[[[ 2,  3],
         [ 2,  1]],

        [[ 4,  6],
         [ 4,  2]]],


       [[[ 6,  9],
         [ 6,  3]],

        [[ 8, 12],
         [ 8,  4]]]])
In [296]: np.einsum('ij,jk',a,b)
Out[296]: 
array([[ 6,  5],
       [14, 13]])
In [297]: np.einsum('ij,ij',a,b)
Out[297]: 18

axis=0 的情况，相当于：

np.dot(a[:,:,None],b[:,None,:])

它添加了一个新的最后一个轴和新的第二个到最后一个轴，并对它们进行传统的点积求和。 但是我们通常用广播做这种“外部”乘法：

a[:,:,None,None]*b[None,None,:,:]

虽然对轴使用 0,1,2 很有趣，但它实际上并没有增加新的计算能力。 轴的元组形式更强大和有用。

代码总结（大步骤）

1 - 将axes转换为axes_a和axes_b如上述foo函数中摘录的

2 - 将a和b组成数组，并获得形状和 ndim

3 - 检查将相加的轴上的匹配大小（收缩）

4 - 构造一个newshape_a和newaxes_a ； b相同（复杂步骤）

5 - at = a.transpose(newaxes_a).reshape(newshape_a) ; b

6 - res = dot(at, bt)

7 - 将res重塑为所需的返回形状

5和6是计算核心。 4 是概念上最复杂的步骤。 对于所有axes值，计算都是相同的， dot积，但设置不同。

超过 0,1,2

虽然文档只提到了标量轴的 0,1,2，但代码不限于这些值

In [331]: foo(3)
Out[331]: ([-3, -2, -1], [0, 1, 2])

如果输入为 3，则轴 = 3 应该可以工作：

In [330]: np.tensordot(np.ones((2,2,2)), np.ones((2,2,2)), axes=3)
Out[330]: array(8.)

或更一般地说：

In [325]: np.tensordot(np.ones((2,2,2)), np.ones((2,2,2)), axes=0).shape
Out[325]: (2, 2, 2, 2, 2, 2)
In [326]: np.tensordot(np.ones((2,2,2)), np.ones((2,2,2)), axes=1).shape
Out[326]: (2, 2, 2, 2)
In [327]: np.tensordot(np.ones((2,2,2)), np.ones((2,2,2)), axes=2).shape
Out[327]: (2, 2)
In [328]: np.tensordot(np.ones((2,2,2)), np.ones((2,2,2)), axes=3).shape
Out[328]: ()

如果输入为 0d，则 axes=0 有效（axes = 1 无效）：

In [335]: np.tensordot(2,3, axes=0)
Out[335]: array(6)

你能解释一下吗？

In [363]: np.tensordot(np.ones((4,2,3)),np.ones((2,3,4)),axes=2).shape
Out[363]: (4, 4)

我已经尝试过 3d 数组的其他标量轴值。 虽然可以提出有效的形状对，但更明确的元组轴值更容易使用。 0,1,2选项是仅适用于特殊情况的捷径。 元组方法更容易使用 - 尽管我仍然更喜欢einsum表示法。

Answer 2

示例 1-0： np.tensordot([1, 1], [2, 2], axes=0)

在这种情况下， a和b都具有单个轴并具有形状(2,) 。

所述axes=0参数可以被转换为（（最后0轴的一个），（第一个0轴线B的）），或者在这种情况下((), ()) 这些是将要收缩的轴。

所有其他轴都不会收缩。 由于a和b 中的每一个都有第 0 个轴而没有其他轴，因此这些轴是((0,), (0,)) 。

然后tensordot操作如下（大致）：

[
    [x*y for y in b]  # all the non-contraction axes in b
    for x in a        # all the non-contraction axes in a
]

请注意，由于a和b之间共有 2 个可用轴，并且由于我们收缩了其中的 0 个，因此结果有 2 个轴。 形状是(2,2)因为它们是a和b 中各个非收缩轴的形状（按顺序）。

示例 1-1： np.tensordot([1, 1], [2, 2], axes=1)

所述axes=1参数可以被转换为（（最后1轴的），（在第一1个轴线B的）），或者在这种情况下((0,), (0,)) 这些是将要收缩的轴

所有其他轴都不会收缩。 由于我们已经在收缩每个轴，剩下的轴是((), ()) 。

然后tensordot操作如下：

sum(  # summing over contraction axis
    [x*y for x,y in zip(a, b)]  # contracted axes must line up
)

请注意，由于我们正在收缩所有轴，因此结果是一个标量（或 0 形张量）。 在 numpy 中，您只会得到一个形状为()的张量，表示 0 轴而不是实际的标量。

示例 1-2： np.tensordot([1, 1], [2, 2], axes=2)

这不起作用的原因是因为a和b都没有两个单独的轴可以收缩。

示例 2-1： np.tensordot([[1,1],[1,1]], [[2,2],[2,2]], axes=1)

我跳过了你的几个例子，因为它们并不复杂到比我认为的前几个更清晰。

在这种情况下， a和b都有两个可用的轴（让这个问题更有趣一些），并且它们都有形状(2,2) 。

所述axes=1论点仍然代表的最后1轴和b的第一1轴，留给我们((1,), (0,)) 这些是将要收缩的轴。

其余的轴不收缩，并有助于最终解决方案的形状。 它们是((0,), (1,)) 。

然后我们可以构建 tensordot 操作。 为了便于论证，假设a和b是 numpy 数组，以便我们可以使用数组属性并使问题更清晰（例如b=np.array([[2,2],[2,2]]) ）。

[
    [
        sum(  # summing the contracted indices
            [x*y for x,y in zip(v,w)]  # axis 1 of a and axis 0 of b must line up for the summation
        )
        for w in b.T  # iterating over axis 1 of b (i.e. the columns)
    ]
    for v in a  # iterating over axis 0 of a (i.e. the rows)
]

结果具有形状(a.shape[0], b.shape[1])因为这些是非收缩轴。