numpy.einsum#

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

对操作数进行爱因斯坦求和约定评估.

使用爱因斯坦求和约定,许多常见的多维、线性代数数组操作可以用一种简单的方式表示.在隐式模式下,`einsum` 计算这些值.

在显式模式中,`einsum` 提供了进一步的灵活性来计算其他可能不被认为是经典爱因斯坦求和操作的数组操作,通过禁用或强制对指定的下标标签进行求和.

请参阅注释和示例以获得澄清.

参数:

subscriptsstr

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

operands数组类列表

这些是操作的数组.

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’}, 可选

控制可能发生的数据转换类型.将其设置为 ‘不安全’ 是不推荐的,因为它可能会对累积产生不利影响.

‘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 <https://optimized-einsum.readthedocs.io/en/stable/>`_ 以不依赖后端的方式优化类似 einsum 表达式的收缩顺序.

备注

在 1.6.0 版本加入.

爱因斯坦求和约定可以用来计算许多多维、线性代数数组操作.`einsum` 提供了一种简洁的表示方法.

这些操作的一个非详尽列表,可以通过 einsum 计算,如下所示,并附有示例:

数组的迹, numpy.trace.
返回一个对角线,:py:func:numpy.diag.
数组轴求和, numpy.sum.
转置和排列,:py:func:numpy.transpose.
矩阵乘法和点积, numpy.matmul
numpy.dot.
向量内积和外积,:py:func: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) 将返回一个 2D 数组对角线的可写视图.

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