原文: http://numba.pydata.org/numba-doc/latest/user/generated-jit.html
虽然 jit()
装饰器在许多情况下都很有用,但有时您希望根据其输入类型编写具有不同实现的函数。 generated_jit()
装饰器允许用户在编译时控制特化的选择,同时保持 JIT 函数的运行时执行速度。
假设您要编写一个函数,该函数根据某些约定返回给定值是否为“缺失”值。为了示例,我们采用以下定义:
- 对于浮点参数,缺失值是
NaN
- 对于 Numpy datetime64 和 timedelta64 参数,缺失值为
NaT
- 其他类型没有缺失值的概念。
使用 generated_jit()
装饰器可以轻松实现编译时逻辑:
import numpy as np
from numba import generated_jit, types
@generated_jit(nopython=True)
def is_missing(x):
"""
Return True if the value is missing, False otherwise.
"""
if isinstance(x, types.Float):
return lambda x: np.isnan(x)
elif isinstance(x, (types.NPDatetime, types.NPTimedelta)):
# The corresponding Not-a-Time value
missing = x('NaT')
return lambda x: x == missing
else:
return lambda x: False
这里有几点需要注意:
- 使用参数的 Numba 类型调用修饰函数,而不是它们的值。
- 修饰函数实际上并不计算结果,它返回一个 callable,实现给定类型的函数的实际定义。
- 可以在编译时预先计算一些数据(上面的
missing
变量),以便在编译的实现中重用它们。 - 函数定义使用与装饰函数相同的参数名称,这是确保按名称传递参数按预期工作所必需的。
generated_jit()
装饰器支持与 jit()
装饰器相同的仅关键字参数,例如nopython
和cache
选项。