numpy的einsum：嵌套點產品

Question

我有兩個n -by- k -by- 3陣列a和b，例如，

import numpy as np 

a = np.array([ 
    [ 
     [1, 2, 3], 
     [3, 4, 5] 
    ], 
    [ 
     [4, 2, 4], 
     [1, 4, 5] 
    ] 
    ]) 
b = np.array([ 
    [ 
     [3, 1, 5], 
     [0, 2, 3] 
    ], 
    [ 
     [2, 4, 5], 
     [1, 2, 4] 
    ] 
    ]) 

，並希望計算所有對的點積「三胞胎」，即

np.sum(a*b, axis=2) 

一個更好的辦法來做到這一點也許是einsum，但我似乎無法得到指數。

這裏有什麼提示嗎？

Answer 1

您正在鬆開這兩個輸入數組中的第三個軸，並保持前兩個軸對齊。因此，在np.einsum的情況下，我們可以將前兩個字符串與第三個字符串相同，但是會在輸出字符串符號中跳過，表示兩個輸入都沿着該軸減少。因此，解決方案將是 -

np.einsum('ijk,ijk->ij',a,b)