numba的有趣行为-使用argmax()的guvectorized函数

时间:2018-11-23 20:50:00

标签: vectorization numba argmax

考虑以下脚本:

from numba import guvectorize, u1, i8
import numpy as np

@guvectorize([(u1[:],i8)], '(n)->()')
def f(x, res):
    res = x.argmax()

x = np.array([1,2,3],dtype=np.uint8)
print(f(x))
print(x.argmax())
print(f(x))

运行它时,我得到以下信息:

4382569440205035030
2
2

为什么会这样?有没有办法解决这个问题?

1 个答案:

答案 0 :(得分:0)

Python没有引用,因此res = ...实际上并未分配给输出参数,而是重新绑定了名称res。我相信res指向的是未初始化的内存,这就是为什么您的首次运行会​​给出看似随机的值。

Numba使用切片语法([:])解决此问题,该语法确实会改变res-您还需要将类型声明为数组。工作功能是:

@guvectorize([(u1[:], i8[:])], '(n)->()')
def f(x, res):
    res[:] = x.argmax()