Multiple output and numba signatures

Question

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 ?

Asked By: Ipse Lium

||

Source

Answer 1

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.

Answered By: DavidW

Answer 2

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

Answered By: Jean Lescut

Multiple output and numba signatures

Question:

Answers:

Tuple with homogeneous types :

Tuple with heterogeneous types :