python - 如何在 python 中使用 numba.jit 将计算值传递给列表排序？

Question

我正在尝试使用 Python 中 numba-jit 函数中的自定义键对列表进行排序。简单的自定义键工作，例如我知道我可以使用这样的东西按绝对值排序:

import numba

@numba.jit(nopython=True)
def myfunc():
    mylist = [-4, 6, 2, 0, -1]
    mylist.sort(key=lambda x: abs(x))
    return mylist  # [0, -1, 2, -4, 6]

但是，在以下更复杂的示例中，我收到了一个我不明白的错误。

import numba
import numpy as np


@numba.jit(nopython=True)
def dist_from_mean(val, mu):
    return abs(val - mu)

@numba.jit(nopython=True)
def func():
    l = [1,7,3,9,10,-4,-2,0]
    avg_val = np.array(l).mean()
    l.sort(key=lambda x: dist_from_mean(x, mu=avg_val))
    return l

它报告的错误如下:

Traceback (most recent call last):
  File "testitout.py", line 18, in <module>
    ret = func()
  File "/.../python3.6/site-packages/numba/core/dispatcher.py", line 415, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/.../python3.6/site-packages/numba/core/dispatcher.py", line 358, in error_rewrite
    reraise(type(e), e, None)
  File "/.../python3.6/site-packages/numba/core/utils.py", line 80, in reraise
    raise value.with_traceback(tb)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: convert make_function into JIT functions)
Cannot capture the non-constant value associated with variable 'avg_val' in a function that will escape.

File "testitout.py", line 14:
def func():
    <source elided>
    l.sort(key=lambda x: dist_from_mean(x, mu=avg_val))
                                                ^

你知道这里发生了什么吗？

Answer 1

Do you know what is happening here?

通过使用参数 nopython = True您停用了对象模式，因此 Numba 无法将所有值作为 Python 对象处理(请参阅: https://numba.pydata.org/numba-doc/latest/glossary.html#term-object-mode)。 (引用其实是我今天偶然写的另一个帖子: How call a `@guvectorize` inside a `@guvectorize` in numba?)

@numba.jit(nopython=True)
def func():
    l = [1,7,3,9,10,-4,-2,0]
    avg_val = np.array(l).mean()
    l.sort(key=lambda x: dist_from_mean(x, mu=avg_val))
    return l

无论如何， lambda对于 numba jit 函数来说“太”复杂了——至少当它作为参数传递时(比较 https://github.com/numba/numba/issues/4481)。与 nopython激活模式后，您只能使用有限数量的库 - 完整列表可在此处找到: https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
这就是它抛出以下错误的原因:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step:convert make_function into JIT functions) Cannot capture thenon-constant value associated with variable 'avg_val' in a functionthat will escape.

此外，您在另一个中引用了 jit 加速函数 - 当具有 nopython = True 时.这也可能是问题的根源。
我强烈建议您查看以下教程: http://numba.pydata.org/numba-doc/latest/user/5minguide.html#will-numba-work-for-my-code;它应该可以帮助您解决类似的问题!

进一步阅读和来源:

https://github.com/numba/numba/issues/5120

http://numba.pydata.org/numba-doc/latest/user/5minguide.html#will-numba-work-for-my-code

TypingError: Failed in nopython mode pipeline (step: nopython frontend)

https://github.com/numba/numba/issues/4481

https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html

How call a `@guvectorize` inside a `@guvectorize` in numba?