python - 努巴。如何使用元组轴参数编写 np.sum ？

Question

此代码:

@jit(nopython=True)
def foo(x):
    return x.sum(axis=(1,2))

x=np.linspace(0,1)
x=x.reshape(5,5,-1)
print(foo(x))

返回此错误:

NotImplementedError: No definition for lowering array.sum(array(float64, 3d, C), Tuple(Literal[int](1), Literal[int](2))) -> array(float64, 2d, C)

当 axis 参数只是一个整数，而不是整数元组( https://numpy.org/doc/stable/reference/generated/numpy.ndarray.sum.html#numpy.ndarray.sum )时，numba 似乎支持 np.sum 。 因此，我使用此解决方法 return x.sum(axis=1).sum(axis=1) 但如果我考虑代码优化，这不是一个好的解决方案。

是否存在其他解决方案或者我必须等待 future 的 numba 版本？

Answer 1

只需将数组重新调整为比所需结果多一维即可。在您的示例中:

x.sum(axis=(1,2))

可以替换为:

x.reshape(5,-1).sum(axis=1)

它产生相同的结果并且可以由 Numba 执行。