Numba即时编译：列表类型与多返回值的函数签名配置问题

一、异构返回列表的签名问题 #

你的Cal函数返回的是包含整数、二维numpy数组、整数的异构列表，Numba的ListType要求列表元素类型必须同构，所以直接用nb.ListType()是无效的。这里提供两种处理方案：

1. 使用object类型列表（不推荐，性能差） #

如果必须用列表返回异构元素，可以指定元素类型为nb.pyobject，但这种方式会失去Numba的编译优化优势，因为object类型无法被高效处理：

import numba as nb
import numpy as np

@nb.jit(nb.ListType(nb.pyobject))
def Cal():
    return [1, np.zeros(shape=[5, 5]), 3]

a = Cal()

2. 改用元组返回（推荐，性能最优） #

Numba对异构元组的支持远好于异构列表，你可以把返回值改成元组，然后用nb.types.Tuple定义每个元素的具体类型：

import numba as nb
import numpy as np

# 签名对应：(int64, 二维float64数组, int64)
@nb.jit(nb.types.Tuple((nb.int64, nb.float64[:, :], nb.int64)))
def Cal():
    return 1, np.zeros(shape=[5, 5]), 3

a = Cal()

二、多返回值的函数签名指定 #

Numba中多返回值本质是返回一个元组，所以直接用nb.types.Tuple包裹每个返回值的类型即可。以你的TwoOutput函数为例：

import numba as nb

# 签名对应返回两个int64类型的值
@nb.jit(nb.types.Tuple((nb.int64, nb.int64)))
def TwoOutput():
   return 1, 2

x, y = TwoOutput()

如果想简化写法，也可以用Numba的平台相关类型（比如nb.int_表示当前平台的默认整数类型）：

@nb.jit(nb.types.Tuple((nb.int_, nb.int_)))
def TwoOutput():
   return 1, 2

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