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

Question

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

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

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

Answer 1

编辑：此答案的最初重点是 where axesis a tuple的情况，为每个参数指定一个或多个轴。这种用法允许我们对传统的 进行变体dot，特别是对于大于 2d 的数组（我在链接问题中的回答也是， /sf/answers/2930968631/）。作为标量的轴是一种特殊情况，它被翻译成元组版本。所以从本质上讲，它仍然是一个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

正在检查 的轴 0a的大小是否与 的轴 1 的大小匹配b。但由于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])

在我的最新示例中，与在 2 个数组的所有轴上axes=2指定 a 相同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)

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

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)

这是-2在foo(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表示法。