计算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:])