numpy.einsum#

numpy.einsum(subscripts, *operands, out=None, dtype=None, order='K', casting='safe', optimize=False)[source]#

对操作数评估爱因斯坦求和约定。

使用爱因斯坦求和约定，许多常见的多分量线性代数数组运算可以用简单的方式表示。在隐式模式下，einsum计算这些值。

在显式模式下，einsum提供了进一步的灵活性来计算可能不被认为是经典爱因斯坦求和运算的其他数组运算，方法是禁用或强制对指定的下标标签进行求和。

请参阅注释和示例以获得说明。

参数：

subscriptsstr

指定求和的下标，作为下标标签的逗号分隔列表。除非显式指示符“->”以及精确输出形式的下标标签也包含在内，否则将执行隐式（经典爱因斯坦求和）计算。

operandsarray_like 列表

这些是操作的数组。

outndarray，可选

如果提供，则计算结果将写入此数组。

dtype{数据类型，None}，可选

如果提供，则强制计算使用指定的数据类型。请注意，您可能还必须提供更宽松的casting参数以允许转换。默认为 None。

order{‘C’，‘F’，‘A’，‘K’}，可选

控制输出的内存布局。“C”表示它应该是 C 连续的。“F”表示它应该是 Fortran 连续的，“A”表示如果所有输入都是“F”，则它应该是“F”，否则为“C”。“K”表示它应该尽可能接近输入的布局，包括任意置换的轴。默认为“K”。

casting{‘no’，‘equiv’，‘safe’，‘same_kind’，‘unsafe’}，可选

控制可能发生哪种数据转换。不建议将此设置为“unsafe”，因为它可能会对累积产生不利影响。

“no”表示数据类型根本不应该转换。
“equiv”表示仅允许字节顺序更改。
“safe”表示仅允许可以保留值的转换。
“same_kind”表示仅允许安全转换或同类转换，例如 float64 到 float32。
“unsafe”表示可以执行任何数据转换。

默认为“safe”。

optimize{False，True，‘greedy’，‘optimal’}，可选

控制是否应发生中间优化。如果为 False，则不会发生优化，如果为 True，则默认为“greedy”算法。还接受来自np.einsum_path函数的显式收缩列表。有关更多详细信息，请参阅np.einsum_path。默认为 False。

返回值：

outputndarray: 基于爱因斯坦求和约定的计算结果。

另请参阅

einsum_path，dot，inner，outer，tensordot，linalg.multi_dot
einsum: einops 包提供了类似的详细接口来涵盖其他操作：转置、重塑/展平、重复/平铺、压缩/解压缩和约简。opt_einsum 以与后端无关的方式优化类似 einsum 表达式的收缩顺序。

注释

版本 1.6.0 中的新功能。

爱因斯坦求和约定可用于计算许多多分量线性代数数组运算。einsum 提供了一种简洁的方式来表示这些运算。

下面显示了可以通过 einsum 计算的这些操作的非详尽列表以及示例

数组的迹，numpy.trace。
返回对角线，numpy.diag。
数组轴求和，numpy.sum。
转置和排列，numpy.transpose。
矩阵乘法和点积，numpy.matmul
numpy.dot.
向量内积和外积，numpy.inner
numpy.outer.
广播、逐元素和标量乘法，
numpy.multiply.
张量收缩，numpy.tensordot。
链接数组操作，以高效的计算顺序，
numpy.einsum_path.

下标字符串是下标标签的逗号分隔列表，其中每个标签引用对应操作数的维度。每当标签重复时，它就会被求和，因此 np.einsum('i,i', a, b) 等效于 np.inner(a,b)。如果标签只出现一次，则不会求和，因此 np.einsum('i', a) 会产生 a 的视图，没有任何更改。另一个示例 np.einsum('ij,jk', a, b) 描述了传统的矩阵乘法，等效于 np.matmul(a,b)。一个操作数中重复的下标标签会取对角线。例如，np.einsum('ii', a) 等效于 np.trace(a)。

在隐式模式下，所选下标非常重要，因为输出的轴按字母顺序重新排序。这意味着 np.einsum('ij', a) 不会影响 2D 数组，而 np.einsum('ji', a) 会对其进行转置。此外，np.einsum('ij,jk', a, b) 返回矩阵乘法，而 np.einsum('ij,jh', a, b) 返回乘法的转置，因为下标“h”在“i”之前。

在显式模式下，可以通过指定输出下标标签来直接控制输出。这需要标识符“->”以及输出下标标签列表。此功能增加了函数的灵活性，因为可以在需要时禁用或强制求和。np.einsum('i->', a) 调用类似于 np.sum(a)，如果 a 是一个一维数组，而 np.einsum('ii->i', a) 类似于 np.diag(a)，如果 a 是一个方形的二维数组。不同之处在于 einsum 默认不允许广播。此外，np.einsum('ij,jh->ih', a, b) 直接指定了输出下标标签的顺序，因此返回矩阵乘法，与上面隐式模式中的示例不同。

要启用和控制广播，请使用省略号。默认的 NumPy 样式广播是通过在每个项的左侧添加省略号来完成的，例如 np.einsum('...ii->...i', a)。np.einsum('...i->...', a) 类似于 np.sum(a, axis=-1)，用于任何形状的数组 a。要沿着第一个和最后一个轴取迹，您可以执行 np.einsum('i...i', a)，或者要执行以最左侧索引而不是最右侧索引的矩阵-矩阵积，您可以执行 np.einsum('ij...,jk...->ik...', a, b)。

当只有一个操作数、没有轴被求和且没有提供输出参数时，将返回对操作数的视图而不是新数组。因此，将对角线作为 np.einsum('ii->i', a) 会产生一个视图（在版本 1.10.0 中更改）。

einsum 还提供了一种替代方法来提供下标和操作数，如 einsum(op0, sublist0, op1, sublist1, ..., [sublistout])。如果在此格式中未提供输出形状，则 einsum 将在隐式模式下计算，否则将显式执行。下面的示例具有使用两种参数方法的相应 einsum 调用。

版本 1.10.0 中的新功能。

只要输入数组是可写的，从 einsum 返回的视图现在就是可写的。例如，np.einsum('ijk...->kji...', a) 现在将具有与 np.swapaxes(a, 0, 2) 相同的效果，而 np.einsum('ii->i', a) 将返回二维数组对角线的可写视图。

版本 1.12.0 中的新功能。

添加了 optimize 参数，该参数将优化 einsum 表达式的收缩顺序。对于具有三个或更多操作数的收缩，这可以在计算过程中以更大的内存占用为代价，大大提高计算效率。

通常会应用一种“贪婪”算法，经验测试表明该算法在大多数情况下都能找到最佳路径。在某些情况下，“optimal”将通过更昂贵的穷举搜索返回最佳路径。对于迭代计算，建议计算一次最佳路径，然后通过将其作为参数提供来重复使用该路径。下面给出了一个示例。

有关更多详细信息，请参阅 numpy.einsum_path。

示例

>>> a = np.arange(25).reshape(5,5)
>>> b = np.arange(5)
>>> c = np.arange(6).reshape(2,3)

矩阵的迹

>>> np.einsum('ii', a)
60
>>> np.einsum(a, [0,0])
60
>>> np.trace(a)
60

提取对角线（需要显式形式）

>>> np.einsum('ii->i', a)
array([ 0,  6, 12, 18, 24])
>>> np.einsum(a, [0,0], [0])
array([ 0,  6, 12, 18, 24])
>>> np.diag(a)
array([ 0,  6, 12, 18, 24])

沿轴求和（需要显式形式）

>>> np.einsum('ij->i', a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [0,1], [0])
array([ 10,  35,  60,  85, 110])
>>> np.sum(a, axis=1)
array([ 10,  35,  60,  85, 110])

对于更高维度的数组，可以使用省略号对单个轴求和

>>> np.einsum('...j->...', a)
array([ 10,  35,  60,  85, 110])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 10,  35,  60,  85, 110])

计算矩阵转置或重新排序任意数量的轴

>>> np.einsum('ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum('ij->ji', c)
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.einsum(c, [1,0])
array([[0, 3],
       [1, 4],
       [2, 5]])
>>> np.transpose(c)
array([[0, 3],
       [1, 4],
       [2, 5]])

向量内积

>>> np.einsum('i,i', b, b)
30
>>> np.einsum(b, [0], b, [0])
30
>>> np.inner(b,b)
30

矩阵向量乘法

>>> np.einsum('ij,j', a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum(a, [0,1], b, [1])
array([ 30,  80, 130, 180, 230])
>>> np.dot(a, b)
array([ 30,  80, 130, 180, 230])
>>> np.einsum('...j,j', a, b)
array([ 30,  80, 130, 180, 230])

广播和标量乘法

>>> np.einsum('..., ...', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(',ij', 3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.einsum(3, [Ellipsis], c, [Ellipsis])
array([[ 0,  3,  6],
       [ 9, 12, 15]])
>>> np.multiply(3, c)
array([[ 0,  3,  6],
       [ 9, 12, 15]])

向量外积

>>> np.einsum('i,j', np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.einsum(np.arange(2)+1, [0], b, [1])
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])
>>> np.outer(np.arange(2)+1, b)
array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

张量收缩

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> np.einsum('ijk,jil->kl', a, b)
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
>>> np.einsum(a, [0,1,2], b, [1,0,3], [2,3])
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])
>>> np.tensordot(a,b, axes=([1,0],[0,1]))
array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

可写返回数组（自版本 1.10.0 起）

>>> a = np.zeros((3, 3))
>>> np.einsum('ii->i', a)[:] = 1
>>> a
array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

省略号用法的示例

>>> a = np.arange(6).reshape((3,2))
>>> b = np.arange(12).reshape((4,3))
>>> np.einsum('ki,jk->ij', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('ki,...k->i...', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])
>>> np.einsum('k...,jk', a, b)
array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

链式数组运算。对于更复杂的收缩，可以通过重复计算“贪婪”路径或预先计算“最优”路径并重复应用它来提高速度，使用 einsum_path 插入（自版本 1.12.0 起）。对于更大的数组，性能改进可能特别显著

>>> a = np.ones(64).reshape(2,4,8)

基本 einsum：约 1520 毫秒（在 3.1GHz Intel i5 上进行基准测试）。

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)

次优 einsum（由于重复路径计算时间）：约 330 毫秒

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a,
...         optimize='optimal')

贪婪 einsum（更快的最优路径近似）：约 160 毫秒

>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy')

最优 einsum（在某些用例中的最佳使用模式）：约 110 毫秒

>>> path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, 
...     optimize='optimal')[0]
>>> for iteration in range(500):
...     _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)