A few weekends and numba signatures

Question

This may be trivial, but I was wondering how to write signatures in the jit decorator when there are multiple 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?

Answer 1

You can use explicit declarations or a string declaration:

Tuple with homogeneous types:

 @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 

Tuple with heterogeneous types:

 @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 Numba source code: https://github.com/numba/numba

Of course, the solution suggested by DavidW is a great solution if you don't know the exact type:

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

Answer 2

According to this newsgroup post, you can specify using numba.typeof(<an example of your tuple>)

for instance

 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 can also compile them from the components specified in numba.types , but this is probably more work than using typeof

The fact that he can do this in nopython mode means that the performance should be in order (unpacking tuples is explicitly indicated as a supported function http://numba.pydata.org/numba-doc/dev/reference/pysupported .html ). However, I did not actually test the performance.