Support arguments that are functions in (n)jit

[mroeschke]

Feature request

Would it be possible to support arguments that are functions when using jit that don't execute in object mode?

In [19]: @jit
    ...: def f(arg, func):
    ...:     return func(arg)
    ...:

In [20]: f([1], sum)
Compilation is falling back to object mode WITH looplifting enabled because Function "f" failed type inference due to: non-precise type pyobject
[1] During: typing of argument at <ipython-input-19-32db0f8a0a3e> (3)

File "<ipython-input-19-32db0f8a0a3e>", line 3:
def f(arg, func):
    return func(arg)
    ^
...
Out[20]: 1

In [21]: @njit
    ...: def f(arg, func):
    ...:     return func(arg)
    ...:

In [22]: f([1], sum)
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
[1] During: typing of argument at <ipython-input-21-76dc76439c5b> (3)

File "<ipython-input-21-76dc76439c5b>", line 3:
def f(arg, func):
    return func(arg)
    ^

[stuartarchibald]

Thanks for the query. Numba supports JIT compiled functions as arguments, e.g.:

from numba import njit
from numpy import sum
import numpy as np

@njit
def f(arg, func):
    return func(arg)

jit_sum = njit(lambda x: sum(np.asarray(x)))

f([1], jit_sum)

does this help?

[mroeschke]

Thanks @stuartarchibald! Just what I was looking for. Thanks for answering all my usage questions.

[stuartarchibald]

@mroeschke no problem, glad to be of help!