如何解决嵌套for循环场景下的numba lowering错误

报错原因 #

Numba的@njit装饰器要求函数内所有变量的类型、长度必须在编译阶段固定，不能在运行过程中动态变更。你的代码里：

• 初始时res = [0]，是长度为1的整型列表
• 后续赋值res = [x, y]，变成了长度为2的浮点型列表
  类型和长度都发生了改变，超出了Numba静态类型推断的支持范围，因此触发lowering报错。

解决方案 #

方案1：用独立变量存储最优坐标（最简单高效） #

不需要用列表存结果，直接声明两个单独的浮点变量分别存最优的x、y值，完全避免列表类型变更的问题：

import numpy as np
from numba import njit

@njit()
def test():
    xs = np.arange(0, 1, 0.1)
    ys = np.arange(0, 1, 0.1)

    best_x = 0.0
    best_y = 0.0
    smallest_value = np.inf
    for x in xs:
        for y in ys:
            value = x*x + y*y
            if value < smallest_value:
                smallest_value = value
                best_x = x
                best_y = y
    return [best_x, best_y], smallest_value

if __name__ == "__main__":
    print(test())

方案2：初始化时就固定res的类型和长度 #

一开始就把res初始化为长度为2的浮点数组，后续只修改元素值，不重新赋值整个列表：

import numpy as np
from numba import njit

@njit()
def test():
    xs = np.arange(0, 1, 0.1)
    ys = np.arange(0, 1, 0.1)

    res = np.zeros(2, dtype=np.float64) # 固定长度2，类型为浮点数
    smallest_value = np.inf
    for x in xs:
        for y in ys:
            value = x*x + y*y
            if value < smallest_value:
                smallest_value = value
                res[0] = x
                res[1] = y
    return res, smallest_value

if __name__ == "__main__":
    print(test())

方案3：使用Numba自带的typed.List #

如果确实需要用动态列表，可以用Numba支持的类型化列表明确声明类型：

import numpy as np
from numba import njit
from numba.typed import List

@njit()
def test():
    xs = np.arange(0, 1, 0.1)
    ys = np.arange(0, 1, 0.1)

    res = List()
    res.append(0.0)
    res.append(0.0) # 提前固定长度和类型
    smallest_value = np.inf
    for x in xs:
        for y in ys:
            value = x*x + y*y
            if value < smallest_value:
                smallest_value = value
                res[0] = x
                res[1] = y
    return res, smallest_value

if __name__ == "__main__":
    print(test())

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