矢量化 numpy.einsum

Vectorising numpy.einsum

我有以下四个张量

  1. H (h, r)
  2. A(a,r)
  3. D (d, r)
  4. T (a, t, r)

对于a中的每个i,都有一个对应的T[i]形状(t, r)

我需要执行 np.einsum 以产生以下结果 (pred):

pred = np.einsum('hr, ar, dr, tr ->hadt', H, A, D, T[0])
for i in range(a):
    pred[:, i:i+1, :, :] = np.einsum('hr, ar, dr, tr ->HADT', H, A[i:i+1], D, T[i])

但是,我想在不使用 for 循环的情况下进行此计算。原因是我使用的 autograd 当前不适用于项目分配!

一种方法是使用 T -

的所有维度
np.einsum('Hr, Ar, Dr, ATr ->HADT', H, A, D, T)

因为,我们需要对所有输入进行求和归约轴r,同时在输出中保留所有其他(轴),我看不到任何中间方法 it/bringing在任何基于点的工具中利用 BLAS。