python - 写双(三)和作为内积？

Question

由于我的 np.dot 由 OpenBlas 和 Openmpi 加速，我想知道是否有可能写出双和

for i in range(N):
     for j in range(N):
         B[k,l] += A[i,j,k,l] * X[i,j]

作为内积。就在我正在使用的那一刻

B = np.einsum("ijkl,ij->kl",A,X)

但不幸的是它很慢并且只使用一个处理器。有什么想法吗？

编辑:我用一个简单的例子对到目前为止给出的答案进行了基准测试，看起来它们都处于相同的数量级:

A = np.random.random([200,200,100,100])
X = np.random.random([200,200])
def B1():
    return es("ijkl,ij->kl",A,X) 
def B2():
    return np.tensordot(A, X, [[0,1], [0, 1]])
def B3():
    shp = A.shape
    return np.dot(X.ravel(),A.reshape(shp[0]*shp[1],1)).reshape(shp[2],shp[3])

%timeit B1()
%timeit B2()
%timeit B3()

1 loops, best of 3: 300 ms per loop
10 loops, best of 3: 149 ms per loop
10 loops, best of 3: 150 ms per loop

从这些结果得出结论，我会选择 np.einsum，因为它的语法仍然是最易读的，并且与其他两个相比的改进仅为 2 倍。我想下一步是将代码外部化为 C 或 Fortran。

Answer 1

你可以使用np.tensordot():

np.tensordot(A, X, [[0,1], [0, 1]])

它确实使用了多个内核。

---

编辑:看到 np.einsum 和 np.tensordot 在增加输入数组的大小时如何缩放是很有趣的:

In [18]: for n in range(1, 31):
   ....:     A = np.random.rand(n, n+1, n+2, n+3)
   ....:     X = np.random.rand(n, n+1)
   ....:     print(n)
   ....:     %timeit np.einsum('ijkl,ij->kl', A, X)
   ....:     %timeit np.tensordot(A, X, [[0, 1], [0, 1]])
   ....:
1
1000000 loops, best of 3: 1.55 µs per loop
100000 loops, best of 3: 8.36 µs per loop
...
11
100000 loops, best of 3: 15.9 µs per loop
100000 loops, best of 3: 17.2 µs per loop
12
10000 loops, best of 3: 23.6 µs per loop
100000 loops, best of 3: 18.9 µs per loop
...
21
10000 loops, best of 3: 153 µs per loop
10000 loops, best of 3: 44.4 µs per loop

对于更大的数组，使用 tensordot 的优势变得很明显。