在 numba 函数中将 np.min 与列表输入一起使用

Question

这里使用np.min有什么问题？为什么 numba 不喜欢在该函数中使用列表，有没有其他方法可以让 np.min 工作？

from numba import njit
import numpy as np

@njit
def availarray(length):
    out=np.ones(14)
    if length>0:
        out[0:np.min([int(length),14])]=0
    return out

availarray(3)

该函数在 min 上运行良好，但 np.min 应该更快...

Answer 1

为了使您的代码与 numba 一起工作，您必须在 NumPy 数组上应用 np.min，这意味着您必须将列表 [int(length),14] 转换为 NumPy 数组，如下所示正在关注

from numba import njit
import numpy as np

@njit
def availarray(length):
    out=np.ones(14)
    if length>0:
        out[0:np.min(np.array([int(length),14]))]=0   
    return out

availarray(3)
# array([0., 0., 0., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

Answer 2

问题是 np.min 的 numba 版本需要 array 作为输入。

from numba import njit
import numpy as np

@njit
def test_numba_version_of_numpy_min(inp):
    return np.min(inp)

>>> test_numba_version_of_numpy_min(np.array([1, 2]))  # works
1

>>> test_numba_version_of_numpy_min([1, 2]) # doesn't work
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<function amin at 0x000001B5DBDEE598>) with argument(s) of type(s): (reflected list(int64))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.

更好的解决方案是只使用 Python 的 numba 版本 min:

from numba import njit
import numpy as np

@njit
def availarray(length):
    out = np.ones(14)
    if length > 0:
        out[0:min(length, 14)] = 0
    return out

因为 np.min 和 min 实际上都是这些函数的 Numba 版本（至少在 njitted 函数中）min 在这种情况下也应该快得多.然而，它不太可能被注意到，因为数组的分配和将一些元素设置为零将是这里的主要运行时贡献者。

请注意，您甚至不需要在此处调用 min - 因为即使使用更大的停止索引，切片也会隐式停止在数组的末尾：

from numba import njit
import numpy as np

@njit
def availarray(length):
    out = np.ones(14)
    if length > 0:
        out[0:length] = 0
    return out

在 numba 函数中将 np.min 与列表输入一起使用

Using np.min with list input in a numba function

python

numpy

list

numba