矢量化numpy.einsum

Question

我有以下四个张量

1. H（h，r）
2. A（a，r）
3. D（d，r）
4. T（a，t，r）

对于i中的每个a，形状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。