python - 了解 NumPy 的 einsum

Question

我正在努力理解究竟是如何einsum工作的。我查看了文档和一些示例，但似乎并没有坚持下去。

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

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

对于两个数组：A和B.

我认为这需要A^T * B，但我不确定（它正在对其中一个进行转置，对吗？）。谁能告诉我这里发生了什么（以及一般使用时einsum）？

score 562 · Accepted Answer

（注意：这个答案是基于我不久前写的一篇简短的博客文章。）einsum

做什么`einsum`？

假设我们有两个多维数组，A并且B. 现在让我们假设我们想要...

以特定方式相乘A以创建新的产品数组；B然后也许
沿特定轴对这个新数组求和；然后也许
以特定顺序转置新数组的轴。

与 NumPy 函数（如 , ）的组合相比，这很有可能einsum会帮助我们更快、更高效地执行此操作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这两个轴应该相乘。换句话说，我们将 array 与 arrayA的每一列相乘B，就像这样A[:, np.newaxis] * B做一样。
请注意，j它不会在我们想要的输出中显示为标签；我们刚刚使用过i（我们希望得到一个一维数组）。通过省略标签，我们告诉einsum沿该轴求和。换句话说，我们正在对产品的行求和，就像这样.sum(axis=1)做一样。

这基本上就是您使用einsum. 它有助于玩一点；如果我们将两个标签都留在输出中，'i,ij->ij'，我们将返回一个二维产品数组（与相同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)。这是一张图片，显示了我们从函数中获得的Aand和输出数组的标签：B

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

将此结果与标签未求和的数组进行比较可能会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)ornp.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)    
```

score 70 · Accepted Answer

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

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

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

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

这里下标字符串 ij对应于数组A，而下标字符串 jk对应于数组B。另外，这里要注意的最重要的一点是，每个下标字符串中的字符数必须与数组的维度相匹配。（即 2D 数组的两个字符，3D 数组的三个字符，等等。）如果您在下标字符串之间重复字符（在我们的例子中），那么这意味着您希望总和沿着这些维度发生。因此，它们将被归约。（即那个维度将消失） jein

this 之后的下标字符串->，将是我们的结果数组。如果你把它留空，那么所有的东西都会被求和，并返回一个标量值作为结果。否则，结果数组将根据下标字符串具有尺寸。在我们的示例中，它将是ik. 这很直观，因为我们知道对于矩阵乘法，数组中的列数必须与数组A中的行数相匹配，B这就是这里发生的情况（即，我们通过重复下标字符串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)or 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) Trace（即主对角元素之和）（类似于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(hadamard-product) 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

奖励：在这里阅读更多数学：爱因斯坦求和，绝对在这里：张量符号

score 13 · Accepted Answer

在阅读 einsum 方程时，我发现能够在精神上将它们归结为命令式版本是最有帮助的。

让我们从以下（强加的）声明开始：

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

首先处理标点符号，我们看到我们有两个 4 字母逗号分隔的 blob -bhwi和bhwj，在箭头之前，在它之后有一个 3 字母 blob bij。因此，该等式从两个 rank-4 张量输入产生一个 rank-3 张量结果。

现在，让每个 blob 中的每个字母成为范围变量的名称。字母出现在 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 公式如何映射到上面片段中的最终求和语句。for-loops 和 range bounds 只是绒毛，最后的语句是你真正需要了解发生了什么的全部内容。

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

# 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]

score 8 · Accepted Answer

我发现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]])
>>>

score 7 · Accepted Answer

让我们制作 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的最后一个B，第一个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更改为：Cijk

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]

score 7 · Accepted Answer

另一种观点`np.einsum`

这里的大多数答案都是通过例子来解释的，我想我会给出一个额外的观点。

您可以将einsum其视为广义矩阵求和运算符。给定的字符串包含下标，它们是表示轴的标签。我喜欢将其视为您的操作定义。下标提供了两个明显的约束：

每个输入数组的轴数，
输入之间的轴大小相等。

让我们以最初的例子为例：np.einsum('ij,jk->ki', A, B). 这里的约束1.转换为 A.ndim == 2and B.ndim == 2，以及2. to A.shape[1] == B.shape[0]。

正如您稍后将看到的，还有其他限制。例如：

输出下标中的标签不能出现多次。
输出下标中的标签必须出现在输入下标中。

在查看时ij,jk->ki，您可以将其视为：

输入数组中的哪些组件将有助于[k, i]输出数组的组件。

下标包含输出数组的每个组件的操作的确切定义。

我们将坚持 operation ，以及andij,jk->ki的以下定义：AB

>>> A = np.array([[1,4,1,7], [8,1,2,2], [7,4,3,4]])
>>> A.shape
(3, 4)

>>> B = np.array([[2,5], [0,1], [5,7], [9,2]])
>>> B.shape
(4, 2)

输出 ,Z的形状为，(B.shape[1], A.shape[0])并且可以通过以下方式简单地构造。从的空白数组开始Z：

Z = np.zeros((B.shape[1], A.shape[0]))
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        for k range(B.shape[0]):
           Z[k, i] += A[i, j]*B[j, k] # ki <- ij*jk

np.einsum是关于在输出数组中累积贡献。每一对都对每个组件(A[i,j], B[j,k])都有贡献。Z[k, i]

您可能已经注意到，它看起来与您计算一般矩阵乘法的方式非常相似......

最少的实施

这是np.einsumPython 中的最小实现。这应该有助于了解幕后的真实情况。

在我们继续进行的过程中，我将继续参考前面的示例。定义inputs为[A, B]。

np.einsum实际上可以接受两个以上的输入。下面，我们将关注一般情况：n 个输入和n 个输入下标。主要目标是找到迭代的域，即我们所有范围的笛卡尔积。

我们不能依赖手动编写for循环，仅仅因为我们不知道会有多少。主要思想是：我们需要找到所有唯一的标签（我将使用它们key并keys引用它们），找到相应的数组形状，然后为每个标签创建范围，并计算范围的乘积itertools.product以获得学习。

指数	`keys`	约束	`sizes`	`ranges`
1	`'i'`	`A.shape[0]`	3	`range(0, 3)`
2	`'j'`	`A.shape[1] == B.shape[0]`	4	`range(0, 4)`
0	`'k'`	`B.shape[1]`	2	`range(0, 2)`

研究领域是笛卡尔积：range(0, 2) x range(0, 3) x range(0, 4).

下标处理：

>>> expr = 'ij,jk->ki'
>>> qry_expr, res_expr = expr.split('->')
>>> inputs_expr = qry_expr.split(',')
>>> inputs_expr, res_expr
(['ij', 'jk'], 'ki')

在输入下标中找到唯一键（标签）：

>>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) 
               for key, size in list(zip(keys, input.shape))])
{('i', 3), ('j', 4), ('k', 2)}

我们应该检查约束（以及输出下标）！使用set是一个坏主意，但它适用于本示例的目的。

获取相关的大小（用于初始化输出数组）并构造范围（用于创建我们的迭代域）：

>>> sizes = dict(keys)
{'i': 3, 'j': 4, 'k': 2}

>>> ranges = [range(size) for _, size in keys]
[range(0, 2), range(0, 3), range(0, 4)]

我们需要一个包含键（标签）的列表：
```
>>> to_key = sizes.keys()
['k', 'i', 'j']
```
计算ranges的笛卡尔积
```
>>> domain = product(*ranges)
```
注意：返回一个随时间消耗[itertools.product][1]的迭代器。

将输出张量初始化为：

>>> res = np.zeros([sizes[key] for key in res_expr])

我们将循环domain：
```
>>> for indices in domain:
...     pass
```
对于每次迭代，indices将包含每个轴上的值。在我们的示例中，这将提供i,j和k作为元组: (k, i, j)。对于每个输入 (A和B)，我们需要确定要获取哪个组件。那是A[i, j]和B[j, k]，是的！但是，从字面上讲，我们没有变量i,j和。k

indices我们可以使用zipto_key来创建每个键（标签）与其当前值之间的映射：
```
>>> vals = dict(zip(to_key, indices))
```
要获取输出数组的坐标，我们使用vals并循环遍历键：[vals[key] for key in res_expr]。但是，要使用这些索引输出数组，我们需要将其包装tuple并zip沿每个轴分隔索引：
```
>>> res_ind = tuple(zip([vals[key] for key in res_expr]))
```
输入索引相同（尽管可以有多个）：
```
>>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr]
```
我们将使用 aitertools.reduce来计算所有贡献组件的乘积：
```
>>> def reduce_mult(L):
...     return reduce(lambda x, y: x*y, L)
```

总体而言，域上的循环如下所示：

>>> for indices in domain:
...     vals = {k: v for v, k in zip(indices, to_key)}
...     res_ind = tuple(zip([vals[key] for key in res_expr]))
...     inputs_ind = [tuple(zip([vals[key] for key in expr])) 
...                     for expr in inputs_expr]
...
...     res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])

>>> res
array([[70., 44., 65.],
       [30., 59., 68.]])

这与返回的结果非常接近np.einsum('ij,jk->ki', A, B)！

score 0 · Accepted Answer

我认为最简单的例子是在tensorflow 文档中

将方程转换为 einsum 符号有四个步骤。让我们以这个方程为例C[i,k] = sum_j A[i,j] * B[j,k]

首先我们删除变量名。我们得到ik = sum_j ij * jk
我们删除该sum_j术语，因为它是隐含的。我们得到ik = ij * jk
我们替换*为,. 我们得到ik = ij, jk
输出在 RHS 上并用->符号分隔。我们得到ij, jk -> ik

einsum 解释器只是反向运行这 4 个步骤。结果中缺少的所有索引都被求和。

以下是文档中的更多示例

# Matrix multiplication
einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Trace
einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]

# Batch matrix multiplication
einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]

score 0 · Accepted Answer

0

于 2021-02-15T10:39:48.267 回答

python - 了解 NumPy 的 einsum

8 回答 8

做什么einsum？

如何einsum工作？

一个稍微大一点的例子

一些练习

另一种观点np.einsum

最少的实施

Related

Reference

做什么`einsum`？

如何`einsum`工作？

另一种观点`np.einsum`