Numba typed lists with homogeneous-signature functions

Numba Support: How do I do ...?
1

I have some numba-jitted event functions with equal signature, i.e.:

from numba import jit

@jit("Tuple((float64,float64))(float64[::1])", nopython=True)
def event_1(y): 
    return 1.1, 1.2 # (random values for this example)

@jit("Tuple((float64,float64))(float64[::1])", nopython=True)
def event_2(y): 
    return 2.1, 2.2 # (random values for this example)

My goal is to create a jitted-function which returns a List of event functions. The event_handler output is supposed to be the input of another cache-compiled jitted-function and so its signature must be fixed:

from numba.typed import List
@jit("ListType(FunctionType(Tuple((float64, float64))(float64[::1])))()", nopython=True)
def event_handler(): 
    return List([event_1, event_2])

However, the code above only works if the List returned by event_handler has at least two different event functions. If the List has only 1 event-function item, or multiple items of the same function (e.g., either List([event_1]) or List([event_1, event_1]) ), the code above cannot compile and produces the following error:

No conversion from ListType[type(CPUDispatcher(<function event_1 at 0x7f83c2a22430>))] to ListType[FunctionType[UniTuple(float64 x 2)(array(float64, 1d, C))]]

I believe the reason is because in the latter case the List item-type is set equal to the function CPUDispatcher rather than the typed function signature. I’ve already tried the solution in the topic named “typed list of jitted functions in jitclass” (I am unable to add links) to initialise the list but unfortunately it does not work when the list is passed from one function to another.

How can I solve this issue? It seems strange to me that the list type is inferred correctly when two different functions are provided but suddendly inherits a CPUDispatcher type whenever a single item is given.

2

Does this example with explicit typing help work around the problem at all?

1 Like
4

Actually, it did help solve the problem, thank you!
I’m posting here the solution for anyone who stumbles upon the same issue. With reference to my original post, I can initialize the List with the empty_list method, passing as argument its signature (defined outside the scope of any jitted function):

from numba import types
from numba.typed import List

vector_sig = types.Array(dtype=types.float64, ndim=1, layout="C")
function_sig = types.FunctionType(types.Tuple((types.float64, types.float64))(vector_sig))

@jit("ListType(FunctionType(Tuple((float64, float64))(float64[::1])))()", nopython=True)
def event_handler(): 
    event_list = List.empty_list(function_sig)
    event_list.append(event_1)
    return event_list