Возможно, это тривиально, но мне было интересно, как писать сигнатуры в декораторе 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
Как насчет выступлений? Лучше ли писать две разные функции?
Вы можете использовать явные объявления или строковое объявление:
@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
Согласно этому сообщению группы новостей, вы можете указать, используя 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). Тем не менее, я на самом деле не проверял производительность.