如何为Numba设置嵌套列表类型签名？

你这段代码在使用Numba的parallel=True并手动指定三层嵌套列表返回类型时，触发了类型转换错误，具体表现为无法将并行环境中生成的嵌套列表转换为你指定的ListType(ListType(ListType(int32)))类型。

问题根源 #

Numba在并行模式下对可变嵌套列表的类型推断和内存管理存在限制，手动指定复杂的嵌套列表类型签名时，无法和并行循环中动态生成的列表结构匹配，从而抛出类型转换异常。

解决方案 #

方案1：移除手动类型标注，让Numba自动推断 #

删掉装饰器里的类型签名，Numba会自动处理嵌套列表的类型匹配：

from numba import njit, prange

@njit(fastmath=True, parallel=True, cache=True)
def test(x, y):
    a = []
    for i in prange(10):
        b = []
        for j in range(4):
            c = []
            for k in range(5):
                c.append(k)
            b.append(c)
        a.append(b)
    return a

方案2：改用多维数组替代嵌套列表（推荐） #

Numba对numpy数组的并行支持更稳定高效，嵌套列表在并行场景下的性能和兼容性都不如固定形状的数组：

from numba import njit, prange
import numpy as np

@njit(fastmath=True, parallel=True, cache=True)
def test(x, y):
    # 预先分配固定形状的三维数组
    a = np.zeros((10, 4, 5), dtype=np.int32)
    for i in prange(10):
        for j in range(4):
            for k in range(5):
                a[i, j, k] = k
    return a

注意事项 #

• 并行模式下尽量避免使用可变长度的嵌套列表，手动指定嵌套ListType类型签名极易触发类型不匹配问题。
• 如果必须使用嵌套列表，优先依赖Numba的自动类型推断，不要手动指定复杂的嵌套类型。

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