如何使用 numba 进行 function 切换

Question

我有 3 个 jitted 函数， a(x) ， b(x)和c(x) 。

我需要一个开关 function 来执行此操作：

@nb.njit
def switch(i, x):
    if i == 0:
        return a(x)
    elif i == 1:
        return b(x)
    elif i == 2:
        return c(c)

但是我想写的更简洁一些，不会有性能损失，比如：

functions = (a, b, c)

@nb.njit
def switch(i, x):
    functions[i](x)

但是，nopython jit 编译无法处理函数元组。 我怎样才能做到这一点？

Answer 1

只要函数共享相同的类型签名就应该是可能的，请参见以下示例：
https://numba.discourse.group/t/typed-list-of-jitted-functions-in-jitclass/413/4

所以一个例子是：

from numba import njit
from numba.typed import List
from numba.types import float64

a = njit()(lambda x,y: x-y)
b = njit()(lambda x,y: x+y)
c = njit()(lambda x,y: x/y)

function_type = float64(float64, float64).as_type()

funcs = List.empty_list(function_type)
funcs.append(a)
funcs.append(b)
funcs.append(c)

@njit
def switch(i, funcs, *args):
    return funcs[i](*args)

switch(2, funcs, 1, 2)
# 0.5