numpy.einsum

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

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

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

在 明确的 模式， einsum 通过在指定的下标标签上禁用或强制求和，可以进一步灵活地计算可能不被视为经典爱因斯坦求和操作的其他数组操作。

请参阅注释和示例以进行澄清。

参数

subscriptsSTR

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

operands数组列表

这些是操作的数组。

outndarray，可选

如果提供，计算将在此数组中完成。

dtype数据类型，无，可选

如果提供，则强制计算使用指定的数据类型。注意，你可能还得给一个更自由的 casting 参数以允许转换。默认值为无。

order'C'、'F'、'A'、'K'，可选

控制输出的内存布局。'C'表示它应该是C连续的。'“f”表示它应该是Fortran连续的，“a”表示如果输入都是“f”，“c”，则应该是“f”。k'表示应尽可能靠近布局，包括任意排列的轴。默认值是“k”。

casting'no'、'equiv'、'safe'、'same'u kind'、'unsafe'、可选

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

• “否”表示完全不应强制转换数据类型。

• “equiv”表示只允许更改字节顺序。

• “safe”表示只允许保留值的强制转换。

• “相同类型”意味着只允许安全的类型或类型内的类型，如float64到float32。

• “不安全”表示可以进行任何数据转换。

默认值为“安全”。

optimize假，真，'贪婪'，'最佳'，可选

控制是否应进行中间优化。如果“贪婪”算法默认为“假”和“真”，则不会进行优化。还接受来自 np.einsum_path 功能。见 np.einsum_path 了解更多详细信息。默认为false。

返回

output恩达雷

基于爱因斯坦求和约定的计算。

参见

einsum_path, dot, inner, outer, tensordot, linalg.multi_dot

einops

类似的详细接口由 einops 包括附加操作的软件包：转置、整形/展平、重复/平铺、挤压/取消挤压和缩小。

opt_einsum

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) 不影响二维数组，而 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, axis=-1) 和 np.einsum('ii->i', a) 就像 np.diag(a) . 区别在于 einsum 默认情况下不允许广播。另外 np.einsum('ij,jh->ih', a, b) 直接指定输出下标标签的顺序，因此返回矩阵乘法，与上面隐式模式中的示例不同。

要启用和控制广播，请使用省略号。默认的numpy样式的广播是通过在每个术语的左侧添加省略号来完成的，例如 np.einsum('...ii->...i', 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表达式的收缩顺序的参数。对于具有三个或更多操作数的收缩，这可以极大地提高计算效率，同时在计算过程中牺牲更大的内存占用。

通常采用“贪婪”算法，经验测试表明，在大多数情况下，该算法返回最佳路径。在某些情况下，“最优”会通过更昂贵、更详尽的搜索返回最高级的路径。对于迭代计算，建议只计算一次最佳路径，并通过提供该路径作为参数来重用该路径。下面给出一个例子。

见 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 ：~1520ms（以3.1GHz Intel i5为基准）

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

次优 einsum （由于重复路径计算时间）：~330ms

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

贪婪的 einsum （更快的最佳路径近似）：~160ms

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

最优的 einsum （某些用例中的最佳使用模式）：~110ms

>>> 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)

numpy.tensordot numpy.einsum_path