Несколько выходных и numba-сигнатур

Возможно, это тривиально, но мне было интересно, как писать сигнатуры в декораторе 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

Как насчет выступлений? Лучше ли писать две разные функции?

Ответы

Ответ 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

Ответ 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). Тем не менее, я на самом деле не проверял производительность.