使用Numba加速复数NumPy数组运算时遭遇LoweringError问题

使用Numba加速复数数组计算时触发错误：

numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)

复现代码如下：

import numpy as np
from numba import jit
from numpy import array

@jit(nopython=True)
def func(x):
   a = 1j
   v = x*array([[1.,a],
                [2.,3.]])
   return v
func_vec = np.vectorize(func)

print(func_vec(10.))

已知条件：

• 当a为实数时程序运行正常
• 尝试为v指定dtype=np.complex128后问题仍未解决

环境信息：

• Numba版本：0.51.0
• NumPy版本：1.22.3
• Python版本：3.8.10
• 系统：Ubuntu 20.04

问题根源 #

Numba 0.51.0的nopython模式在函数内部动态创建混合实数与复数的numpy数组时，类型推断逻辑存在缺陷，导致后端编译失败。此外，np.vectorize并非真正的向量化实现，与Numba的jit结合会额外引入兼容性问题。

修复方案 #

方案1：提前定义复数数组（推荐） #

将包含复数的数组移到jit函数外部，作为全局常量使用，避免函数内的动态类型推断：

import numpy as np
from numba import jit

# 预先定义复数数组并指定dtype
mat = np.array([[1., 1j],
                [2., 3.]], dtype=np.complex128)

@jit(nopython=True)
def func(x):
    v = x * mat
    return v

# 直接调用，无需np.vectorize
print(func(10.))

方案2：函数内明确指定数组dtype #

在jit函数内部创建数组时，显式指定dtype=np.complex128，强制数组类型为复数：

import numpy as np
from numba import jit

@jit(nopython=True)
def func(x):
    mat = np.array([[1., 1j],
                    [2., 3.]], dtype=np.complex128)
    v = x * mat
    return v

print(func(10.))

额外优化建议 #

• 放弃使用np.vectorize：Numba的jit函数原生支持向量化输入，直接传入数组即可批量计算：

  print(func(np.array([10., 20., 30.])))
  

• 升级Numba版本：新版本（如0.57及以上）对复数类型的类型推断逻辑做了优化，能从根源减少这类问题。

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