矢量化numpy.einsum

Question

我有四个张量

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

对于每个iin 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。