Numba函数类型匹配错误修复：numpy字符串数组传参报错解决

下面是简化后的程序代码：

import numpy as np
from numba import types, typed, njit


@njit(signature_or_function=types.unicode_type(types.unicode_type))
def my_func(symbol):
    return "*" + symbol + "*"


@njit
def run_numba(symbols):
    for symbol in symbols:
        result = my_func(symbol=symbol)  # 将symbol设为字面量"a"时可正常运行
        print(result)

symbols = np.array(["a", "b"], dtype="<U2")
run_numba(symbols)

运行时触发如下错误：

Traceback (most recent call last):
  File "/home/me/.pycharm_helpers/pydev/pydevconsole.py", line 364, in runcode
    coro = func()
           ^^^^^^
  File "<input>", line 17, in <module>
  File "/home/me/.pyenv/versions/3.11.1/envs/my_env/lib/python3.11/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/me/.pyenv/versions/3.11.1/envs/my_env/lib/python3.11/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of type(CPUDispatcher(&lt;function my_func at 0x7f172e1499e0&gt;)) with parameters (symbol=[unichr x 2])
Known signatures:
 * (unicode_type,) -> unicode_type
During: resolving callee type: type(CPUDispatcher(&lt;function my_func at 0x7f172e1499e0&gt;))
During: typing of call at <input> (13)
File "<input>", line 13:
<source missing, REPL/exec in use?>

注：如果把symbol变量替换为字符串字面量（比如直接写my_func(symbol="a")），程序可以正常运行，请问如何修复这段代码？

错误核心是Numba的类型匹配不兼容：给my_func指定的签名是通用unicode_type(unicode_type)，但从numpy<U2数组中取出的元素是固定长度的Unicode类型（[unichr x 2]），两者无法匹配。

以下是三种可行的修复方式：

方式一：修改my_func签名适配固定长度Unicode #

针对numpy数组的<U2类型，将my_func的签名改为对应长度的UnicodeType(2)：

import numpy as np
from numba import types, typed, njit


@njit(signature_or_function=types.UnicodeType(2)(types.UnicodeType(2)))
def my_func(symbol):
    return "*" + symbol + "*"


@njit
def run_numba(symbols):
    for symbol in symbols:
        result = my_func(symbol=symbol)
        print(result)

symbols = np.array(["a", "b"], dtype="<U2")
run_numba(symbols)

方式二：移除my_func的显式签名，让Numba自动推导类型 #

删除my_func的固定签名，Numba会根据传入的参数自动识别类型，适配不同长度的Unicode字符串：

import numpy as np
from numba import types, typed, njit


@njit
def my_func(symbol):
    return "*" + symbol + "*"


@njit
def run_numba(symbols):
    for symbol in symbols:
        result = my_func(symbol=symbol)
        print(result)

symbols = np.array(["a", "b"], dtype="<U2")
run_numba(symbols)

方式三：将numpy数组转为Numba typed.List #

如果需要保留my_func的通用签名，可将numpy数组转换为Numba的typed.List，列表内的元素会被识别为通用unicode_type：

import numpy as np
from numba import types, typed, njit


@njit(signature_or_function=types.unicode_type(types.unicode_type))
def my_func(symbol):
    return "*" + symbol + "*"


@njit
def run_numba(symbols):
    for symbol in symbols:
        result = my_func(symbol=symbol)
        print(result)

symbols = np.array(["a", "b"], dtype="<U2")
typed_symbols = typed.List(symbols.tolist())
run_numba(typed_symbols)

内容的提问来源于stack exchange，提问作者GlaceCelery