Python 努比·艾因苏姆的黑色巫毒

我得到了一些使用einsum函数的工作代码。但由于einsum目前对我来说仍然像是黑色巫毒。我想知道，这段代码实际上在做什么，是否可以使用

np.dot

我的数据看起来像这样

n, p, q = 40000, 8, 4
a = np.random.rand(n, p, q)
b = np.random.rand(n, p)

我现有的函数einsum函数如下所示

f1 = np.einsum("ijx,ijy->ixy", a, a)
f2 = np.einsum("ijx,ij->ix", a, b)

但它到底有什么作用呢？我一直到这里：每个维度（轴）都由一个标签表示，

I

等于第一个轴

n

，

j

对于第二个轴

p

和

x

和

y

对于同一个轴

q

是不同的标签。 因此，

f1

的输出数组的顺序是

ixy

，因此输出形状是

40000,4,4（n，q，q）

---

但这就是我能做到的。还有

让我们玩几个小数组

In [110]: a=np.arange(2*3*4).reshape(2,3,4)

In [111]: b=np.arange(2*3).reshape(2,3)

In [112]: np.einsum('ijx,ij->ix',a,b)
Out[112]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])

In [113]: np.diagonal(np.dot(b,a)).T
Out[113]: 
array([[ 20,  23,  26,  29],
       [200, 212, 224, 236]])

np.dot

在第一个阵列的最后一个dim上运行，在第二个阵列的最后一个dim上运行。因此，我必须切换参数，以便

3

维度对齐

dot（b，a）

生成一个（2,2,4）数组<代码>对角线选择其中的两行，并转置以清理。另一个

einsum

很好地表达了这一点：

In [122]: np.einsum('iik->ik',np.dot(b,a))

由于

np.dot

生成的数组比原始的

einsum

更大，因此不太可能更快，即使底层的C代码更紧

（奇怪的是，我在用

einsum

复制

np.dot（b，a）

时遇到了问题；它不会生成（2,2，…）数组）

对于

a，a

的情况，我们必须做类似的事情-滚动一个数组的轴，使最后一个维度与另一个数组的倒数第二个维度对齐，执行

点

，然后使用

对角线

和

转置

进行清理：

In [157]: np.einsum('ijx,ijy->ixy',a,a).shape
Out[157]: (2, 4, 4)
In [158]: np.einsum('ijjx->jix',np.dot(np.rollaxis(a,2),a))
In [176]: np.diagonal(np.dot(np.rollaxis(a,2),a),0,2).T

tensordot

是在选定轴上取点的另一种方式

np.tensordot(a,a,(1,1))
np.diagonal(np.rollaxis(np.tensordot(a,a,(1,1)),1),0,2).T  # with cleanup

---

维基百科关于的文章应该可以帮助你理解

einsum

的功能@菲利普：谢谢你的链接！这确实有帮助。但也让我更加意识到，我自己无法优化这一点，因为这超出了我的想象。

einsum

通常与

dot

一样好。两者都是经过编译的

dot

对于特定的维度组合有更严格的代码。这是一个很好的解释！要知道，

einsum

真的很强大，而且现在知道它在做什么，事实上更“可读”。至于你的评论，

np.einsum（'ij，kjl'，b，a）==np.dot（b，a）

。