CODEWITHSUNDEEP
pythontypesprotocolssignaturenumba
Post169
Post169

Reputation: 718

Numba signatures protocol

Despite searching Stack Overflow, and the internet as a whole, and reading several Stack Overflow questions and numba.pydata.org pages, and learning some clues to how to tell Numba what types I want to give to and get from functions, I'm not finding the actual logic of how it works.

For example, I experimented with a function that processes a list of integers and puts out another list of integers, and while the decorator @numba.jit(numba.int64[:](numba.int64[:])) worked, the decorators @numba.njit(numba.int64[:](numba.int64[:])) and @numba.vectorize(numba.int64[:](numba.int64[:])) did not work.

(njit successfully went past the decorator and stumbled on the function itself; I'm guessing that concatenating elements to a list is not an available function in 'no python' mode. vectorize, however, complains about the signature, TypeError: 'Signature' object is not iterable; maybe it's worried that a 1D array could include a single element without brackets, which is not an iterable?)

Is there a simple way to understand how Numba works to enough depth to anticipate how I should express the signature?

Upvotes: 1

Views: 2613

Answers (1)

chrisb
chrisb

Reputation: 52276

The simplest answer for jit (and njit which is just and alias for nopython=True) is try to avoid writing signatures altogether - in the common cases, type inference will get you there.

Specific to your question numba.int64[:](numba.int64[:]) is a valid signature and works for jit.

numba.vectorize- expects an iterable of signatures (error message), so your signature(s) need to be wrapped in a list. Additional, vectorize creates a numpy ufunc, which is defined by scalar operations (which are then broadcast), so your signature must be of scalar types. E.g.

@numba.vectorize([numba.int64(numba.int64)])
def add_one(v):
    return v + 1

add_one(np.array([4, 5, 6], dtype=np.int64))
# Out[117]: array([5, 6, 7], dtype=int64)

Upvotes: 2

Related Questions