Skip to content

1.5。使用@generated_jit 进行灵活的专业化

原文： http://numba.pydata.org/numba-doc/latest/user/generated-jit.html

虽然 jit() 装饰器在许多情况下都很有用，但有时您希望根据其输入类型编写具有不同实现的函数。 generated_jit() 装饰器允许用户在编译时控制特化的选择，同时保持 JIT 函数的运行时执行速度。

1.5.1。示例

假设您要编写一个函数，该函数根据某些约定返回给定值是否为“缺失”值。为了示例，我们采用以下定义：

  • 对于浮点参数，缺失值是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变量），以便在编译的实现中重用它们。
  • 函数定义使用与装饰函数相同的参数名称，这是确保按名称传递参数按预期工作所必需的。

1.5.2。编译选项

generated_jit() 装饰器支持与 jit() 装饰器相同的仅关键字参数，例如nopython和cache选项。

回到顶部