多输出和numba签名

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

您可以使用explicite声明或字符串声明:

同类型的元组:

例如,我们可以:

@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

根据此新闻组帖子,您可以指定使用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).但是,我实际上没有测试过性能.