Maybe it is trivial, but I was wondering how to write signatures in the jit decorator when there are several outputs.
For instance :
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
Source : my own experiments, and the source code of Numba : https://github.com/numba/numba
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
According to this newsgroup post you can specify using numba.typeof(<an example of your tuple>)
For example
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
You could also build them from the components given in numba.types, but that's probably more work than using typeof
The fact it can do this in nopython mode suggests performance should be OK (tuple unpacking is explicitly listed as a supported feature http://numba.pydata.org/numba-doc/dev/reference/pysupported.html). However, I haven't actually tested the performance.
