多个输出和 numba 签名 [英] Multiple output and numba signatures

问题描述

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

Maybe it is trivial, but I was wondering how to write signatures in the jit decorator when there are several outputs.

例如:

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

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

What about the performances ? Is it better to write two different functions ?

推荐答案

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

You can either use explicit declarations or string declaration :

@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

Source : my own experiments, and the source code of Numba : https://github.com/numba/numba

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

Of course, the solution proposed by DavidW is an excellent workaround when you do not know the exact type :

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

这篇关于多个输出和 numba 签名的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！