在Numba jitclass中调用numpy.round出现TypingError报错

Numba的jitclass对NumPy的原地操作函数（比如numpy.round_）支持有局限：

• round_是原地修改输入数组的方法，jitclass在类型推断时，没法正确处理这类原地操作对类成员数组的内存修改逻辑，直接导致类型匹配失败。
• 相比非原地版本（比如numpy.round），Numba对这类原地ufunc的支持本来就不完善，在jitclass的类方法上下文里，类型检查会更严格，更容易触发报错。

改用非原地的numpy.round()再赋值 #

这是最省事的方案，把原地修改换成返回新数组后重新赋值给类成员：

from numba import jitclass, float64
import numpy as np

spec = [
    ('arr', float64[:])
]

@jitclass(spec)
class TestClass:
    def __init__(self):
        self.arr = np.array([1.234, 5.678])
    
    def process(self):
        self.arr = np.round(self.arr, decimals=2)

obj = TestClass()
obj.process()
print(obj.arr)

自定义原地round函数（必须原地操作时用） #

如果业务逻辑要求必须原地修改数组，可以自己写个简单的原地round函数，用njit装饰后在jitclass方法里调用：

from numba import jitclass, float64, njit
import numpy as np

@njit
def inplace_round(arr, decimals):
    scale = 10 ** decimals
    for i in range(arr.shape[0]):
        arr[i] = np.round(arr[i] * scale) / scale

spec = [
    ('arr', float64[:])
]

@jitclass(spec)
class TestClass:
    def __init__(self):
        self.arr = np.array([1.234, 5.678])
    
    def process(self):
        inplace_round(self.arr, decimals=2)

obj = TestClass()
obj.process()
print(obj.arr)

确认jitclass类型定义精准 #

检查spec里的类型声明是否和实际使用的完全匹配：

• 比如二维数组要写成float64[:, :]，别用泛泛的float64[:]
• 类成员数组的dtype要和spec声明一致，比如别用np.float32数组却声明float64

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