python - 多个输出和 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

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

具有同类类型的元组:

@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