Numba jit的问题："Typing error"和"All templates rejected with/without literals"

Question

我正在实现一个解决Python 3.7.3中微分方程的程序，并且有一个函数我无法使用Numba编译。它的最新版本是：

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

它应该做的是，取方程的"右侧"(rhs)并更新x(x有3个分量是实场，并且代码在傅里叶空间中"更新"它，这就是为什么最后一个轴是用舒尔补码方法m//2+1而不是m)。当我运行代码时，我收到以下消息：

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 FilesPython37libsite-packagesnumbadecorators.py", line 186, in wrapper
disp.compile(sig)
File "C:Program FilesPython37libsite-packagesnumbacompiler_lock.py", line 32, in _acquire_compile_lock
return func(*args, **kwargs)
File "C:Program FilesPython37libsite-packagesnumbadispatcher.py", line 659, in compile
cres = self._compiler.compile(args, return_type)
File "C:Program FilesPython37libsite-packagesnumbadispatcher.py", line 83, in compile
pipeline_class=self.pipeline_class)
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 955, in compile_extra
return pipeline.compile_extra(func)
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 377, in compile_extra
return self._compile_bytecode()
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 886, in _compile_bytecode
return self._compile_core()
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 873, in _compile_core
res = pm.run(self.status)
File "C:Program FilesPython37libsite-packagesnumbacompiler_lock.py", line 32, in _acquire_compile_lock
return func(*args, **kwargs)
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 254, in run
raise patched_exception
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 245, in run
stage()
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 501, in stage_nopython_frontend
self.locals)
File "C:Program FilesPython37libsite-packagesnumbacompiler.py", line 1105, in type_inference_stage
infer.propagate()
File "C:Program FilesPython37libsite-packagesnumbatypeinfer.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]))
^

我不明白为什么错误消息表明变量类型不受支持，我也不知道我需要纠正什么错误。我使用的版本是numba==0.44.1，numpy==1.16.1。

谢谢。

Answer 1

看起来 Numba 无法确定输出x的类型是什么，所以我在x中添加了一个dtype。然后你会遇到混合np.int64和uint16，在大小参数中np.zeros3因为被解释为i64。因此，以下内容将编译：

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

另外，请注意，我在传递给@jit的签名中添加了一个返回类型，尽管我相信这不是必需的。

所以我使用输入：

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)

我认为，这回馈了一些明智的东西。