动态编译numba jit

Question

我有许多功能来支持单线程和多线程执行，这使得我需要双行代码，从而产生冗余的重复：

@jit(nopython=True, cache=True)
def fn1(x):
    for i in range(len(x)):
        # [code body]

@jit(nopython=True, cache=True, parallel=True)
def fn1_par(x):
    for i in prange(len(x)):
        # [code body]

两者的代码体完全相同，唯一不同的是@jit(parallel=True)和range -> prange。是否有一种方法可以在一个函数中表达这个逻辑，同时保留cache=True的优点？

注意:使用“包装器包装器”(osvil的答案)的解决方法在cache=True中失败；第一个调用的信任持续存在(例如，第一个函数的parallel=False将覆盖第二个的parallel=True)。更改包装函数的__name__没有帮助。打开问题。

Answer 1

numba.jit只是个装潢师。您可以定义函数一次，然后在同一函数上多次应用装饰符“手动”。

从我的脑海里看上去会是这样的：

def python_function(x):
   ...

numba_nopython_function = numba.jit(nopython=True)(python_function)
numba_nopython_parallel_function = numba.jit(nopython=True, parallel=True)(python_function)

通过numba_nopython_function使用该函数将使用普通的numba版本。使用numba_nopython_parallel_function将使用numba并行函数。如果您愿意，还可以通过python_function本身访问原始python函数。它符合干燥原则。按你的意愿命名，因为我的名字太冗长了，太不切实际了。

不过，我不知道这将如何与函数的缓存进行交互。