python中两个以上矩阵的快速内积

Question

我现在正在写代码，需要尽可能快地计算三个二维数组之间的内积。你知道吗

我们叫它们a，b，c。它们都有相同的尺寸（N x M）。你知道吗

我想计算以下三维数组op，大小为（nxnnxn），这样op[I，j，k]就是a[I，m]b[j，m]c[k，m]的和

(click here for the nice Latex formula)

这基本上是将np.inner扩展到3个输入，而不是2个。你知道吗

实际上，我将遇到的维数是N=100和M=300000。矩阵不会是稀疏的，所以op包含大约一百万个非零值。你知道吗

到目前为止，我尝试了两种方法。 第一种使用广播：

import numpy as np
N = 100
M = 300000
a = np.random.randn(N, M)
b = np.random.randn(N, M)
c = np.random.randn(N, M)

def method1(a, b, c):
    a_i = a[:, None, None, :]
    b_j = b[None, :, None, :]
    c_k = c[None, None, :, :]
    return np.sum(a_i * b_j * c_k, axis=3)

问题是，它首先计算a_i * b_j * c_k，这是一个nxnxnmarray，所以在我的例子中，它太难处理了。你知道吗

我尝试了另一种使用np.einsum的方法，它比以前的方法快得多：

def method2(a, b, c):
    return np.einsum('im,jm,km', a, b, c)

我的问题是它还是太慢了。对于N=100和M=30000，在我的计算机上运行已经需要95秒了，所以把M取到它的实际值300000是不可能的。你知道吗

我的问题是：你知道用什么方法来解决我的问题吗（也许是一个神奇的numpy函数？），或者我必须求助于cython或numba之类的东西才能真正使这个计算可行吗？你知道吗

提前感谢您的帮助！你知道吗

Answer 1

非常有趣的一个和^{}有关。你知道吗

方法#1：对于适当大小的阵列

基于上述Q&A的获胜方法，这里有一个解决方案-

np.tensordot(a[:,None]*b,c,axes=(2,1))

说明：

1）a[:,None]*b：获取形状(N, N, M)的3D数组。因此，对于用例来说，它应该是(100, 100, 30000)，这对于常规系统来说可能有点太多了，但是如果有一些额外的系统内存juice，它可能就可以工作了。你知道吗

2）np.tensordot(..)：下一步，我们将上一步的最后一个轴与第三个数组c进行求和，从而得到一个(100, 100, 100)形状的输出数组。你知道吗

---

方法#2：对于非常大的阵列，b与c相同

out = np.zeros((N, N, N))
for i in range(N):
    for j in range(N):
        for k in range(j+1):
            out[i,j,k] = np.einsum('i,i,i->',a[i],b[j],b[k])

r,c = np.triu_indices(N,1)
out[np.arange(N)[:,None], r,c] = out[np.arange(N)[:,None], c,r]