如何强制Numba返回NumPy类型？

嘿，我太懂你这种困惑了——Numba自动把NumPy整数类型转成Python原生int的行为确实有点反直觉，不过别发愁，咱们有办法让它乖乖返回NumPy类型！

先把你给出的测试代码补全，复现这个问题看看：

import numba as nb
import numpy as np

print(f"Numba version: {nb.__version__}")  # 0.59.0
print(f"NumPy version: {np.__version__}")  # 1.23.5

# 显式定义签名
sig = nb.uint32(nb.uint32, nb.uint32)

@nb.njit(sig, cache=False)
def test(a, b):
    return a + b

# 测试
a_np = np.uint32(3)
b_np = np.uint32(7)
result = test(a_np, b_np)
print(f"当前返回类型: {type(result)}")  # 输出是<class 'int'>，不是我们要的np.uint32

解决办法：显式用NumPy类型构造函数包裹返回值 #

这是最直接可靠的方式，在函数返回的时候，用你想要的NumPy类型把结果包起来就行，Numba会尊重这个显式的转换操作：

import numba as nb
import numpy as np

print(f"Numba version: {nb.__version__}")  # 0.59.0
print(f"NumPy version: {np.__version__}")  # 1.23.5

sig = nb.uint32(nb.uint32, nb.uint32)

@nb.njit(sig, cache=False)
def test(a, b):
    # 关键：用目标NumPy类型包裹计算结果
    return np.uint32(a + b)

# 再测试一次
a_np = np.uint32(3)
b_np = np.uint32(7)
result = test(a_np, b_np)
print(f"现在的返回类型: {type(result)}")  # 输出<class 'numpy.uint32'>，达成目标！

补充说明 #

Numba之所以默认把NumPy标量转成Python原生类型，主要是为了让JIT函数和普通Python代码的交互更“无缝”——毕竟Python里大家更常用原生int/float。但如果你需要严格保持NumPy类型的一致性，上面的显式转换方法就完全够用，而且对性能几乎没有影响，Numba会把这个转换操作编译成高效的机器码。

备注：内容来源于stack exchange，提问作者Raven