将 function 及其所有子函数传递给 njit

Question

所以我最近发现了 Numba，我对此感到非常惊讶。 When trying it out I've used a bubblesort function as the test function, but since my bubblesort function calls another function I get errors when calling njit on it.

我已经解决了这个问题，首先在我的bubblesort子函数上调用njit，然后让我的bubblesort调用njit子函数，它可以工作，但它迫使我在尝试比较时定义两个bubblesort函数。 我想知道是否有另一种方法可以做到这一点。

这就是我正在做的事情：

def bytaintill(l):
    changed = False
    for i in range(len(l) - 1):
        if l[i] > l[i+1]:
            l[i], l[i+1] = l[i+1], l[i]
            changed = True
    return changed


bytaintill_njit = njit()(bytaintill)

def bubblesort(l):
    not_done = True
    while not_done:
        not_done = bytaintill_njit(l)
    return

def bubble(l):
    not_done = True
    while not_done:
        not_done = bytaintill(l)
    return

bubblesort_njit = njit()(bubblesort)

Answer 1

要扩展我的评论，您不需要定义新功能，但也可以将 jit-ed 版本 map 改成同名。 通常，最方便的方法是使用@jit装饰器（或@njit ，它是@jit(nopython=True)的缩写）。

from numba import njit

@njit
def bytaintill(l):
    changed = False
    for i in range(len(l) - 1):
        if l[i] > l[i+1]:
            l[i], l[i+1] = l[i+1], l[i]
            changed = True
    return changed

@njit
def bubble(l):
    not_done = True
    while not_done:
        not_done = bytaintill(l)
    return

出于基准测试的目的，您可以简单地注释掉装饰器。 如果您希望能够在 jit-ed 和 python 版本之间来回切换 go，您可以尝试这样的操作：

from numba import njit

do_jit = True  # set to True or False

def bytaintill(l):
    changed = False
    for i in range(len(l) - 1):
        if l[i] > l[i+1]:
            l[i], l[i+1] = l[i+1], l[i]
            changed = True
    return changed

def bubble(l):
    not_done = True
    while not_done:
        not_done = bytaintill(l)
    return

if do_jit:
    bytaintill = njit()(bytaintill)
    bubble = njit()(bubble)