Numba njit编译后运行numpy矩阵求逆结果与纯Python不一致如何解决

问题根源 #

Numba实现的np.linalg.inv和NumPy原生接口的底层依赖库不同：

• NumPy默认调用系统绑定的成熟LAPACK后端（如MKL、OpenBLAS），这类库针对硬件做了大量精度适配和优化，计算精度更高
• Numba内置的线性代数函数基于简化的LAPACK端口实现，默认开启的快速数学优化也会牺牲部分精度，当输入矩阵是病态矩阵时，误差会被明显放大

解决方案 #

不需要自行编写Numba包装函数，可按优先级选择以下方案：

1. 开启Numba精度对齐参数
   给njit装饰器添加精度相关配置，关闭有损优化，对齐NumPy的计算逻辑：

@numba.njit(error_model='numpy', fastmath=False)
def cal_Test_jit(A,b):
    c = np.linalg.inv(A)@b
    return c, np.linalg.inv(A)

该方案无需修改业务逻辑，大部分场景下可大幅缩小两个版本的计算偏差。

2. 用objmode上下文调用原生NumPy接口
   如果调整参数后精度仍不满足要求，可直接在Numba函数中指定求逆逻辑走原生Python实现，完全保留NumPy的计算精度，仅该片段不享受Numba加速：

from numba import njit, objmode

@njit
def cal_Test_jit(A,b):
    # 括号内指定返回值的类型和维度，和你实际数据类型匹配即可
    with objmode(Ai='float64[:,:]', c='float64[:]'):
        Ai = np.linalg.inv(A)
        c = Ai @ b
    return c, Ai

如果你的函数整体耗时大头不是矩阵求逆，该方案对整体性能的影响几乎可以忽略。

额外优化建议 #

显式求逆后乘向量的写法本身数值稳定性较差，无论是原生Python还是Numba场景，都建议替换为直接求解线性方程组的np.linalg.solve(A, b)，不需要显式计算逆矩阵，数值误差会小很多，Numba对solve接口的实现精度也远高于inv，修改后两个版本的误差基本可以控制在浮点精度误差范围内：

# 原生版本
def cal_Test(A,b):
    c = np.linalg.solve(A, b)
    Ai = np.linalg.inv(A)
    return c, Ai

# Numba版本
@numba.njit(error_model='numpy', fastmath=False)
def cal_Test_jit(A,b):
    c = np.linalg.solve(A, b)
    Ai = np.linalg.inv(A)
    return c, Ai

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