How to return list for eager compilation?

The numba allows eager compilation by telling the function signature. But, I can not find some information about the list type. My test code is:

import numba as nb
import numpy as np

# @nb.jit(nb.ListType())
def Cal():
    return [1, 2, 3]

a = Cal()

What’s the function signature for Cal? @nb.jit(nb.ListType()) is wrong.

In addition, what if there are three outputs? How to provide the function signature? For example:

def TwoOutput():
   return 1, 2

Any suggestion is appreciated

Dear @zhang-qiang-github,

There are a few things wrong with your example, but I understand that getting the typing correct can be difficult at the start. I highly encourage you to take a few hours and carefully read the documentation.

The signature must provide both return type information and argument information. A valid signature for a function that takes two integers and returns an integer looks like this:

import numba as nb 

sig = nb.int64(nb.int64, nb.int64)

Note that you call the return type with the argument types. This is incorrect in your first code snippet. You instantiate the type instance, but don’t call it with any return type. You must call the type instance even when not taking any arguments:

sig = nb.int64()

If you don’t return, you return void:

sig = nb.void(nb.int64)

I think it would be most helpful to just provide a few examples that you can study along with the documentation (e.g. here to learn more about lists). Note, I use nb.types.List instead of nb.types.ListType. Here the examples:

import numba as nb
from numba.types import List, int64

arg_types = ()
ret_type = List(int64, reflected=False, initial_value=[1, 2, 3])
sig = ret_type(*arg_types)

@nb.jit(sig)
def f1(): 
    return [1, 2, 3]


arg_types = ()
ret_type = List(int64, reflected=False)
sig = ret_type(*arg_types)

@nb.jit(sig)
def f2(): 
    a = []
    a.extend([1, 2, 3])
    return a

arg_types = List(int64, reflected=True), List(int64, reflected=True)
ret_type = List(int64, reflected=True)
sig = ret_type(*arg_types)

@nb.jit(sig)
def f3(a, b): 
    a.extend(b)
    return a


f1()
f2()
f3([1, 2, 3], [4, 5, 6])

For your second example, you must understand that Python also cannot return multiple values. What feels like multiple return values is the returning of a tuple followed by tuple unpacking. So you also return a tuple in jitted functions. If all the return types are the same, you return a UniTuple and a Tuple otherwise.

This looks like this:

import numba as nb
from numba.types import UniTuple, Tuple, int64, float64, bool_, void

arg_types = ()
ret_type = UniTuple(int64, 3)
sig = ret_type(*arg_types)

@nb.jit(sig)
def f4(): 
    return (1, 2, 3)

arg_types = ()
ret_type = Tuple((int64, float64, bool_))
sig = ret_type(*arg_types)

@nb.jit(sig)
def f5(): 
    return (1, 2.0, True)


f4()
f5()

Finally, I would like to give you a small hack that may be helpful when working with types in Numba. You can even obtain complex types by using the nb.typeof:

var = ((1, 2), [1, 2], True, None)
nb.typeof(var)

# gives: Tuple(UniTuple(int64 x 2), reflected list(int64)<iv=None>, bool, none)

Another convenient feature is to call your function without providing the signature and then inspect the signature Numba created:

var = ((1, 2), [1, 2], 3, True, None)

@nb.njit
def foo(arg):
    pass

foo(var)
foo.signatures
# gives: [(Tuple(UniTuple(int64 x 2), reflected list(int64)<iv=None>, int64, bool, none),)]

Hope this helps you get started.

2 Likes

Many thanks for this kindly reply. It help me a lot.