多个输出和 numba 签名

Posted 2023-02-16

【中文标题】多个输出和 numba 签名

【英文标题】：Multiple output and numba signatures

【发布时间】：2015-08-02 12:03:18

【问题描述】：

也许这很简单，但我想知道当有多个输出时如何在 jit 装饰器中编写签名。

例如：

import numba as nb

@nb.jit(['???(int32, int32, float(:,:), float(:,:))'], nopython=True)
def foo(nx, ny, a, b):
    for i in range(nx):
        for i in range(ny):
            do stuff with a & b
    return a, b

表演怎么样？写两个不同的函数会更好吗？

【参考方案1】：

您可以使用显式声明或字符串声明：

具有同质类型的元组：

@nb.jit(nb.types.UniTuple(nb.float64[:],2)(nb.float64[:]),nopython=True)
def f(a) :
    return a,a

@nb.jit('UniTuple(float64[:], 2)(float64[:])',nopython=True)
def f(a) :
    return a,a

具有异构类型的元组：

@nb.jit(nb.types.Tuple((nb.float64[:], nb.float64[:,:]))(nb.float64[:], nb.float64[:,:]),nopython=True)
def f(a, b) :
    return a, b

@nb.jit('Tuple((float64[:], float64[:,:]))(float64[:], float64[:,:])',nopython=True)
def f(a, b) :
    return a, b

来源：我自己的实验，Numba的源代码：https://github.com/numba/numba

当然，当您不知道确切的类型时，DavidW 提出的解决方案是一个很好的解决方法：

@nb.jit(nb.typeof((1.0,1.0))(nb.double),nopython=True)
def f(a):
  return a,a

【参考方案2】：

根据this newsgroup post，您可以使用numba.typeof(<an example of your tuple>)指定

例如

import numba as nb

# I've put "nopython=True" just to demonstrate it still works
# whether you need it is your choice
@nb.jit(nb.typeof((1.0,1.0))(nb.double),nopython=True)
def f(a):
  return a,a

print f(5.0) # returns 5.0,5.0

您也可以使用numba.types 中给出的组件来构建它们，但这可能比使用typeof 需要更多的工作

它可以在 nopython 模式下执行此操作的事实表明性能应该没问题（元组解包被明确列为支持的功能http://numba.pydata.org/numba-doc/dev/reference/pysupported.html）。但是，我还没有实际测试过性能。