向量化 numpy.einsum

Question

我有以下四个张量

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，它目前不适用于项目分配！

Answer 1

一种方法是使用 T 的所有维度 -

np.einsum('Hr, Ar, Dr, ATr ->HADT', H, A, D, T)

因为，我们需要在所有输入中对轴-r 求和，同时在输出中保留所有其他（轴），我看不到任何中间方法/引入任何基于点的工具以此来利用 BLAS。