假設我有一個大小為 (M,N,D,D) 的陣列 A 和一個大小為 (M,N,D) 的向量 x。我想生成大小為 (M,N,D) 的陣列 b，其中 b[m,n,:] = np.matmul(A[m,n,:,:],x[m,n,:]) . 所以基本上，包含在 A[m,n,:] 中的 D×D 矩陣和包含在 x[m,n,:] 中的大小為 D 的向量應該按照標準矩陣向量乘法相乘。一個天真的 for 回圈實作將是
for m in range(M):
for n in range(N):
b[m,n,:] = np.matmul(A[m,n,:],x[m,n,:])
顯然，這對于大 N 和 M 是低效的（我就是這種情況；N 的數量級為 100,000，M 的數量級為 1000）。我的問題是：如何矢量化這個回圈？我試過使用 np.tensordot 但無論我如何嘗試構造它，我都會遇到形狀不匹配，大概是因為 A 有 4 個軸而 x 有 3 個（請注意，就我所知，np.tensordot是正確的使用函式，我只是沒有正確使用它）
編輯
嘗試使用 np.einsum，但我沒有得到預期的結果（應該是 [5,11]：
uj5u.com熱心網友回復：
這就是np.einsum亮點：
np.einsum('mnij,mnj->mni', A,B)
檢查
np.allclose(np.einsum('mnij,mnj->mni', A,x), b)
回傳 True
uj5u.com熱心網友回復：
廣播可以達到同樣的結果：
x1 = x.reshape(M,N,d,1)
b1 = np.matmul(A,x1)
b2 = b1.reshape(M,N,d)
轉載請註明出處，本文鏈接：https://www.uj5u.com/net/338746.html