尝试 numba 的包装器

Question

我编写此包装器是否使用 numba 取决于可用性（“可以导入 numba 吗？”）和用户设置（“numba_enabled() 是否返回 True？”）。

这对于测试特别有用，我想在有和没有 numba 的情况下运行测试。

def try_numba_autojit(func):
    def func_(*args):
        if numba_enabled():
            try:
                import numba
            except:
                pass
            else:
                return numba.autojit(func)(*args)
        return func(*args)
    return func_

问题：包装器导致原本可以工作的函数抛出此错误：

*** OverflowError: can't convert negative value to unsigned char

如果我设置try_numba_autojit = numba.autojit，覆盖包装器，则没有错误。我的结论是 *args 可能正在做一些 numba 不喜欢的事情。

在 pdb 中，我尝试过手动解包 args：

(Pdb) numba.autojit(func)
<specializing numba function(<function _refine at 0xa53ec6c>)>
(Pdb) numba.autojit(func)(*args)
*** OverflowError: can't convert negative value to unsigned char

(Pdb) numba.autojit(func)(args[0], args[1], args[2], args[3], args[4], args[5], args[6])
*** OverflowError: can't convert negative value to unsigned char

同样的错误。

是否有推荐定义一个类似try_numba_autojit 的函数以避免打扰 numba？

Answer 1

我不知道为什么您正在做的事情不起作用，但是如果您可以在包装时间而不是函数调用时间测试 numba-availibility/enability，这可能会起作用（未经测试，抱歉)

def try_numba_autojit():
    if numba_enabled():
        try:
            import numba
            return numba.autojit
        except ImportError:
            pass
    return lambda func: func

并像这样使用：

@try_numba_autojit()
def f(a,b,c): ...