python - 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 代码更紧凑，它也不太可能更快。

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

对于 a,a 情况，我们必须做类似的事情 - 滚动一个数组的轴，使最后一个维度与另一个维度的倒数第二个对齐，执行 点，然后用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 是另一种在选定轴上绘制点的方法。

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