Question

也许这是微不足道的，但我想知道如果有几个输出，如何在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

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

Answer 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

Answer 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）。但是，我还没有真正测试过它的表现。

多输出和numba签名

2 个答案:

具有同类型的元组：

具有异构类型的元组：