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

时间:2014-10-26 23:10:36

标签: python arrays numpy linear-algebra inner-product

我试图制作像numpy.inner这样的函数,但它会对两个数组的第一个轴而不是最后一个轴求和。目前我正在tensordot使用rollaxis

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

但我想知道:有更好的方法吗?或许不需要我滚动轴?

我觉得einsum应该可以做到这一点,但我不确定如何在这里使用它。
当我指定下标字符串时,似乎要求我对ab的维度进行硬编码,因为输入维度没有特别要求,所以我无法在此处执行此操作。

(注意:我 意识到在第一个轴而不是最后一个轴上进行求和存在性能影响,但我在这里忽略它们。)

2 个答案:

答案 0 :(得分:2)

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

答案 1 :(得分:1)

这与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:])