了解NumPy的einsum

我正在努力了解确切的einsum工作原理。我看了一下文档和一些示例，但看起来似乎并不固定。

这是我们在课堂上讲的一个例子：

C = np.einsum("ij,jk->ki", A, B)

对于两个数组A和B

我认为可以A^T * B，但是我不确定（正在对其中之一进行移调吗？）。谁能告诉我这里到底发生了什么（以及使用时的一般情况einsum）？

（注：这个答案是基于短的博客文章约einsum我写了前一阵子。）

怎么einsum办？

假设我们有两个多维数组，A和B。现在假设我们要...

• 乘 A用B一种特殊的方式来创造新的产品阵列; 然后也许
• 沿特定轴求和该新数组；然后也许
• 以特定顺序转置新数组的轴。

有一个很好的机会，einsum可以帮助我们做到这一点更快，内存更是有效的NumPy的功能组合，喜欢multiply，sum和transpose允许。

einsum工作如何？

这是一个简单（但并非完全无关紧要）的示例。取以下两个数组：

A = np.array([0, 1, 2])

B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

我们将逐个元素相乘A，B然后沿着新数组的行求和。在“普通” NumPy中，我们将编写：

>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])

因此，此处的索引操作A将两个数组的第一个轴对齐，以便可以广播乘法。然后将乘积数组中的行相加以返回答案。

现在，如果我们想使用它einsum，我们可以这样写：

>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])

该签名字符串'i,ij->i'是这里的关键，需要解释的一点。您可以将其分为两半。在左侧（的左侧->），我们标记了两个输入数组。在的右侧->，我们标记了要结束的数组。

接下来会发生以下情况：

• A有一个轴 我们已经标记了它i。并且B有两个轴；我们将轴0标记为i，将轴1 标记为j。

• 通过在两个输入数组中重复标签i，我们告诉我们einsum这两个轴应该相乘。换句话说，就像A数组B一样，我们将array 与array 的每一列相乘A[:, np.newaxis] * B。

• 请注意，j它不会在所需的输出中显示为标签；我们刚刚使用过i（我们想以一维数组结尾）。通过省略标签，我们告诉einsum来总结沿着这条轴线。换句话说，我们就像对行进行求和.sum(axis=1)。

基本上，这是您需要了解的所有信息einsum。玩一会会有所帮助；如果我们将两个标签都留在输出中，则会'i,ij->ij'返回2D产品数组（与相同A[:, np.newaxis] * B）。如果我们说没有输出标签，'i,ij->我们将返回一个数字（与相同(A[:, np.newaxis] * B).sum()）。

einsum但是，最重要的是，它不会首先构建临时产品系列；它只是对产品进行累加。这样可以节省大量内存。

一个更大的例子

为了解释点积，这里有两个新数组：

A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])

B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

我们将使用计算点积np.einsum('ij,jk->ik', A, B)。这是一张图片，显示了从函数获得的A和B和输出数组的标签：

您会看到j重复的标签-这意味着我们会将的行A与的列相乘B。此外，j输出中不包含标签-我们对这些产品进行求和。标签i和k被保留用于输出，因此我们得到一个2D数组。

这一结果与其中标签阵列比较可能是更加明显j的不求和。在下面的左侧，您可以看到写入产生的3D数组np.einsum('ij,jk->ijk', A, B)（即，我们保留了label j）：

求和轴j给出了预期的点积，如右图所示。

一些练习

为了获得更多的感觉einsum，使用下标符号实现熟悉的NumPy数组操作可能会很有用。任何涉及乘法和求和轴组合的内容都可以使用编写 einsum。

令A和B为两个具有相同长度的一维数组。例如A = np.arange(10)和B = np.arange(5, 15)。

• 的总和A可以写成：

  np.einsum('i->', A)

• A * B可以按元素写成：

  np.einsum('i,i->i', A, B)

• 内积或点积np.inner(A, B)或np.dot(A, B)可以写成：

  np.einsum('i,i->', A, B) # or just use 'i,i'

• 外部乘积np.outer(A, B)可以写成：

  np.einsum('i,j->ij', A, B)

对于2D数组，C和D，只要轴是兼容的长度（相同长度或其中之一具有长度1），下面是一些示例：

• C（主对角线总和）的轨迹np.trace(C)可以写成：

  np.einsum('ii', C)

• 的元素方式乘法C和转置D，C * D.T可以写成：

  np.einsum('ij,ji->ij', C, D)

• 可以将每个元素乘以C该数组D（以构成4D数组）C[:, :, None, None] * D，可以写成：

  np.einsum('ij,kl->ijkl', C, D)  

numpy.einsum()如果您直观地理解它的想法，将非常容易。作为示例，让我们从涉及矩阵乘法的简单描述开始。

使用时numpy.einsum()，您要做的就是传递所谓的下标字符串作为参数，然后传递输入数组。

假设您有两个2D数组A和B，并且想要进行矩阵乘法。所以你也是：

np.einsum("ij, jk -> ik", A, B)

在这里，下标字符串 ij对应于array，A而下标字符串 jk对应于array B。另外，这里要注意的最重要的一点是，每个下标字符串中的字符数必须与数组的大小匹配。（例如，对于2D数组为2个字符，对于3D数组为3个字符，依此类推。）如果您在下标字符串之间重复字符（在我们的示例中），则意味着您希望总和沿着这些维度发生。因此，它们将减少总和。（即该维度将消失） jein

此之后的下标字符串->将成为我们的结果数组。如果将其保留为空，则将对所有内容求和，并返回标量值作为结果。否则，所得数组将具有根据下标字符串的尺寸。在我们的示例中，它将为ik。这很直观，因为我们知道对于矩阵乘法，数组中的列数A必须与数组中的行数相匹配，B这就是这里发生的情况（即，我们通过在下标字符串中重复char j来编码此知识）

这里还有一些其他示例，简要说明了np.einsum()实现某些常见张量或nd数组操作的用途/功能。

输入项

# a vector
In [197]: vec
Out[197]: array([0, 1, 2, 3])

# an array
In [198]: A
Out[198]: 
array([[11, 12, 13, 14],
       [21, 22, 23, 24],
       [31, 32, 33, 34],
       [41, 42, 43, 44]])

# another array
In [199]: B
Out[199]: 
array([[1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3],
       [4, 4, 4, 4]])

1）矩阵乘法（类似于np.matmul(arr1, arr2)）

In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]: 
array([[130, 130, 130, 130],
       [230, 230, 230, 230],
       [330, 330, 330, 330],
       [430, 430, 430, 430]])

2）沿主对角线提取元素（类似于np.diag(arr)）

In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])

3）Hadamard乘积（即两个数组的按元素乘积）（类似于arr1 * arr2）

In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]: 
array([[ 11,  12,  13,  14],
       [ 42,  44,  46,  48],
       [ 93,  96,  99, 102],
       [164, 168, 172, 176]])

4）逐元素平方（类似于np.square(arr)或arr ** 2）

In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]: 
array([[ 1,  1,  1,  1],
       [ 4,  4,  4,  4],
       [ 9,  9,  9,  9],
       [16, 16, 16, 16]])

5）痕迹（即主对角元素的总和）（类似于np.trace(arr)）

In [217]: np.einsum("ii -> ", A)
Out[217]: 110

6）矩阵转置（类似于np.transpose(arr)）

In [221]: np.einsum("ij -> ji", A)
Out[221]: 
array([[11, 21, 31, 41],
       [12, 22, 32, 42],
       [13, 23, 33, 43],
       [14, 24, 34, 44]])

7）（向量的）外积（类似于np.outer(vec1, vec2)）

In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]: 
array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6],
       [0, 3, 6, 9]])

8）（向量的）内积（类似于np.inner(vec1, vec2)）

In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14

9）沿轴0求和（类似于np.sum(arr, axis=0)）

In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])

10）沿轴1的总和（类似于np.sum(arr, axis=1)）

In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4,  8, 12, 16])

11）批矩阵乘法

In [287]: BM = np.stack((A, B), axis=0)

In [288]: BM
Out[288]: 
array([[[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]],

       [[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3],
        [ 4,  4,  4,  4]]])

In [289]: BM.shape
Out[289]: (2, 4, 4)

# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)

In [293]: BMM
Out[293]: 
array([[[1350, 1400, 1450, 1500],
        [2390, 2480, 2570, 2660],
        [3430, 3560, 3690, 3820],
        [4470, 4640, 4810, 4980]],

       [[  10,   10,   10,   10],
        [  20,   20,   20,   20],
        [  30,   30,   30,   30],
        [  40,   40,   40,   40]]])

In [294]: BMM.shape
Out[294]: (2, 4, 4)

12）沿轴2的总和（类似于np.sum(arr, axis=2)）

In [330]: np.einsum("ijk -> ij", BM)
Out[330]: 
array([[ 50,  90, 130, 170],
       [  4,   8,  12,  16]])

13）对数组中的所有元素求和（类似于np.sum(arr)）

In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480

14）多轴总和（即边际化）
（类似于np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7))）

# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))

# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)

# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))

In [365]: np.allclose(esum, nsum)
Out[365]: True

15）双点积（类似于np.sum（哈达玛积） cf. 3）

In [772]: A
Out[772]: 
array([[1, 2, 3],
       [4, 2, 2],
       [2, 3, 4]])

In [773]: B
Out[773]: 
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])

In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124

16）2D和3D阵列乘法

在要验证结果的线性方程组（Ax = b）求解时，这种乘法可能非常有用。

# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)

# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)

# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)

# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True

相反，如果必须使用np.matmul()此验证，则我们必须执行几项reshape操作才能获得相同的结果，例如：

# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)

# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True

奖金：在这里阅读更多数学：爱因斯坦求和，当然在这里：张量表示法

让我们制作2个数组，它们具有不同但兼容的维度，以突出它们之间的相互作用

In [43]: A=np.arange(6).reshape(2,3)
Out[43]: 
array([[0, 1, 2],
       [3, 4, 5]])


In [44]: B=np.arange(12).reshape(3,4)
Out[44]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

您的计算将（2,3）的“点”（乘积之和）与（3,4）相乘，以生成（4,2）数组。 i是第一个昏暗的A，最后一个C; k最后B1个，第1个C。 j通过求和“消耗”。

In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]: 
array([[20, 56],
       [23, 68],
       [26, 80],
       [29, 92]])

这与np.dot(A,B).T-是转置的最终输出相同。

要查看更多情况j，请将C下标更改为ijk：

In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  3,  6,  9],
        [16, 20, 24, 28],
        [40, 45, 50, 55]]])

这也可以通过以下方式产生：

A[:,:,None]*B[None,:,:]

即，添加一个k维度的端部A，以及i与前部B，产生了（2,3,4）阵列。

0 + 4 + 16 = 20，9 + 28 + 55 = 92等; 求和j转置以获得较早的结果：

np.sum(A[:,:,None] * B[None,:,:], axis=1).T

# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,)  j,k]

我发现NumPy：交易技巧（第二部分）具有启发性

我们使用->指示输出数组的顺序。因此，将“ ij，i-> j”视为具有左侧（LHS）和右侧（RHS）。LHS上标签的任何重复都会明智地计算乘积元素，然后求和。通过更改RHS（输出）端的标签，我们可以相对于输入数组定义要在其中进行处理的轴，即沿轴0、1求和，依此类推。

import numpy as np

>>> a
array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)

请注意，存在三个轴，即i，j，k，并且重复了j（在左侧）。 i,j代表的行和列a。j,k为b。

为了计算乘积并对齐j轴，我们需要在上添加一个轴a。（b将沿第一个轴广播？）

a[i, j, k]
   b[j, k]

>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 0,  2,  4],
        [ 6,  8, 10],
        [12, 14, 16]],

       [[ 0,  3,  6],
        [ 9, 12, 15],
        [18, 21, 24]]])

j在右侧不存在，因此我们求和j是3x3x3数组的第二个轴

>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
       [18, 24, 30],
       [27, 36, 45]])

最后，索引在右侧（按字母顺序）相反，因此我们进行了转置。

>>> c.T
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])

>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])
>>>

在阅读einsum方程式时，我发现最简单的方法就是将它们简化为必要的形式。

让我们从以下（强加）语句开始：

C = np.einsum('bhwi,bhwj->bij', A, B)

首先通过标点符号进行操作，我们看到在箭头前有两个4个逗号分隔的斑点- bhwi和bhwj，在箭头后有一个3个字母斑点bij。因此，该方程从两个4级张量输入产生3级张量结果。

现在，让每个斑点中的每个字母成为范围变量的名称。字母在Blob中出现的位置是该张量范围内的轴的索引。因此，产生C的每个元素的命令式求和必须从三个嵌套的for循环开始，每个C的索引一个。

for b in range(...):
    for i in range(...):
        for j in range(...):
            # the variables b, i and j index C in the order of their appearance in the equation
            C[b, i, j] = ...

因此，从本质for上讲，每个C的输出索引都有一个循环。我们现在暂时不确定范围。

接下来，我们看一下左侧-是否有没有出现在右侧的范围变量？在我们的情况下-是，h并且w。for为每个此类变量添加一个内部嵌套循环：

for b in range(...):
    for i in range(...):
        for j in range(...):
            C[b, i, j] = 0
            for h in range(...):
                for w in range(...):
                    ...

现在，在最内层的循环中，我们定义了所有索引，因此我们可以编写实际的求和并完成转换：

# three nested for-loops that index the elements of C
for b in range(...):
    for i in range(...):
        for j in range(...):

            # prepare to sum
            C[b, i, j] = 0

            # two nested for-loops for the two indexes that don't appear on the right-hand side
            for h in range(...):
                for w in range(...):
                    # Sum! Compare the statement below with the original einsum formula
                    # 'bhwi,bhwj->bij'

                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]

如果到目前为止您已经能够遵循该代码，那么恭喜您！这就是您需要阅读einsum方程式所需的全部。请特别注意原始einsum公式如何映射到以上代码段中的最终sumsum语句。for循环和范围边界只是模糊不清的，最终声明是您真正需要了解的所有内容。

为了完整起见，让我们看看如何确定每个范围变量的范围。嗯，每个变量的范围只是它索引的维度的长度。显然，如果变量在一个或多个张量中索引一个以上的维度，则每个维度的长度必须相等。这是上面带有完整范围的代码：

# C's shape is determined by the shapes of the inputs
# b indexes both A and B, so its range can come from either A.shape or B.shape
# i indexes only A, so its range can only come from A.shape, the same is true for j and B
assert A.shape[0] == B.shape[0]
assert A.shape[1] == B.shape[1]
assert A.shape[2] == B.shape[2]
C = np.zeros((A.shape[0], A.shape[3], B.shape[3]))
for b in range(A.shape[0]): # b indexes both A and B, or B.shape[0], which must be the same
    for i in range(A.shape[3]):
        for j in range(B.shape[3]):
            # h and w can come from either A or B
            for h in range(A.shape[1]):
                for w in range(A.shape[2]):
                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]