python - 为什么依赖 numba jitt'ed 函数的顺序很重要？

Question

在 python 中，您可以定义多个以任意顺序相互调用的函数，并且在运行时将调用这些函数。这些函数在脚本中定义的顺序并不重要，只要它们存在即可。例如，下面是有效的并且可以工作

import numpy as np

def func1(arr):
    out = np.empty_like(arr)
    for i in range(arr.shape[0]):
        out[i] = func2(arr[i])  # calling func2 here which is defined below
    return out

def func2(a):
    out = a + 1
    return out

func1 可以调用 func2，即使 func2 是在 func1 之后定义的。

但是，如果我用 numba 修饰这些函数, 我得到一个错误

import numpy as np
import numba as nb


@nb.jit("f8[:](f8[:])", nopython=True)
def func1(arr):
    out = np.empty_like(arr)
    for i in range(arr.shape[0]):
        out[i] = func2(arr[i])
    return out

@nb.jit("f8(f8)", nopython=True)
def func2(a):
    out = a + 1
    return out

>>> TypingError: Failed in nopython mode pipeline (step: nopython frontend)
    Untyped global name 'func2': cannot determine Numba type of <class 
    'numba.ir.UndefinedType'>

所以numba在使用JIT编译func1时并不知道func2是什么。不过，只需切换这些函数的顺序即可，因此 func2 位于 func1

之前

@nb.jit("f8(f8)", nopython=True)
def func2(a):
    out = a + 1
    return out

@nb.jit("f8[:](f8[:])", nopython=True)
def func1(arr):
    out = np.empty_like(arr)
    for i in range(arr.shape[0]):
        out[i] = func2(arr[i])
    return out

这是为什么？我有一种感觉，纯 python 模式可以工作，因为 python 是动态类型的而不是编译的，而 numba，使用 JIT，根据定义编译函数(因此可能需要完全了解每个函数中发生的一切？)。但是我不明白为什么 numba 在遇到它没有见过的函数时不在范围内搜索所有函数。

Answer 1

精简版 - 删除 "f8[:](f8[:])"

您的直觉是正确的。 Python 函数在调用 时被查找，这就是为什么它们可以乱序定义的原因。查看带有 dis(反汇编)模块的 python 字节码可以清楚地看到名称 b 每次都被查找为全局函数 a被称为。

def a():
    return b()

def b():
    return 2

import dis
dis.dis(a)
#  2           0 LOAD_GLOBAL              0 (b)
#              2 CALL_FUNCTION            0
#              4 RETURN_VALUE

在 nopython 模式下，numba 需要静态知道每个被调用函数的地址——这使得代码速度更快(不再进行运行时查找)，并且也为其他函数打开了大门优化，如内联。

也就是说，numba 可以处理这种情况。通过指定类型签名 ("f8[:](f8[:])")，您可以强制提前编译。省略它，一个数字将推迟到第一个调用它的函数，它会起作用。