【导读】 einsum全称Einstein summation convention（爱因斯坦求和约定），又称为爱因斯坦标记法。 能够计算任何维度的张量收缩。einsum的写法省去了求和符号，显得更加简洁。
对于一维张量，也即向量（Vector）。其收缩为零维张量，也即标量（Scalar）：
c = ∑ i a i b i c =\sum_i a_i b_i c =i ∑a i b i 所以einsum 的写法就是：
c = a i b i c= a_i b_i c =a i b i 用代码表示为：
>>> a = np.arange(10)
>>> b = np.arange(10)+1
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> b
array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> np.einsum('i,i', a, b)
330
同样 对于二维张量，也即矩阵（Matrix）。其收缩为一维张量或零维张量：
2.1 收缩到零维张量
c = a i j a i j = ∑ i , j a i j a i j c=a_{ij }a_{ij}= \sum_{i,j} a_{ij }a_{ij}c =a ij a ij =i ,j ∑a ij a ij 也即求矩阵内积
array([[0, 1],
[2, 3],
[4, 5]])
>>> np.einsum('ij,ij', a, b)
70
2.2 收缩到一维张量
c i = ∑ j A i j B i j c_i = \sum_j A_{ij} B_{ij}c i =j ∑A ij B ij 或c j = ∑ i A i j B i j c_j = \sum_i A_{ij} B_{ij}c j =i ∑A ij B ij
其中，c i c_i c i 为两矩阵对应位置元素相乘后再把列元素相加；c j c_j c j 为两矩阵对应位置元素相乘后把行相加
>>> np.einsum('ij,ij->i', a, b)
array([ 8, 62])
>>> np.einsum('ij,ij->j', a, b)
array([12, 22, 36])
再给出一个例子
>>> a1 = np.array([[0,1],[2,3],[4,5]])
>>> a1
array([[0, 1],
[2, 3],
[4, 5]])
>>> np.einsum('ji,ij->i', a1, a)
array([10, 40])
再给出一个矩阵与向量的例子
>>> a2 = np.array([[1,2],[3,4],[5,6]])
>>> b2 = np.array([1,2])
>
>>> np.einsum('ij,j', a2, b2)
array([ 5, 11, 17])
>>> np.einsum('ij,j->...', a2, b2)
33
再给出一个多个向量的例子
>>> c2 = np.array([1,2,3])
>>> np.einsum('ij,j,i->i', a2, b2, c2)
array([ 5, 22, 51])
>>> np.einsum('ij,j,i->...', a2, b2, c2)
78
由于其组合非常多样，如三维张量与矩阵、向量的爱因斯坦求和运算，收缩到二维、一维、零维，我们不一一讨论。这里我们只给出三维张量之间进行爱因斯坦求和、收缩到二维张量的例子
3.1 例1
>>> a = np.arange(6).reshape(1,2,3)
>>> b = np.arange(24).reshape(2,3,4)
>>> a
array([[[0, 1, 2],
[3, 4, 5]]])
>>> b
array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
>>> np.einsum('ijk,jkl->il', a, b)
array([[220, 235, 250, 265]])
其中，i = 1 , j = 2 , k = 3 , l = 4 i=1\;,j=2\;,k=3\;,l=4 i =1 ,j =2 ,k =3 ,l =4，可知会产生1 × 4 1\times4 1 ×4的二维张量（矩阵）。
由公式：
C i , l = a i j k b j k l = ∑ j = 0 1 ∑ k = 0 2 a i j k b j k l C_{i,l} = a_{ijk}b_{jkl}=\sum_{j=0}^1 \sum_{k=0}^2a_{ijk}b_{jkl}C i ,l =a ijk b jk l =j =0 ∑1 k =0 ∑2 a ijk b jk l
不失一般性，我们验证：C 0 , 0 = 220 C_{0,0} =220 C 0 ,0 =220
C 0 , 0 = a 000 ∗ b 000 + a 001 ∗ b 010 + a 002 ∗ b 020 + a 010 ∗ b 100 + a 011 ∗ b 110 + a 012 ∗ b 120 \begin{aligned}C_{0,0}=a_{000}b_{000}+a_{001}b_{010}+a_{002}b_{020}\+a_{010}b_{100}+a_{011}b_{110}+a_{012}b_{120}\end{aligned}C 0 ,0 =a 000 ∗b 000 +a 001 ∗b 010 +a 002 ∗b 020 +a 010 ∗b 100 +a 011 ∗b 110 +a 012 ∗b 120
即
C 0 , 0 = 0 ∗ 0 + 1 ∗ 4 + 2 ∗ 8 + 3 ∗ 12 + 4 ∗ 16 + 5 ∗ 20 = 220 C_{0,0}=00+14+28+312+416+520=220 C 0 ,0 =0 ∗0 +1 ∗4 +2 ∗8 +3 ∗12 +4 ∗16 +5 ∗20 =220
矩阵形式的计算过程：
[ 0 1 2 ] [ 0 1 2 3 4 5 6 7 8 9 10 11 ] = [ 20 23 26 29 ] \begin{bmatrix}0&1&2\end{bmatrix}\begin{bmatrix} 0&1&2&3\4&5&6&7\8&9&10&11\end{bmatrix}=\begin{bmatrix}20&23&26&29\end{bmatrix}[0 1 2 ]⎣⎡0 4 8 1 5 9 2 6 10 3 7 11 ⎦⎤=[20 23 26 29 ]
[ 3 4 5 ] [ 12 13 14 15 16 17 18 19 20 21 22 23 ] = [ 200 212 224 236 ] \begin{bmatrix}3&4&5\end{bmatrix}\begin{bmatrix} 12&13&14&15\16&17&18&19\20&21&22&23\end{bmatrix}=\begin{bmatrix}200&212&224& 236\end{bmatrix}[3 4 5 ]⎣⎡12 16 20 13 17 21 14 18 22 15 19 23 ⎦⎤=[200 212 224 236 ]
[ 220 235 250 265 ] = [ 20 23 26 9 ] + [ 200 212 224 236 ] \begin{bmatrix}220&235&250&265\end{bmatrix} =\begin{bmatrix}20&23&26&9\end{bmatrix} +\begin{bmatrix}200&212&224& 236\end{bmatrix}[220 235 250 265 ]=[20 23 26 9 ]+[200 212 224 236 ]
对于更高维的张量收缩或大规模张量，不方便再用矩阵表示计算过程表示，建议直接由公式计算相关元素。
再给出一个复杂的例子：
3.2 例2
>>> 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.]])
该例中i = 3 , j = 4 , k = 5 , l = 2 i=3\;,j=4\;,k=5\;,l=2 i =3 ,j =4 ,k =5 ,l =2，可知会产生5 × 2 5\times2 5 ×2的二维张量
由公式：
C k , l = a i j k b j i l = ∑ i = 0 2 ∑ j = 0 3 a i j k b j i l C_{k,l} = a_{ijk}b_{jil}=\sum_{i=0}^2 \sum_{j=0}^3a_{ijk}b_{jil}C k ,l =a ijk b ji l =i =0 ∑2 j =0 ∑3 a ijk b ji l
同样，不失一般性验证C 0 , 0 = 4400 C_{0,0}=4400 C 0 ,0 =4400
C 0 , 0 = a 000 ∗ b 000 + a 010 ∗ b 100 + a 020 ∗ b 200 + a 030 ∗ b 300 + a 100 ∗ b 010 + a 110 ∗ b 110 + a 120 ∗ b 210 + a 130 ∗ b 310 + a 200 ∗ b 020 + a 210 ∗ b 120 + a 220 ∗ b 220 + a 230 ∗ b 320 \begin{aligned}C_{0,0}=a_{000}b_{000}+a_{010}b_{100}+a_{020}b_{200}+a_{030}b_{300}\+a_{100}b_{010}+a_{110}b_{110}+a_{120}b_{210}+a_{130}b_{310}\+a_{200}b_{020}+a_{210}b_{120}+a_{220}b_{220}+a_{230}b_{320}\end{aligned}C 0 ,0 =a 000 ∗b 000 +a 010 ∗b 100 +a 020 ∗b 200 +a 030 ∗b 300 +a 100 ∗b 010 +a 110 ∗b 110 +a 120 ∗b 210 +a 130 ∗b 310 +a 200 ∗b 020 +a 210 ∗b 120 +a 220 ∗b 220 +a 230 ∗b 320
也即
C 0 , 0 = 0 ∗ 0 + 5 ∗ 6 + 10 ∗ 12 + 15 ∗ 18 + 20 ∗ 2 + 25 ∗ 8 + 30 ∗ 14 + 35 ∗ 20 + 40 ∗ 4 + 45 ∗ 10 + 50 ∗ 16 + 55 ∗ 22 = 4400 \begin{aligned}C_{0,0} &=00+56+1012+1518\&+202+258+3014+3520\&+404+4510+5016+5522\&=4400\end{aligned}C 0 ,0 =0 ∗0 +5 ∗6 +10 ∗12 +15 ∗18 +20 ∗2 +25 ∗8 +30 ∗14 +35 ∗20 +40 ∗4 +45 ∗10 +50 ∗16 +55 ∗22 =4400
功能一：求矩阵的迹
>>> a = np.arange(25).reshape(5,5)
>>> a
>>> array([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])
>>> b = np.arange(5)
>>> b
array([0, 1, 2, 3, 4])
>>> c = np.arange(6).reshape(2,3)
>>> c
array([[0, 1, 2],
[3, 4, 5]])
Trace of a matrix:
>>> np.einsum('ii', a)
60
>>> np.einsum(a, [0,0])
60
>>> np.trace(a)
60
功能二：取对角元素
Extract the diagonal (requires explicit form):
>>> 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])
功能三 ：对某个维度求和
Sum over an axis (requires explicit form):
>>> 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])
对于更高维度的矩阵，对某一维度求和可以使用省略号
>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 9, 10, 11],
[12, 13, 14],
[15, 16, 17]],
[[18, 19, 20],
[21, 22, 23],
[24, 25, 26]]])
>>> np.einsum('...j->...', a)
array([[ 3, 12, 21],
[30, 39, 48],
[57, 66, 75]])
功能四：矩阵转置
>>> c = np.arange(6).reshape(2,3)
>>> c
array([[0, 1, 2],
[3, 4, 5]])
>>> 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]])
功能五：矩阵或向量的内积
>>> b = np.arange(5)
>>> b
array([0, 1, 2, 3, 4])
>>> np.einsum('i,i', b, b)
30
>>> np.einsum(b, [0], b, [0])
30
>>> np.inner(b,b)
30
功能六：矩阵和向量乘法
>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
[3, 4, 5]])
>>> b = np.arange(3)
>>> b
array([0, 1, 2])
>>> np.einsum('ij,j', a, b)
array([ 5, 14])
>>> np.dot(a,b)
array([ 5, 14])
为减少不必要的重复，接下来代码中不再说明同样功能的其他函数或enisum函数的另一种使用形式
如：
>>> np.einsum('...j->...', a)
array([ 3, 12])
>>> np.einsum(a, [Ellipsis,1], [Ellipsis])
array([ 3, 12])
功能七：向量外积
>>> a
array([0, 1, 2])
>>>
>>> b
array([1, 2, 3])
>>> np.einsum('i,j', a, b)
array([[0, 0, 0],
[1, 2, 3],
[2, 4, 6]])
Original: https://blog.csdn.net/qq_40806950/article/details/125900065
Author: Yuriy_Xiong
Title: Numpy库中einsum函数用法
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/762796/
转载文章受原作者版权保护。转载请注明原作者出处！