einsum

paddle. einsum ( equation, *operands ) [源代码]

该函数用于对一组输入 Tensor 进行 Einstein 求和,该函数目前仅适用于动态图

Einstein 求和是一种采用 Einstein 标记法描述的 Tensor 求和,输入单个或多个 Tensor,输出单个 Tensor。

如下的 Tensor 操作或运算均可视为 Einstein 求和的特例

  • 单操作数
    • 迹:trace

    • 对角元:diagonal

    • 转置:transpose

    • 求和:sum

  • 双操作数
    • 内积:dot

    • 外积:outer

    • 广播乘积:mul,*

    • 矩阵乘:matmul

    • 批量矩阵乘:bmm

  • 多操作数
    • 广播乘积:mul,*

    • 多矩阵乘:A.matmul(B).matmul(C)

关于求和标记的约定

  • 维度分量下标:Tensor 的维度分量下标使用英文字母表示,不区分大小写,如'ijk'表示 Tensor 维度分量为 i,j,k

  • 下标对应输入操作数:维度下标以`,`分段,按顺序 1-1 对应输入操作数

  • 广播维度:省略号`...`表示维度的广播分量,例如,'i...j'表示首末分量除外的维度需进行广播对齐

  • 自由标和哑标:输入标记中仅出现一次的下标为自由标,重复出现的下标为哑标,哑标对应的维度分量将被规约消去

  • 输出:输出 Tensor 的维度分量既可由输入标记自动推导,也可以用输出标记定制化
    • 自动推导输出
      • 广播维度分量位于维度向量高维位置,自由标维度分量按字母顺序排序,位于维度向量低纬位置,哑标维度分量不输出

    • 定制化输出
      • 维度标记中`->`右侧为输出标记

      • 若输出包含广播维度,则输出标记需包含`...`

      • 输出标记为空时,对输出进行全量求和,返回该标量

      • 输出不能包含输入标记中未出现的下标

      • 输出下标不可以重复出现

      • 哑标出现在输出标记中则自动提升为自由标

      • 输出标记中未出现的自由标被降为哑标

  • 例子
    • '...ij, ...jk',该标记中 i,k 为自由标,j 为哑标,输出维度'...ik'

    • 'ij -> i',i 为自由标,j 为哑标

    • '...ij, ...jk -> ...ijk',i,j,k 均为自由标

    • '...ij, ...jk -> ij',若输入 Tensor 中的广播维度不为空,则该标记为无效标记

求和规则

Einsum 求和过程理论上等价于如下四步,但实现中实际执行的步骤会有差异。

  • 第一步,维度对齐:将所有标记按字母序排序,按照标记顺序将输入 Tensor 逐一转置、补齐维度,使得处理后的所有 Tensor 其维度标记保持一致

  • 第二步,广播乘积:以维度下标为索引进行广播点乘

  • 第三步,维度规约:将哑标对应的维度分量求和消除

  • 第四步,转置输出:若存在输出标记,则按标记进行转置,否则按广播维度+字母序自由标的顺序转置,返回转之后的 Tensor 作为输出

关于 trace 和 diagonal 的标记约定(待实现功能)

  • 在单个输入 Tensor 的标记中重复出现的下标称为对角标,对角标对应的坐标轴需进行对角化操作,如'i...i'表示需对首尾坐标轴进行对角化

  • 若无输出标记或输出标记中不包含对角标,则对角标对应维度规约为标量,相应维度取消,等价于 trace 操作

  • 若输出标记中包含对角标,则保留对角标维度,等价于 diagonal 操作

参数

equation (str):求和标记

operands (Tensor, [Tensor, ...]):输入 Tensor

返回

Tensor:输出 Tensor

代码示例

>>> import paddle
>>> paddle.seed(102)
>>> x = paddle.rand([4])
>>> y = paddle.rand([5])

>>> # sum
>>> print(paddle.einsum('i->', x))
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
1.81225157)

>>> # dot
>>> print(paddle.einsum('i,i->', x, x))
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
1.13530672)

>>> # outer
>>> print(paddle.einsum("i,j->ij", x, y))
Tensor(shape=[4, 5], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[0.26443148, 0.05962684, 0.25360870, 0.21900642, 0.56994802],
        [0.20955276, 0.04725220, 0.20097610, 0.17355499, 0.45166403],
        [0.35836059, 0.08080698, 0.34369346, 0.29680005, 0.77240014],
        [0.00484230, 0.00109189, 0.00464411, 0.00401047, 0.01043695]])

>>> A = paddle.rand([2, 3, 2])
>>> B = paddle.rand([2, 2, 3])

>>> # transpose
>>> print(paddle.einsum('ijk->kji', A))
Tensor(shape=[2, 3, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[[0.50882483, 0.56067896],
         [0.84598064, 0.36310029],
         [0.55289471, 0.33273944]],
        [[0.04836850, 0.73811269],
         [0.29769155, 0.28137168],
         [0.84636718, 0.67521429]]])

>>> # batch matrix multiplication
>>> print(paddle.einsum('ijk, ikl->ijl', A,B))
Tensor(shape=[2, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[[0.36321065, 0.42009076, 0.40849245],
         [0.74353045, 0.79189068, 0.81345987],
         [0.90488225, 0.79786193, 0.93451476]],
        [[0.12680580, 1.06945944, 0.79821426],
         [0.07774551, 0.55068684, 0.44512171],
         [0.08053084, 0.80583858, 0.56031936]]])

>>> # Ellipsis transpose
>>> print(paddle.einsum('...jk->...kj', A))
Tensor(shape=[2, 2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[[0.50882483, 0.84598064, 0.55289471],
         [0.04836850, 0.29769155, 0.84636718]],
        [[0.56067896, 0.36310029, 0.33273944],
         [0.73811269, 0.28137168, 0.67521429]]])

>>> # Ellipsis batch matrix multiplication
>>> print(paddle.einsum('...jk, ...kl->...jl', A,B))
Tensor(shape=[2, 3, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
       [[[0.36321065, 0.42009076, 0.40849245],
         [0.74353045, 0.79189068, 0.81345987],
         [0.90488225, 0.79786193, 0.93451476]],
        [[0.12680580, 1.06945944, 0.79821426],
         [0.07774551, 0.55068684, 0.44512171],
         [0.08053084, 0.80583858, 0.56031936]]])