Numpy 计算一个张量与另一个张量的所有滚动之间的成对矩阵积的有效方法

假设我们有两个张量：

形状为（d，m，n）的张量A

形状为（d，n，l）的张量B

如果我们想得到A和B最右边矩阵的成对矩阵积，我想我们可以使用大小为（d，d，m，l）的np.einsum（'dmn，…nl->d…ml'，A，B）。然而，我想得到的不是所有对的成对乘积

导入参数k，1d…ml'，a，B），然后使用掩码提取（d，k）对

首先平铺B，然后以某种方式使用einsum

---

但我认为还有更好的办法

我怀疑你能比for循环做得更好。例如，与双for循环相比，使用einsum和stride_技巧的矢量化版本如下：

代码：

from simple_benchmark import BenchmarkBuilder, MultiArgument
import numpy as np
from numpy.lib.stride_tricks import as_strided
B = BenchmarkBuilder()

@B.add_function()
def loopy(A,B,k): 
    d,m,n = A.shape                                   
    l = B.shape[-1]                     
    out = np.empty((d,k,m,l),int)                      
    for i in range(d):                         
        for j in range(k):                     
            out[i,j] = A[i]@B[(i+j)%d]                      
    return out                     

@B.add_function()
def vectory(A,B,k):                                            
    d,m,n = A.shape                                            
    l = B.shape[-1]                                            
    BB = np.concatenate([B,B[:k-1]],0)                         
    BB = as_strided(BB,(d,k,n,l),np.repeat(BB.strides,(2,1,1)))
    return np.einsum("ikl,ijln->ijkn",A,BB)                    

@B.add_arguments('d x k x m x n x l')
def argument_provider():
    for exp in range(10):
        d,k,m,n,l = (np.r_[1.6,1.5,1.5,1.5,1.5]**exp*(4,2,2,2,2)).astype(int)
        print(d,k,m,n,l)
        A = np.random.randint(0,10,(d,m,n))                            
        B = np.random.randint(0,10,(d,n,l))
        yield k*d*m*n*l,MultiArgument([A,B,k])

r = B.run()
r.plot()

import pylab
pylab.savefig('diagwa.png')