CODEWITHSUNDEEP
pythonnumba
Ipse Lium
Ipse Lium

Reputation: 1080

Multiple output and numba signatures

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 ?

Upvotes: 53

Views: 21353

Answers (2)

Jean Lescut
Jean Lescut

Reputation: 1467

You can either use explicit declarations or 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 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

Upvotes: 74

DavidW
DavidW

Reputation: 30930

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.

Upvotes: 15

Related Questions