Numba njit方法赋值单元素集合触发AssertionError问题排查

错误原因 #

Numba的njit采用静态类型推断机制，编译时就会确定函数内变量的类型。你的代码里，初始赋值s = {1,2,3}会让Numba将s推断为多元素集合类型；而else分支中s = {2}会被推断为单元素集合类型——在Numba 0.60.0版本中，这两种被视为不同类型，静态类型系统不允许同一变量在不同分支持有不同类型，因此触发AssertionError。

虽然官方文档说明支持所有集合操作，但早期Numba版本对集合的类型推断存在局限，没有将单/多元素集合统一为同一种可变集合类型。

解决办法 #

有三种可行的修复方式：

1. 统一集合的初始化形式 #

给单元素集合添加尾逗号，让Numba将其推断为多元素集合类型，和初始赋值的类型保持一致：

import numba

@numba.njit
def foo(n: int):
    s = {1, 2, 3}
    if n == 1:
        pass
    else:
        s = {2,}  # 加逗号，强制推断为多元素集合类型

    res = sum(s)
    return res

def main():
    res = foo(2)
    print(res)

2. 用空集合初始化，后续动态添加元素 #

先定义空集合，再通过add或update方法填充元素，这样变量类型始终是统一的可变集合类型，不受元素数量影响：

import numba

@numba.njit
def foo(n: int):
    s = set()
    if n == 1:
        s.update({1, 2, 3})
    else:
        s.add(2)

    res = sum(s)
    return res

def main():
    res = foo(2)
    print(res)

3. 升级Numba版本 #

Numba后续版本（如0.61及以上）修复了集合类型推断的这个问题，允许同一变量赋值不同元素数量的集合（只要元素类型一致）。升级到较新版本后，你的原始代码就能正常运行。

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