python - 在第一个（而不是最后一个）轴上计算 numpy.inner()

Question

我正在尝试制作一个类似的函数，但它在两个数组的第一个轴而不是最后一个轴上numpy.inner求和。目前我正在使用：tensordotrollaxis

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

但我想知道：有没有更好的方法？也许不需要我滚动轴？

我觉得einsum应该使这成为可能，但我不确定如何在这里使用它。
似乎需要我硬编码的维度a以及b当我指定下标字符串时，我在这里无法真正做到这一点，因为对输入维度没有特殊要求。

（注意：我知道在第一个轴而不是最后一个轴上求和会影响性能，但我在这里忽略了它们。）

score 2 · Accepted Answer

2

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

于 2014-10-26T23:15:51.017 回答

score 1 · Accepted Answer

这不如tensordot解决方案漂亮，但您可以从输入构造einsum字符串ndim：

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])

'...'在这里不起作用，因为这意味着重复的轴（在可能的范围内），而您想要唯一的轴（第一个轴除外）。

虽然编写起来比较麻烦，但该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:])

2 回答 2