我可以将 Python 枚举编译成 numba jitclass 吗？

Question

我正在尝试找到如何将 @jitclass 与 Enum class 一起使用。 参考手册明确表示支持它们，但我无法弄清楚，而且我在任何地方都找不到代码示例。

当我尝试跑步时

from numba import jitclass
from numba.types import string

from enum import Enum

type_spec = [
    ('A', string),
    ('B', string)
]


@jitclass(type_spec)
class Type(Enum):
    A = 'A'
    B = 'B'

我收到TypeError: class members are not yet supported: _missing_, name, value, _convert, _member_names_, _member_map_, _member_type_, _value2member_map_, A, B

当我尝试在没有规范的情况下进行编译时，如下所示：

from numba import jitclass

from enum import Enum


@jitclass
class Type(Enum):
    A = 'A'
    B = 'B'

当我调用构造函数时，我得到一个AttributeError: items 。 我似乎无法找到正确的语法将我的 Enum class 转换为 jitclass。 我怎样才能做到这一点？

Answer 1

您可以在njit function 中使用Enum / IntEnum ，而无需在枚举 class 上使用jitclass ：

import numba
from enum import Enum, IntEnum

class MyEnum(Enum):
    Red = 123
    Blue = 456

class IEnum(IntEnum):
    one = 1
    two = 2

@numba.njit
def check_isblue(e):
    assert e == MyEnum.Blue

@numba.njit
def use_cases(a):
    """The following use-cases are supported"""
    # comparison
    a == MyEnum.Red
    a != MyEnum.Red
    a is MyEnum.Red
    a is not MyEnum.Red
    # getitem
    a is MyEnum["Red"]
    # conditional
    MyEnum.Red if True else MyEnum.Blue
    # coersion to int
    IEnum.one + 1


check_isblue(MyEnum.Blue)
use_cases(MyEnum.Red)

其他用例，例如与IntEnum值的矢量化比较（也支持通过@vectorize装饰器）。

这在pull-request 1829中包含的numba/tests/test_enums.py测试模块中“记录”。

Answer 2

我目前面临同样的问题（无法对jitclass枚举进行 jitclass 分类）。 从@user2357112-supports-Monica 似乎暗示的内容来看，Numba 文档似乎可以得到严重改进。

在https://github.com/numba/numba/pull/1829中添加了对枚举的“支持”，但文档方面严重缺乏该拉取请求。