具有同质签名函数的 Numba 类型列表

Question

我有一些具有相同签名的 numba-jitted 事件函数，即：

from numba import jit

@jit("Tuple((float64,float64))(float64[::1])", nopython=True)
def event_1(y): 
    return 1.1, 1.2 # (random values for this example)

@jit("Tuple((float64,float64))(float64[::1])", nopython=True)
def event_2(y): 
    return 2.1, 2.2 # (random values for this example)

我的目标是创建一个 jitted-function，其中 returns 一个事件函数列表。 event_handler 输出应该是另一个缓存编译的 jitted 函数的输入，因此它的签名必须是固定的：

from numba.typed import List
@jit("ListType(FunctionType(Tuple((float64, float64))(float64[::1])))()", nopython=True)
def event_handler(): 
    return List([event_1, event_2])

但是，只有当 event_handler 返回的列表至少有两个不同的事件函数时，上面的代码才有效。如果列表只有 1 个事件函数项，或同一函数的多个项（例如，List([event_1]) 或 List([event_1, event_1])），则上面的代码无法编译并产生以下错误：

No conversion from ListType[type(CPUDispatcher(<function event_1 at 0x7f83c2a22430>))] to ListType[FunctionType[UniTuple(float64 x 2)(array(float64, 1d, C))]]

我认为原因是因为在后一种情况下，列表项类型设置为等于函数 CPUDispatcher 而不是类型化函数签名。我已经尝试 this 解决方案来初始化列表，但不幸的是它不起作用。

我该如何解决这个问题？我觉得奇怪的是，当提供两个不同的函数时，列表类型被正确推断，但每当给出单个项目时突然继承 CPUDispatcher 类型。

Answer 1

我找到了一个解决方案，它改进了列表参数的显式类型和 typed list 的 empty_list 方法。参考我的原始 post，我可以使用 empty_list 方法初始化列表，将其签名作为参数传递（在任何 jitted 函数的范围之外定义）：

from numba import types
from numba.typed import List

vector_sig = types.Array(dtype=types.float64, ndim=1, layout="C")
function_sig = types.FunctionType(types.Tuple((types.float64, types.float64))(vector_sig))

@jit("ListType(FunctionType(Tuple((float64, float64))(float64[::1])))()", nopython=True)
def event_handler(): 
    event_list = List.empty_list(function_sig)
    event_list.append(event_1)
    return event_list

具有同质签名函数的 Numba 类型列表

Numba typed lists with homogeneous-signature functions

python

list

numba