Feature Request: Function as first-class type

[NicolasHug]

• I am using the latest released version of Numba

~ » python -c "import numba; print(numba.__version__)" 
0.40.0

• I have included below a minimal working reproducer

I am getting a type unification error that I can't logically explain:

• the same code works when I'm not inside a jitclass (see foo() vs C.foo())
• even inside the jitclass, the unification can work under some conditions (see C.bar)

EDIT: as noticed by @ogrisel this isn't related to jitclass but more generally to the nopython mode. I updated the example and the log.

Minimal working example:

from numba import njit, jit, jitclass

@njit
def a():
    return 2

@njit
def b():
    return 3

def c(arg):  # works
    if arg:
        f = a
    else:
        f = b
    out = f()
    return out

@jit
def d(arg):  # works
    if arg:
        f = a
    else:
        f = b
    out = f()
    return out

@njit
def e(arg):  # works
    f = a
    f = b
    out = f()
    return out

@njit
def g(arg):  # fails
    if arg:
        f = a
    else:
        f = b
    out = f()
    return out

# OK
c(True)
c(False)

# OK
d(True)
d(False)

# OK
e(True)
e(False)

# KO
g(True)
g(False)

Error log:

Traceback (most recent call last):
  File "lol.py", line 57, in <module>
    g(True)
  File "/home/nico/.virtualenvs/pygbm/lib/python3.7/site-packages/numba/dispatcher.py", line 348, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/nico/.virtualenvs/pygbm/lib/python3.7/site-packages/numba/dispatcher.py", line 315, in error_rewrite
    reraise(type(e), e, None)
  File "/home/nico/.virtualenvs/pygbm/lib/python3.7/site-packages/numba/six.py", line 658, in reraise
    raise value.with_traceback(tb)
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot unify type(CPUDispatcher(<function a at 0x7fd67bd5af28>)) and type(CPUDispatcher(<function b at 0x7fd66bdb0378>)) for 'f', defined at lol.py (38)

File "lol.py", line 38:
def g(arg):  # fails
    <source elided>
    if arg:
        f = a
        ^

[1] During: typing of assignment at lol.py (40)

File "lol.py", line 40:
def g(arg):  # fails
    <source elided>
    else:
        f = b
        ^

[stuartarchibald]

Thanks for the report. I can reproduce. I think this is because the type information about the result of calling member function foo cannot be discerned in the case with the branch.

[ogrisel]

Note that this is not specific to jitclass: if you decorate the foo function with an @njit decorator you get a similar error message.

[NicolasHug]

You're right, thanks, I updated the example

[stuartarchibald]

Thanks for the update, I think I see what's going on now. I thought that this was something more complicated to do with the jitclass internals but I believe in actual fact the issue is more simple. It is that when doing type inference Numba sees the passed in functions as unique types, at present it doesn't know about what the function itself returns, there is no "function type" type. As a result the type inference mechanism sees something equivalent in objectionableness to:

@njit
def g(arg):  # fails
    if arg:
        f = np.int32(a)
    else:
        f = np.array(10)
    out = f()
    return out

Moving this to a feature request, functions as first-class types.

[luk-f-a]

hi, two questions:

• does this feature cover use cases as in using a list/tuple of jitted functions #2542?

• where would this fall in the roadmap (https://numba.pydata.org/numba-doc/dev/developer/roadmap.html)? will it make it before v1.0?

thanks!

[seibert]

It would enable #2542, I believe.

As for roadmap, the soonest would be in the Medium Term bucket, probably post 1.0.

[luk-f-a]

thanks @seibert
For now I'm using a workaround based on sklam's "chain" suggestion.

[seibert]

@sklam: What do you think is the level of effort to make a function type that matches on type signature? (rather than what we have now, where every function is its own type)

[sklam]

I don't have a good estimate yet. But, I have included some notes on possible solutions to it.

Let's use the following code as an example:

def ex(f, g, choose, args0, args1):
    if choose:
        h = f         # pos A
    else:
        h = g         # pos B
    
    x = h(*args0)     # pos C
    y = h(*args1)     # pos D
    return x, y

Option 1. Continuing to allow FunctionType to reference the actual function objects.

Thinking through type inference:

• At A, h is typed as FunctionType(f).
• At B, h is unified to MultiFunctionType([f, g])
• At C, the h(*args0) call causes type-check on f(*args0) and g(*args0).
• At D, the h(*args1) call causes type-check on f(*args1) and g(*args1).

Thinking through lowering:

• The call site at C and D may invoke different versions of h (and different versions of f and g).
• At C and D, the call can be lowered into a switch that contains different call sequence to invoke f and g respectively.

Notes:

• This approach is not scalable. The caller function must know all possible called functions and a new version must be created per new function added.

Option 2. Make functions runtime objects

Thinking through type inference:

• At A, h is typed as FunctionType(f).
• At B, h is unified to MultiFunctionType() (no reference to actual function object)
• At C, the h(*args0) must be able to produce the result type of the call. Without a reference to an actual function object, we can't proceed.

(backing up....)

• At A, h is typed as FunctionType(f).
• At B, h is unified to MultiFunctionType(f) (with reference to the first object only)
• At C, the h(*args0) uses f (the first seen function object) as a typing-template for all possible function value here. It will provide a signature, which we will call sig0.
• At D, the h(*args1) produces a signature, sig1.

Whenever ex is called, arguments binding to f and g must type-check to be compatible to sig0 and sig1. If valid, a call-table is emitted for the respective signatures and a dispatcher object with that call-table is used as the runtime value.

Notes:

• This option is more scalable. No recompile of ex is needed when binding to f and g changes as long as they are compatible to the required call-signatures.
• But, if they are multiple sets of call-signatures, determining which set to use will become expensive.

[luk-f-a]

I'll go out on a limb here, as this is so far from my area of expertise. what about a more limited version that requires the functions to be set in advance?

I am thinking about sklam's chain in #2542 or this example, which is a modified version of that chain.

@nb.njit
def ident(out, x):
    pass

def chain_assign(fs, inner=ident):
    head, tail = fs[-1], fs[:-1]
    item = len(fs)-1
    
    @nb.njit
    def assign(out, x):
        inner(out, x)
        out[item] = head(x)
    
    if tail:
        return chain_assign(tail, assign)
    else:
        return assign

In practice this defines an iterator over functions, but an immutable one. I guess this is what makes it work, avoiding the type-checking and lowering problems described above.
What if we could define in a more convenient way such an iterator or maybe a function-tuple?
Each instance of the iterator/function-tuple would be its own type (similar to how at the moment each function is its own type). Each instance must be created with a reference to actual functions, such that they can be signature-checked against each other.
The example above could be written as:

def ex(func_tuple, choice_index, args0, args1):
    x = func_tuple[choice_index](*args0)     # pos C
    y = func_tuple[choice_index](*args1)     # pos D
    return x, y

Here's where my lack of experience with the internals of Numba stops me. I cannot reproduce the step by step analysis sklam's did. I hope the idea is not too stupid, though.

[ehsantn]

I think function types need to carry the signature to have consistent typing and straight-forward implementation (similar to C). In this way, things like containers of functions as in #2542 will be easy as well. Following @sklam's example (assuming sig0 as typing template of f):

• At A, h is typed as FunctionType(sig0).
• At B, h is unified to FunctionType(sig0), assuming g has the same typing template (is template unification possible?).
• At C, h(*args0) uses sig0 as a typing-template.
• At D, h(*args1) uses sig0 as a typing-template.

At lowering, the signature is known from function variable type, and the function pointer is taken from function variable value and called. Also, supporting more advanced cases like inlining specific cases can be done using literals, value-based specialization, etc. Does this make sense?

[sklam]

@luk-f-a, there is still a difficulty in typing the code:

def ex(func_tuple, choice_index, args0, args1):
    x = func_tuple[choice_index](*args0)     # pos C
    y = func_tuple[choice_index](*args1)     # pos D
    return x, y

It is difficulty to determine the function type from func_tuple[choice_index]. Currently, Numba will only work if choice_index is a literal so that it can determine the type at compile time. If choice_index is a runtime variable, we still need to unify the type of func_tuple[0] and func_tuple[1]. Thus, it will be the equivalent problem as my example in #3405 (comment).

@ehsantn's suggestion will require explicit annotation from the user. For instance,

sig = Signature(...)
f_sig = f.compile(sig)
g_sig = g.compile(sig)
exc(f_sig, g_sig, ...)

In fact, we already have the .compile method. It just doesn't carry the call-signature in it. It sounds doable.

[luk-f-a]

@sklam, thanks, I guess I was confused by the following example:

@nb.njit
def f(a):
    return a

tup = (f, f)

@nb.njit
def ex(tup, ind,a , b):
    x = tup[ind](a)
    y = tup[ind](b)
    return x,y

ex(tup, 0, 1,2)

This works well, despite in theory not being able to determine the function type from func_tuple[choice_index]. I assumed that the only problem was whether incoming tuples are classified as UniTuples and therefore getting a getitem method. So I guess we need some guarantee that the signature is the same and that can only be done after compile?

Would compile be explicitly necessary? or would it be enough to do

@njit(sig_str)
def f

@njit(sig_str)
def g

exc(f, g,...)

[luk-f-a]

@sklam, going back to @seibert 's question, what do you think is the level of effort for someone familiar with Numba internals? thanks!

[sklam]

Explicit compile is probably not necessary and @njit(sig_str) should be the same.

As for effort, I do not expect the initial changes would be big but there may be a lot of subtle points to make it useful. It'd worth to do a quick experimental patch to see how involved it would be.

[pearu]

Heads up: with PR #4967, all the examples in this issue work provided the input functions are cfunc-decorated.

[songololo]

Looking forward to seeing this develop.

I am finding the new Dict and List types quite useful: Out of curiosity, once #4967 is implemented, would it be possible to pass a numba function type to a List or Dict?

[pearu]

I think having List[<FunctionType>], for instance, would be a natural use case of first-class functions. In fact, I don't see reasons atm why it would not work out of the box.