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