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() 装饰器相同的仅关键字参数,例如nopythoncache选项。


我们一直在努力

apachecn/AiLearning

【布客】中文翻译组