作为参数的jitted函数的签名

Question

我查阅了numba文档，但什么也没找到

我有一个函数to jit，它接受一个jitted_函数作为参数。我想通过添加签名进行快速编译，就像：

@jit(float64('jit_func.type', int32, int32...))

“jitted_func.type”应为“函数类型”

当我这样做时：

type(jitted_func)

我得到一个CPUDispatcher对象

谢谢你的帮助

Answer 1

我也在寻找解决办法。不幸的是，@Carbon的建议不起作用，因为函数bar的numba.typeof返回的类型与函数baz的类型不同，即使bar和baz的签名相同

例如：

import numba 

@numba.jit(
    numba.int32(numba.int32),
    nopython=True,
    nogil=True,
)
def bar(a):

    return 2 * a

@numba.jit(
    numba.int32(numba.int32),
    nopython=True,
    nogil=True,
)
def baz(a):

    return 3 * a

@numba.jit(
    numba.int32(numba.typeof(bar), numba.int32),
    nopython=True,
    nogil=True,
)
def foo(fn, a):

    return fn(a)

foo(bar, 2)返回4

foo(baz, 2)返回以下异常：

Traceback (most recent call last):
  File "test_numba.py", line 33, in <module>
    print(foo(baz, 2))
  File "<snip>\Python38\lib\site-packages\numba\core\dispatcher.py", line 656, in _explain_matching_error
    raise TypeError(msg)
TypeError: No matching definition for argument type(s) type(CPUDispatcher(<function baz at 0x000001DFA8C2D1F0>)), int64

我发现的唯一解决办法是完全省略foo的函数签名，让numba来解决这个问题。我不知道这会给你的代码带来什么负面影响（如果有的话）

例如：

import numba 

@numba.jit(
    numba.int32(numba.int32),
    nopython=True,
    nogil=True,
)
def bar(a):

    return 2 * a

@numba.jit(
    numba.int32(numba.int32),
    nopython=True,
    nogil=True,
)
def baz(a):

    return 3 * a

@numba.jit(
    nopython=True,
    nogil=True,
)
def foo(fn, a):

    return fn(a)

foo(bar, 2)返回4

foo(baz, 2)返回6

Answer 2

所以，我不知道如何生成你要寻找的签名，但是如果你有一个带有你要签名的编译函数的样本，你可以使用^ {CD1>}得到预期的签名，考虑，例如：

import numba

@numba.njit(numba.int32(numba.int32))
def x(a):
    return a+1

@numba.njit(numba.int32(numba.typeof(x), numba.int32))
def y(fn,a):
    return fn(a)
    
print(y(x,3))

我查过了，这真是太好了。如果您想进一步处理这个问题，正确的开始位置是numba.core.types.functions，并且Dispatcher类型是在编译中专门处理的，请参见numba.core.typing.context.BaseContext的._resolve_user_function_type

Answer 3

我也对这一点感到好奇，因为我也想对runge-kutta四阶算法进行急切的编译。我还认真地搜索了文档，不幸的是，这似乎还没有实现：（因此，我们现在唯一要做的就是让它进行惰性编译，并在第一次函数调用时体验编译