計算numpy.inner（）在第一個（而不是最後一個）軸

Question

我試圖做一個函數，如numpy.inner，但總結了第一個兩個數組的軸而不是最後一個軸。目前我使用tensordot與rollaxis：

def inner1(a, b): 
    return numpy.tensordot(numpy.rollaxis(a, 0, len(a.shape)), b, 1) 

但我想知道：有沒有更好的辦法？也許一個不需要我滾動軸？

我覺得einsum應該使這成爲可能，但我不知道如何在這裏使用它。
當我指定下標字符串時，似乎需要我對a和b進行硬編碼，在此我無法真正做到，因爲對輸入的維度沒有特別的要求。

（注：我上午知道有到求和第一軸，而不是最後的性能影響，但我在這裏忽略它們）

Answer 1

我想你想要的是np.tensordot(a, b, (0, 0))。

Answer 2

這不是爲tensordot液漂亮，但你可以從輸入ndim構建einsum字符串：

ll = 'abcdefghijklmnopqrstuvw' 
astr = ll[0]+ll[1:a.ndim]+','+ll[0]+ll[a.ndim:a.ndim+b.ndim-1] 
np.einsum(astr,a,b) 

np.einsum，您可以指定軸爲列表而不是字符串

np.einsum(a, [0]+range(1,a.ndim), b, [0]+range(a.ndim,a.ndim+b.ndim-1)) 

對於一對3d和2D陣列，這些產生：

np.einsum('abc,ad', a, b) 
np.einsum(a, [0,1,2], b, [0,3]) 

'...'在這裏不起作用，因爲這意味着重複軸（儘可能），因爲你想要唯一的軸（除了第一個軸）。

在編寫messier時，einsum解決方案比tensordot更快（對於小型測試陣列，速度要快3倍）。

---

einsum與另一種選擇是重塑陣列，減少了「剩餘」的尺寸降低到一個。這爲計算增加了一些時間，但並不是很多：

np.einsum('ij,ik',a.reshape(a.shape[0],-1), b.reshape(a.shape[0],-1)).reshape(a.shape[1:]+b.shape[1:])