矢量化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 。