Cannot compile in nopython mode for a "specialized" function

Hi,

I am working on a physics problem where I assemble a matrix using an analytic function which has parameters of two kinds:

• ones available ahead of time (predef_param below)
• ones that change during the assembling of the matrix (runtime_param below)

To make my assembling code more general, I have tried to separate both concerns (ahead of time vs. runtime) so that I can reuse it.

This results in the fact that I need to have a function simpler_func with less parameters calling the original complex_func.

I have code that looks approximately like this:

PARAM = 5.

def complex_func(runtime_param, predef_param):
    return np.cos(predef_param) ** runtime_param

@nb.jit(forceobj=True)
def simpler_func(runtime_param):
    return complex_func(runtime_param, PARAM)

@nb.njit
def assemble_matrix(N):
    mat = np.zeros((N, N))
    mapping = np.arange(N)
    for i in range(N):
        I = mapping[i]
        for j in range(N):
            J = mapping[j]
            mat[i, j] = simpler_func(I + J)
    return mat

def main():
    assemble_matrix(5)

Running the main function results in the following error with numba 0.49.1:

TypingError                               Traceback (most recent call last)
<ipython-input-25-263240bbee7e> in <module>
----> 1 main()

<ipython-input-24-22ddcf18b8e0> in main()
     20 
     21 def main():
---> 22     assemble_matrix(5)

C:\Anaconda3\envs\pyrus_dev\lib\site-packages\numba\core\dispatcher.py in _compile_for_args(self, *args, **kws)
    399                 e.patch_message(msg)
    400 
--> 401             error_rewrite(e, 'typing')
    402         except errors.UnsupportedError as e:
    403             # Something unsupported is present in the user code, add help info

C:\Anaconda3\envs\pyrus_dev\lib\site-packages\numba\core\dispatcher.py in error_rewrite(e, issue_type)
    342                 raise e
    343             else:
--> 344                 reraise(type(e), e, None)
    345 
    346         argtypes = []

C:\Anaconda3\envs\pyrus_dev\lib\site-packages\numba\core\utils.py in reraise(tp, value, tb)
     78         value = tp()
     79     if value.__traceback__ is not tb:
---> 80         raise value.with_traceback(tb)
     81     raise value
     82 

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of type(CPUDispatcher(<function simpler_func at 0x0000024D428E8EE8>)) with parameters (int64)
 * parameterized
[1] During: resolving callee type: type(CPUDispatcher(<function simpler_func at 0x0000024D428E8EE8>))
[2] During: typing of call at <ipython-input-24-22ddcf18b8e0> (18)


File "<ipython-input-24-22ddcf18b8e0>", line 18:
def assemble_matrix(N):
    <source elided>
            J = mapping[j]
            mat[i, j] = simpler_func(I + J)

Is there a way to make this work? What am I doing wrong? Does this have something to do with the fact that functions can’t be passed around well?

Thank you in advance for your help.

Florian

[HPLegion]

Not a solution, but a clarification as to what may be going wrong here: You cannot call a function that was compiled in object mode (forceobj=True) from a function that was called in nopython mode, as is illustrated by this example:

from numba import jit, njit

@jit(forceobj=True)
def f(x):
    return x

@njit
def g(x):
    return f(x)

g(1)

Invalid use of type(CPUDispatcher(<function f at 0x7f7c5b4ae4d0>)) with parameters (int64)
 * parameterized
[1] During: resolving callee type: type(CPUDispatcher(<function f at 0x7f7c5b4ae4d0>))
[2] During: typing of call at njit_calls_jit.py (9)


File "njit_calls_jit.py", line 9:
def g(x):
    return f(x)
    ^

Maybe there would be a way to do this by explicitly using the object mode context manager but this may bring with it a hefty performance penalty.

[flothesof]

Indeed, this might actually allow me to do it “as is”.

In the case I was working on when I submitted the issue, the blocker really was the scipy.special.factorial2 function. Reimplementing it with numba-accelerated python let me do what I wanted. I will probably be confronted with some more elaborate cases in the near future, where I won’t have the knowledge to reimplement the function in pure-python in which case I’ll try to use obj mode.

Given that my problem is solved and a workaround provided for further cases, I’ll close the issue now. Feel free to reopen if you feel it’s a better choice.

Thanks to all for your help.

Regards Florian