NumPy Einsum的黑巫毒教

Question

我使用einsum函数获得了一些工作代码.但是因为einsum目前仍然black voodoo适合我.我想知道,这段代码实际上是做什么的,如果它可以以某种方式使用优化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)

但就我而言.和

Answer 1

让我们玩几个小数组

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对第一个数组的最后一个暗淡进行操作，对第二个数组的最后一个暗淡进行操作。所以我必须切换参数以使3尺寸对齐。 dot(b,a)生成 (2,2,4) 数组。 diagonal选择其中 2 个“行”，然后转置以进行清理。另一个einsum很好地表达了清理：

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

由于np.dot生成的数组比原始数组更大einsum，因此即使底层 C 代码更紧凑，它也不太可能更快。

np.dot(b,a)（奇怪的是，我在复制时遇到问题einsum；它不会生成 (2,2,...) 数组）。

对于这种a,a情况，我们必须做类似的事情 - 滚动一个数组的轴，使最后一个维度与另一个数组的倒数第二个维度对齐，执行dot，然后使用diagonal和进行清理transpose：

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是接管选定轴的另一种方法dot。

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