Numba jitclass中初始化键为int二元组的字典报错求助

问题回顾 #

尝试在jitclass中使用键为int类型二元组、值为int类型的字典时，两种初始化方式均报错：

• 直接赋值self.dict = {}时，Numba无法推断空字典的具体类型，报TypingError；
• 在__init__内部调用numba.types.UniTuple构造类型并初始化numba.typed.Dict.empty时，因nopython模式不允许调用类型构造器，同样报错。

正确解决方案 #

核心要点是：在jitclass外部提前定义好字典的类型，然后在类内部使用预定义的类型初始化字典。

完整代码示例：

from numba.experimental import jitclass
import numba

# 提前在Python层面定义好键类型和字典类型
key_type = numba.types.UniTuple(numba.int32, 2)
dict_type = numba.types.DictType(key_type, numba.int32)

# jitclass类型注解使用预定义的类型
@jitclass({
    'shape': numba.types.UniTuple(numba.int32, 2),
    'dict': dict_type
})
class BigramCounts:
    def __init__(self, shape):
        self.shape = shape
        # 使用提前定义好的key_type初始化字典
        self.dict = numba.typed.Dict.empty(key_type=key_type, value_type=numba.int32)

# 测试代码
b_c = BigramCounts((2, 3))
b_c.dict[(1, 2)] = 5
print(b_c.dict[(1, 2)])  # 输出：5

错误原因解析 #

• 直接赋值空字典{}：Numba无法将无类型的空字典与jitclass中声明的DictType[UniTuple(int32,2), int32]类型匹配，必须显式初始化指定类型的字典。
• 在__init__内部构造类型：numba.types.UniTuple是Python层面的类型构造工具，jitclass的__init__在nopython模式下运行时，无法执行这类Python代码，必须提前在外部定义好类型并传入。

内容的提问来源于stack exchange，提问作者Adam