时间戳

首页 > 解决方案 > Numba - 如何为可变数量的函数参数声明签名？

python - Numba - 如何为可变数量的函数参数声明签名？

问题描述

Numba 似乎支持 jit-ed 函数中可变数量的参数：

@numba.jit(nopython=True)
def func(*arrays):
    return np.vstack(arrays)

但是，我想知道为这个函数提供签名的正确语法是什么，以便 numba 知道 tuple 的成员应该使用哪种类型arrays。我在文档中找不到对此的引用。

标签： pythonnumba

解决方案

我不确定您是否应该尝试手动提供类型签名，因为 Numba 能够进行类型推断。引用有关高级类型的 Numba 文档：

警告

此处记录的 API 不保证是稳定的。除非必要，否则建议让 Numba 使用 @jit 的无签名变体来推断参数类型。

为了演示，观察我们可以在 nopython 模式下编译这个函数，Numba 会在使用时建立一个签名列表：

import numba
import numpy as np


@numba.jit(nopython=True)
def func(*arrays):
    return np.vstack(arrays)


print("signatures before call", func.nopython_signatures)

arr1 = np.array([1, 2])
arr2 = np.array([3, 4])
arr3 = func(arr1, arr2)
print(arr3)

print("signatures after call", func.nopython_signatures)

结果输出：

signatures before call []
[[1 2]
 [3 4]]
signatures after call [(UniTuple(array(int64, 1d, C) x 2),) -> array(int64, 2d, C)]

推荐阅读