python - 使用函数对象作为 numba njit 函数的参数
问题描述
我想做一个通用函数,它将函数对象作为参数。
最简单的情况之一:
import numpy as np
import numba as nb
@nb.njit()
def test(a, f=np.median):
return f(a)
test(np.arange(10), np.mean)
给出错误,尽管test(np.arange(10))
按预期工作。
错误:
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
[1] During: typing of argument at <ipython-input-54-52cead0f097d> (5)
File "<ipython-input-54-52cead0f097d>", line 5:
def test(a, f=np.median):
return f(a)
^
This error may have been caused by the following argument(s):
- argument 1: cannot determine Numba type of <class 'function'>
This is not usually a problem with Numba itself but instead often caused by
the use of unsupported features or an issue in resolving types.
这是不允许的还是我错过了什么?
解决方案
使用函数作为参数对于 numba 来说很棘手,而且相当昂贵。常见问题解答中提到了这一点:“1.18.1.1。我可以将函数作为参数传递给 jitted 函数吗?” :
1.18.1.1。我可以将函数作为参数传递给 jited 函数吗?
从 Numba 0.39 开始,只要函数参数也经过 JIT 编译,您就可以:
@jit(nopython=True) def f(g, x): return g(x) + g(-x) result = f(jitted_g_function, 1)
但是,使用作为函数的参数进行分派会产生额外的开销。如果这对您的应用程序很重要,您还可以使用工厂函数来捕获闭包中的函数参数:
def make_f(g): # Note: a new f() is created each time make_f() is called! @jit(nopython=True) def f(x): return g(x) + g(-x) return f f = make_f(jitted_g_function) result = f(1)
提高 Numba 中函数的调度性能是一项持续的任务。
这意味着您可以选择使用函数工厂:
import numpy as np
import numba as nb
def test(a, func=np.median):
@nb.njit
def _test(a):
return func(a)
return _test(a)
>>> test(np.arange(10))
4.5
>>> test(np.arange(10), np.min)
0
>>> test(np.arange(10), np.mean)
4.5
或者在将函数参数作为参数传递之前将其包装为 jited-function:
import numpy as np
import numba as nb
@nb.njit()
def test(a, f=np.median):
return f(a)
@nb.njit
def wrapped_mean(a):
return np.mean(a)
@nb.njit
def wrapped_median(a):
return np.median(a)
>>> test(np.arange(10))
4.5
>>> test(np.arange(10), wrapped_mean)
4.5
>>> test(np.arange(10), wrapped_median)
4.5
这两个选项都有相当多的样板文件,并不像人们希望的那样直截了当。
函数工厂方法还重复创建和编译函数,因此如果您经常使用与参数相同的函数调用它,您可以使用字典来存储已知的编译函数:
import numpy as np
import numba as nb
_precompiled_funcs = {}
def test(a, func=np.median):
if func not in _precompiled_funcs:
@nb.njit
def _test(arr):
return func(arr)
result = _test(a)
_precompiled_funcs[func] = _test
return result
return _precompiled_funcs[func](a)
另一种方法(使用包装和 jitted 函数)也有一些开销,但是只要您传入的数组具有大量元素(> 1000),它就不会真正引人注目。
如果您展示的功能确实是您想要使用的功能,我根本不会使用 numba 。使用 Python + NumPy 执行不锻炼 numba 强度的简单任务(索引和迭代数组或繁重的数字运算)应该更快(或同样快)并且更容易调试和理解:
import numba as nb
def test(a, f=np.median):
return f(a)
推荐阅读
- python - 存储值“|” 使用python在字符串中使用双引号
- javascript - 如何根据另一个组件的事件处理程序更新一个组件?
- asp.net - Asp.net Web 窗体不重定向指定的子文件夹触发 404
- javascript - 如何在 jasmine 中为以下函数编写测试套件,因为它包含匿名函数?
- go - 在 Go 中读取 cookie
- mysql - 我可以在 WampServer 上连接另一个本地 MySQL 服务吗?
- javascript - 编辑字符串数组,只保留值的第一个字母
- php - php 的 session_reset 可以用来重新加载修改过的会话吗?
- matlab - 使用迭代查找方程的根
- python - 使用 QToolBar 和 QListView 进行面包屑导航