Problem with Numba jit: "Typing error" and "All templates rejected with/without literals"

Admin 1
Viewed 1786

I am implementing a program that solves a differential equation in Python 3.7.3 and there is one function that I just can't get to compile with Numba. The most recent version of it is:

import numpy as np
from numba import jit, uint16, complex128, prange

# Here is the setup of the program, as well as variable initialization

@jit((complex128[:, :, :], uint16, complex128[:, :], complex128[:, :], complex128[:, :]), nopython=True)
def upd_x(rhs: np.ndarray, m: int, s: np.ndarray, a: np.ndarray, b: np.ndarray) -> np.ndarray:
    x = np.zeros((3, m, m//2+1))
    x[2] = s*(1-a*(rhs[0]+rhs[1]))
    for i in range(2):
        x[i] = a*(rhs[i]+b*x[2])
    return x

What it is supposed to do is, take the "right hand side" (rhs) of the equation and update x (x has 3 components which are real fields, and the code is "updating" it in Fourier space, that is why the last axis is m//2+1 instead of m) with Schur's complement method. When I ran the code I got the following message:

Traceback (most recent call last):
  File "C:/Users/Username/Desktop/Program/Program.py", line 95, in <module>
    @jit((complex128[:, :, :], uint16, complex128[:, :], complex128[:, :], complex128[:, :]), nopython=True)
  File "C:\Program Files\Python37\lib\site-packages\numba\decorators.py", line 186, in wrapper
    disp.compile(sig)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "C:\Program Files\Python37\lib\site-packages\numba\dispatcher.py", line 659, in compile
    cres = self._compiler.compile(args, return_type)
  File "C:\Program Files\Python37\lib\site-packages\numba\dispatcher.py", line 83, in compile
    pipeline_class=self.pipeline_class)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 955, in compile_extra
    return pipeline.compile_extra(func)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 377, in compile_extra
    return self._compile_bytecode()
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 886, in _compile_bytecode
    return self._compile_core()
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 873, in _compile_core
    res = pm.run(self.status)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler_lock.py", line 32, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 254, in run
    raise patched_exception
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 245, in run
    stage()
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 501, in stage_nopython_frontend
    self.locals)
  File "C:\Program Files\Python37\lib\site-packages\numba\compiler.py", line 1105, in type_inference_stage
    infer.propagate()
  File "C:\Program Files\Python37\lib\site-packages\numba\typeinfer.py", line 915, in propagate
    raise errors[0]
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function setitem>) with argument(s) of type(s): (array(float64, 3d, C), Literal[int](2), array(complex128, 2d, C))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
In definition 2:
    All templates rejected with literals.
In definition 3:
    All templates rejected without literals.
In definition 4:
    All templates rejected with literals.
In definition 5:
    All templates rejected without literals.
In definition 6:
    All templates rejected with literals.
In definition 7:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: typing of staticsetitem at C:/Users/User/Desktop/Program/Program.py (98)

File "Programa.py", line 98:
def upd_x(rhs, m, s, a, b):
    <source elided>
    x = np.zeros((3, m, m//2+1))
    x[2] = s*(1-a*(rhs[0]+rhs[1]))
    ^

I don't understand why the error message suggests that the variable type is unsupported, and I also don't know what is wrong that I need to correct. The versions that I am using are numba==0.44.1, numpy==1.16.1.

Thank you very much.

typeerror numba runtime-error python-3.x python
1 Answers

It looks like Numba was unable to determine what the type of the output x, is, so I added a dtype to x. Then you run into mixing np.int64 and uint16, in size arguments to np.zeros since 3 is interpreted as a i64. So the following will compile:

import numpy as np
from numba import jit, uint16, complex128, prange

# Here is the setup of the program, as well as variable initialization

@jit(complex128[:,:,:](complex128[:, :, :], uint16, complex128[:, :], complex128[:, :], complex128[:, :]), nopython=True)
def upd_x(rhs: np.ndarray, m: int, s: np.ndarray, a: np.ndarray, b: np.ndarray) -> np.ndarray:
    mx = np.int64(m)
    x = np.zeros((3, mx, mx//2+1), dtype=np.complex128)
    x[2] = s*(1-a*(rhs[0]+rhs[1]))
    for i in range(2):
        x[i] = a*(rhs[i]+b*x[2])
    return x

Also, note that I added a return type to the signature being passed to @jitalthough this isn't necessary, I believe.

And so I use the inputs:

m = 4
x = np.zeros((3, m, m//2+1), dtype=np.complex128) + 2 + 2j
y = np.zeros((m, m//2 + 1 ), dtype=np.complex128) + 1 + 1j

upd_x(x, np.uint16(m), y, y, y)

and this gives back something sensible, I think.

user_3929490
Related