[英]Compute numpy.inner() over first (instead of last) axis
I'm trying to make a function like numpy.inner
, but which sums over the first axis of both arrays instead of the last axis. 我正在尝试制作一个像
numpy.inner
这样的函数,但是该函数求和两个数组的第一个轴,而不是最后一个轴。 Currently I'm using tensordot
with rollaxis
: 目前,我正在使用
tensordot
与rollaxis
:
def inner1(a, b):
return numpy.tensordot(numpy.rollaxis(a, 0, len(a.shape)), b, 1)
but I'm wondering: is there a better way? 但我想知道:还有更好的方法吗? Perhaps one that doesn't require me to roll the axes?
也许不需要我滚动轴的那一个?
I feel like einsum
should make this possible, but I'm not sure how to use it here. 我觉得
einsum
应该可以做到这一点,但是我不确定如何在这里使用它。
It seems to require me to hard-code the dimensionality of a
and b
when I specify the subscripts string, which I can't really do here because there is no particular requirement on the input dimensionality. 当我指定下标字符串时,似乎要求我对
a
和b
的维进行硬编码,因为在输入维上没有特别要求,所以我在这里实际上不能这样做。
(Note: I am aware that there are performance implications to summing over the first axis instead of the last, but I'm ignoring them here.) (注: 我知道有到求和第一轴,而不是最后的性能影响,但我在这里忽略它们。)
我认为您想要的是np.tensordot(a, b, (0, 0))
。
This isn't as pretty as the tensordot
solution, but you can construct the einsum
string from ndim
of the inputs: 这不是因为漂亮
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
lets you specify axes as lists rather than the string np.einsum
允许您将轴指定为列表而不是字符串
np.einsum(a, [0]+range(1,a.ndim), b, [0]+range(a.ndim,a.ndim+b.ndim-1))
For a pair of 3d and 2d arrays, these produce: 对于一对3d和2d阵列,它们产生:
np.einsum('abc,ad', a, b)
np.einsum(a, [0,1,2], b, [0,3])
'...'
doesn't work here because that implies repeated axes (to the extent possible), where as you want unique axes (except for the 1st). '...'
在这里不起作用,因为这意味着重复的轴(在可能的范围内),在此您需要唯一的轴(第一个轴除外)。
While messier to write, the einsum
solution is faster than the tensordot
one (3x faster for small test arrays). 而混乱写,所述
einsum
溶液比快tensordot
一个(3×更快小测试阵列)。
Another option with einsum
is to reshape the arrays, reducing the 'remaining' dimensions down to one. einsum
另一种选择是对阵列进行einsum
,将“剩余”尺寸减小到一个。 This adds a bit of time to the calculation, but not a lot: 这会给计算增加一些时间,但不会很多:
np.einsum('ij,ik',a.reshape(a.shape[0],-1), b.reshape(a.shape[0],-1)).reshape(a.shape[1:]+b.shape[1:])
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.