如何为用作查找表的Numba字典正确定义类型？

你的代码存在两个核心问题：

1. 函数内变量名错误：python_dict应为传入的pyth_dict
2. Numba作为静态类型JIT编译器，要求字典必须显式声明键值类型，尤其当值为异构类型（整数+二维数组）时，无法自动推断兼容类型

结合你"只读查找表、键值类型固定"的需求，提供两种可行方案：

方案1：在Numba函数内显式声明字典类型 #

使用numba.typed.Dict.empty()创建指定类型的字典，键类型设为字符串，值类型用numba.types.pyobject兼容整数和数组：

import numba
import numpy as np
mat = np.random.normal(size=(3,3))

pyth_dict = {"A": 2, "B": mat}

@numba.njit()
def numbaize(pyth_dict):
    # 明确指定键为字符串类型，值为object类型（兼容整数和数组）
    numba_dict = numba.typed.Dict.empty(
        key_type=numba.types.unicode_type,
        value_type=numba.types.pyobject
    )
    for key in pyth_dict.keys():
        numba_dict[key] = pyth_dict[key]  # 修正变量名错误
    return numba_dict

numba_dict = numbaize(pyth_dict)

方案2：函数外初始化Numba字典（性能更优） #

因为字典是只读结构，可直接在JIT函数外完成初始化，避免函数内的类型推断和循环开销：

import numba
import numpy as np
mat = np.random.normal(size=(3,3))

pyth_dict = {"A": 2, "B": mat}

# 外部创建并填充Numba字典
numba_dict = numba.typed.Dict.empty(
    key_type=numba.types.unicode_type,
    value_type=numba.types.pyobject
)
for k, v in pyth_dict.items():
    numba_dict[k] = v

# 直接在Numba函数中使用该字典
@numba.njit()
def lookup(numba_dict, key):
    return numba_dict[key]

# 测试调用
print(lookup(numba_dict, "A"))
print(lookup(numba_dict, "B"))

关键注意事项 #

• 异构值必须用numba.types.pyobject作为值类型，这是Numba静态类型系统兼容不同Python对象的唯一方式
• 方案2更适合只读场景，能减少JIT函数内的额外开销
• 确保所有键为字符串类型，值的类型与预先定义的一致，否则会触发类型不匹配错误

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