矢量化numpy.einsum

Question

我有四个张量

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

对于i中a每个i ，存在形状(t, r)的对应T[i] 。

我需要做一个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 。