maxframe.tensor.einsum#
- maxframe.tensor.einsum(subscripts, *operands, dtype=None, order='K', casting='safe', optimize=False)[源代码]#
在操作数上计算爱因斯坦求和约定。
使用爱因斯坦求和约定,许多常见的多维线性代数数组运算可以用简单的方式表示。在 隐式 模式下 einsum 计算这些值。
在 显式 模式下,einsum 提供了更大的灵活性来计算其他可能不被视为经典爱因斯坦求和运算的数组操作,通过禁用或强制对指定下标标签进行求和。
请参阅注释和示例以获取说明。
- 参数:
subscripts (str) -- 指定求和的下标为逗号分隔的下标标签列表。除非包含显式指示符 '->' 以及精确输出形式的下标标签,否则将执行隐式(经典爱因斯坦求和)计算。
operands (list of array_like) -- 这些是用于操作的数组。
dtype ({data-type, None}, optional) -- 如果提供,则强制计算使用指定的数据类型。请注意,您可能还需要给出更宽松的 casting 参数以允许转换。默认值为 None。
order ({'C', 'F', 'A', 'K'}, optional) -- 控制输出的内存布局。'C' 表示应为 C 连续。'F' 表示应为 Fortran 连续,'A' 表示如果输入都是 'F' 则为 'F',否则为 'C'。'K' 表示应尽可能接近输入的布局,包括任意排列的轴。默认值为 'K'。
casting ({'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional) -- 控制可能发生的数据类型转换。不建议将此设置为 '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。
- 返回:
output (maxframe.tensor.Tensor) -- 基于爱因斯坦求和约定的计算结果。
爱因斯坦求和约定可以用于计算
许多多维线性代数数组操作。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()。下标字符串是逗号分隔的下标标签列表,
其中每个标签对应于相应操作数的一个维度。
当标签重复时会进行求和,因此
mt.einsum('i,i', a, b)等价于
mt.inner(a,b)。如果一个标签只出现一次,则不会求和,因此
mt.einsum('i', a)生成一个a的视图且不作任何更改。另一个例子mt.einsum('ij,jk', a, b)描述传统的矩阵乘法,等价于
在 *隐式模式 中,所选的下标很重要*
因为输出的轴会按字母顺序重新排序。这
意味着
mt.einsum('ij', a)不会影响二维数组,而mt.einsum('ji', a)会对其转置。此外,mt.einsum('ij,jk', a, b)返回矩阵乘法的结果,而,mt.einsum('ij,jh', a, b)返回的是乘法结果的转置,因为下标 'h' 在 'i' 之前。
在 *显式模式 中,可以通过指定输出下标标签来直接控制输出*
这需要使用标识符 '->' 以及输出下标标签列表。
此功能增强了函数的灵活性,因为可以在需要时禁用或强制求和。调用
此功能增强了函数的灵活性,因为可以在需要时禁用或强制求和。调用
在必要时可以禁用或强制执行求和操作。例如
mt.einsum('i->', a)类似于mt.sum(a, axis=-1),而
mt.einsum('ii->i', a)类似于mt.diag(a)。区别在于 einsum 默认不允许广播。
此外,
mt.einsum('ij,jh->ih', a, b)直接指定了输出下标标签的顺序,因此返回矩阵
乘法,与上面隐式模式的示例不同。
要启用和控制广播,请使用省略号。默认
NumPy 风格的广播是通过添加省略号
在每个项的左侧完成的,例如
mt.einsum('...ii->...i', a)。要沿第一个和最后一个轴取迹,
可以使用
mt.einsum('i...i', a),或者要做矩阵-矩阵乘积但使用最左侧索引而不是最右侧索引,可以
使用
mt.einsum('ij...,jk...->ik...', a, b)。当只有一个操作数、没有轴被求和且未提供输出
参数时,将返回操作数的一个视图
而不是新数组。因此,将对角线提取为
mt.einsum('ii->i', a)会产生一个视图(在 1.10.0 版本中更改)。
einsum 还提供了另一种提供下标的方式
和操作数,形式为
einsum(op0, sublist0, op1, sublist1, ..., [sublistout])。如果未以这种格式提供输出形状,einsum 将
在隐式模式下计算,否则它将显式执行。
下面的示例都对应两种
参数方法的 einsum 调用。
示例
>>> import maxframe.tensor as mt >>> a = mt.arange(25).reshape(5,5) >>> b = mt.arange(5) >>> c = mt.arange(6).reshape(2,3) Trace of a matrix: >>> mt.einsum('ii', a).execute() 60 >>> mt.einsum(a, [0,0]).execute() 60 Extract the diagonal (requires explicit form): >>> mt.einsum('ii->i', a).execute() array([ 0, 6, 12, 18, 24]) >>> mt.einsum(a, [0,0], [0]).execute() array([ 0, 6, 12, 18, 24]) >>> mt.diag(a).execute() array([ 0, 6, 12, 18, 24]) Sum over an axis (requires explicit form): >>> mt.einsum('ij->i', a).execute() array([ 10, 35, 60, 85, 110]) >>> mt.einsum(a, [0,1], [0]).execute() array([ 10, 35, 60, 85, 110]) >>> mt.sum(a, axis=1).execute() array([ 10, 35, 60, 85, 110]) For higher dimensional arrays summing a single axis can be done with ellipsis: >>> mt.einsum('...j->...', a).execute() array([ 10, 35, 60, 85, 110]) >>> mt.einsum(a, [Ellipsis,1], [Ellipsis]).execute() array([ 10, 35, 60, 85, 110]) Compute a matrix transpose, or reorder any number of axes: >>> mt.einsum('ji', c).execute() array([[0, 3], [1, 4], [2, 5]]) >>> mt.einsum('ij->ji', c).execute() array([[0, 3], [1, 4], [2, 5]]) >>> mt.einsum(c, [1,0]).execute() array([[0, 3], [1, 4], [2, 5]]) >>> mt.transpose(c).execute() array([[0, 3], [1, 4], [2, 5]]) Vector inner products: >>> mt.einsum('i,i', b, b).execute() 30 >>> mt.einsum(b, [0], b, [0]).execute() 30 >>> mt.inner(b,b).execute() 30 Matrix vector multiplication: >>> mt.einsum('ij,j', a, b).execute() array([ 30, 80, 130, 180, 230]) >>> mt.einsum(a, [0,1], b, [1]).execute() array([ 30, 80, 130, 180, 230]) >>> mt.dot(a, b).execute() array([ 30, 80, 130, 180, 230]) >>> mt.einsum('...j,j', a, b).execute() array([ 30, 80, 130, 180, 230]) Broadcasting and scalar multiplication: >>> mt.einsum('..., ...', 3, c).execute() array([[ 0, 3, 6], [ 9, 12, 15]]) >>> mt.einsum(',ij', 3, c).execute() array([[ 0, 3, 6], [ 9, 12, 15]]) >>> mt.einsum(3, [Ellipsis], c, [Ellipsis]).execute() array([[ 0, 3, 6], [ 9, 12, 15]]) >>> mt.multiply(3, c).execute() array([[ 0, 3, 6], [ 9, 12, 15]]) Vector outer product: >>> mt.einsum('i,j', mt.arange(2)+1, b).execute() array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]) >>> mt.einsum(mt.arange(2)+1, [0], b, [1]).execute() array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]) >>> mt.outer(mt.arange(2)+1, b).execute() array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]) Tensor contraction: >>> a = mt.arange(60.).reshape(3,4,5) >>> b = mt.arange(24.).reshape(4,3,2) >>> mt.einsum('ijk,jil->kl', a, b).execute() array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]) >>> mt.einsum(a, [0,1,2], b, [1,0,3], [2,3]).execute() array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]) >>> mt.tensordot(a,b, axes=([1,0],[0,1])).execute() array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]) Writeable returned arrays (since version 1.10.0): >>> a = mt.zeros((3, 3)) >>> mt.einsum('ii->i', a)[:] = 1 >>> a.execute() array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) Example of ellipsis use: >>> a = mt.arange(6).reshape((3,2)) >>> b = mt.arange(12).reshape((4,3)) >>> mt.einsum('ki,jk->ij', a, b).execute() array([[10, 28, 46, 64], [13, 40, 67, 94]]) >>> mt.einsum('ki,...k->i...', a, b).execute() array([[10, 28, 46, 64], [13, 40, 67, 94]]) >>> mt.einsum('k...,jk', a, b).execute() array([[10, 28, 46, 64], [13, 40, 67, 94]]) Chained array operations. For more complicated contractions, speed ups might be achieved by repeatedly computing a 'greedy' path or pre-computing the 'optimal' path and repeatedly applying it, using an `einsum_path` insertion (since version 1.12.0). Performance improvements can be particularly significant with larger arrays: >>> a = mt.ones(64).reshape(2,4,8) Basic `einsum`: ~1520ms (benchmarked on 3.1GHz Intel i5.) >>> for iteration in range(500): ... _ = mt.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a) Sub-optimal `einsum` (due to repeated path calculation time): ~330ms >>> for iteration in range(500): ... _ = mt.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal') Greedy `einsum` (faster optimal path approximation): ~160ms >>> for iteration in range(500): ... _ = mt.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy') Optimal `einsum` (best usage pattern in some use cases): ~110ms >>> path = mt.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')[0] >>> for iteration in range(500): ... _ = mt.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)