numpy.einsum

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

在操作数上计算爱因斯坦求和约定.

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

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

有关详细说明,请参见注释和示例.

参数:

subscriptsstr

指定用于求和的下标, 它是以下标标签的逗号分隔列表.除非显式指示符"->"以及精确输出形式的下标标签也被包括在内,否则将执行隐式(经典爱因斯坦求和)计算.

operandslist of array_like

这些是运算的数组.

outndarray,可选

如果提供,则在此数组中完成计算.

dtype{data-type, None}, optional

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

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

控制输出的内存布局."C"表示它应该是 C contiguous."F"表示它应该是 Fortran contiguous,"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’}, optional

控制是否应进行中间优化. 如果为 False,则不会进行优化,如果为 True,则默认使用"greedy"算法. 也接受来自 np.einsum_path 函数的显式收缩列表. 有关更多详细信息,请参见 np.einsum_path . 默认为 False.

返回:

outputndarray

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

参见

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

einsum

einops 包提供了类似的详细接口,以涵盖其他操作:转置,reshape/flatten,repeat/tile,squeeze/unsqueeze 和 reductions. opt_einsum 以后端无关的方式优化 einsum 类表达式的收缩顺序.

注释

爱因斯坦求和约定可用于计算许多多维线性代数数组运算. 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’ 之前.

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

要启用和控制广播,请使用省略号. 默认的 NumPy 风格广播是通过在每个术语的左侧添加一个省略号来完成的,例如 np.einsum('...ii->...i', a) . 对于任何形状的数组 a , np.einsum('...i->...', a) 类似于 np.sum(a, axis=-1) . 要沿着第一个和最后一个轴取迹,可以执行 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 调用.

现在,每当输入数组可写时,从 einsum 返回的视图都是可写的. 例如, np.einsum('ijk...->kji...', a) 现在将具有与 np.swapaxes(a, 0, 2) 相同的效果,并且 np.einsum('ii->i', a) 将返回二维数组对角线的可写视图.

增加了 optimize 参数,该参数将优化 einsum 表达式的收缩顺序.对于具有三个或更多操作数的收缩,这可以大大提高计算效率,但代价是在计算期间占用更大的内存.

通常应用"贪婪"算法,经验测试表明,该算法在大多数情况下都会返回最佳路径.在某些情况下,"optimal"将通过更昂贵的,详尽的搜索返回最佳路径.对于迭代计算,建议计算一次最佳路径,并通过将其作为参数提供来重用该路径.下面给出了一个例子.

有关更多详细信息,请参见 numpy.einsum_path .

示例

Try it in your browser!

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

链式数组运算.对于更复杂的收缩,可以通过重复计算"贪婪"路径或预先计算"optimal"路径并重复应用它,使用 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)

Go BackOpen In Tab