numpy-einsum表示法：一组矩阵与一组向量的点积

Question

我想用n维m*m矩阵堆栈乘以n维向量堆栈（长度m），这样得到的m*n数组包含第n项中矩阵和向量的点乘结果：

vec1=np.array([0,0.5,1,0.5]); vec2=np.array([2,0.5,1,0.5])
vec=np.transpose(n.stack((vec1,vec2)))
mat = np.moveaxis(n.array([[[0,1,2,3],[0,1,2,3],[0,1,2,3],[0,1,2,3]],[[-1,2.,0,1.],[0,0,-1,2.],[0,1,-1,2.],[1,0.1,1,1]]]),0,2)
outvec=np.zeros((4,2))
for i in range(2):
    outvec[:,i]=np.dot(mat[:,:,i],vec[:,i])

受这篇文章Element wise dot product of matrices and vectors的启发，我尝试了einsum中所有不同的索引组合扰动，并发现

^{pr2}$

给出正确的结果。在

不幸的是，我真的不明白这一点-我假设我重复'ijk，jk'部分中的条目k意味着我乘以k并求和。我试图阅读文档https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.einsum.html，但我仍然不明白。在

（我之前的尝试包括

 np.einsum('ijk,il->ik', mat, vec)

我甚至不知道这意味着什么。当我放弃索引l时，它会发生什么变化？）在

提前谢谢！在

Answer 1

einsum很简单（当你玩了一段时间的索引排列时，就是…）。在

让我们来做一些简单的事情，一个三层的2×2矩阵和一个2×，数组的三层堆栈

import numpy as np

a = np.arange(3*2*2).reshape((3,2,2))
b = np.arange(3*2).reshape((3,2))

我们需要知道我们要用einsum来计算什么

^{pr2}$

我们做了什么？我们有一个索引i，当我们在其中一个堆叠矩阵与其中一个堆叠向量（两者都以i为索引）之间执行点积时，该索引是固定的，并且单个输出行表示对堆叠矩阵的最后一个索引和堆叠向量的唯一索引的求和。在

这很容易在einsum指令中编码

• 我们需要相同的i索引来指定矩阵、向量和输出
• 我们要沿着最后一个矩阵索引和剩余的向量索引减少，比如k
• 我们希望输出中的列数与每个堆栈矩阵中的行数相等，比如j

因此

In [102]: np.einsum('ijk,ik->ij', a, b)                                                                   
Out[102]: 
array([[ 1,  3],
       [23, 33],
       [77, 95]])

我希望我关于如何正确执行指令的讨论是清楚、正确和有用的。在

Answer 2

阅读Einstein summation notation。在

基本上，规则是：

没有->

• 输入中重复的任何字母都代表一个要乘和求和的轴
• 输入中没有重复的任何字母都包含在输出中

使用->

• 输入中重复的任何字母都表示要乘以的轴
• 输出中不包含的任何字母都表示要求和的轴

例如，对于形状相同的矩阵A和{}：

np.einsum('ij, ij',       A, B)  # is A ddot B,                returns 0d scalar
np.einsum('ij, jk',       A, B)  # is A dot  B,                returns 2d tensor
np.einsum('ij, kl',       A, B)  # is outer(A, B),             returns 4d tensor
np.einsum('ji, jk, kl',   A, B)  # is A.T @ B @ A,             returns 2d tensor
np.einsum('ij, ij -> ij', A, B)  # is A * B,                   returns 2d tensor
np.einsum('ij, ij -> i' , A, A)  # is norm(A, axis = 1),       returns 1d tensor
np.einsum('ii'             , A)  # is tr(A),                   returns 0d scalar

Answer 3

In [321]: vec1=np.array([0,0.5,1,0.5]); vec2=np.array([2,0.5,1,0.5])
     ...: vec=np.transpose(np.stack((vec1,vec2)))
In [322]: vec1.shape
Out[322]: (4,)
In [323]: vec.shape
Out[323]: (4, 2)

stack函数的一个优点是我们可以指定一个轴，跳过转置：

^{pr2}$

为什么np.和{}的混合？NameError: name 'n' is not defined。那种事差点把我送走。在

In [326]: mat = np.moveaxis(np.array([[[0,1,2,3],[0,1,2,3],[0,1,2,3],[0,1,2,3]],[[-1,2.,0
     ...: ,1.],[0,0,-1,2.],[0,1,-1,2.],[1,0.1,1,1]]]),0,2)
In [327]: mat.shape
Out[327]: (4, 4, 2)

In [328]: outvec=np.zeros((4,2))
     ...: for i in range(2):
     ...:     outvec[:,i]=np.dot(mat[:,:,i],vec[:,i])
     ...:     
In [329]: outvec
Out[329]: 
array([[ 4.  , -0.5 ],
       [ 4.  ,  0.  ],
       [ 4.  ,  0.5 ],
       [ 4.  ,  3.55]])

In [330]: # (4,4,2) (4,2)   'kji,ji->ki'

在循环中，i轴（大小2）的位置在所有3个数组中都是最后一个。只留下一个轴给vec，我们将其称为j。它与最后一个（紧挨着i的mat）配对。k从mat转移到{}。在

In [331]: np.einsum('kji,ji->ki', mat, vec)
Out[331]: 
array([[ 4.  , -0.5 ],
       [ 4.  ,  0.  ],
       [ 4.  ,  0.5 ],
       [ 4.  ,  3.55]])

通常，einsum字符串会自己写入。例如，mat被描述为（m，n，k），而{}被描述为（n，k），结果是（m，k）

在本例中，只有j维被求和-它出现在左边，但在右边。最后一个维度，i在我的符号中，没有求和，因为if出现在两边，就像在迭代中一样。我认为这是“顺其自然”。它不是dot产品的积极组成部分。在

实际上，你是在最后一个维度上堆积，尺寸2。通常我们把第一个放在第一个上，但是你把两者都换过来放最后一个。在

---

您的“失败”尝试运行，并且可以复制为：

In [332]: np.einsum('ijk,il->ik', mat, vec)
Out[332]: 
array([[12. ,  4. ],
       [ 6. ,  1. ],
       [12. ,  4. ],
       [ 6. ,  3.1]])
In [333]: mat.sum(axis=1)*vec.sum(axis=1)[:,None]
Out[333]: 
array([[12. ,  4. ],
       [ 6. ,  1. ],
       [12. ,  4. ],
       [ 6. ,  3.1]])

j和l维度没有出现在右侧，因此它们被相加。它们可以在相乘之前求和，因为它们只出现在一个项中。我添加了None以启用广播（将ik与i相乘）。在

np.einsum('ik,i->ik', mat.sum(axis=1), vec.sum(axis=1))

---

如果在第一个上堆叠，并为vec（2,4,1）添加一个维度，那么它将使用一个（2,4,4）mat matmul。mat @ vec[...,None]。在

In [337]: m1 = mat.transpose(2,0,1)
In [338]: m1@v1[...,None]
Out[338]: 
array([[[ 4.  ],
        [ 4.  ],
        [ 4.  ],
        [ 4.  ]],

       [[-0.5 ],
        [ 0.  ],
        [ 0.5 ],
        [ 3.55]]])
In [339]: _.shape
Out[339]: (2, 4, 1)