使用Numba JIT合并Numpy数组遭遇TypingError问题求助

问题重现 #

以下函数尝试用Numba的@nb.njit装饰器加速两个NumPy数组的合并：

import numpy as np
import numba as nb

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.array([a, b])

• 传入单个浮点数时运行正常：

  >>> combine(1., 2.)
  array([1., 2.])
  

• 传入一维数组时触发TypingError，核心报错信息：

  TypingError: array(float64, 1d, C) not allowed in a homogeneous sequence

错误原因 #

Numba的nopython模式下，np.array([a, b])要求输入的序列是同质标量序列，而不是包含数组的序列。当传入两个一维数组时，[a, b]是数组的列表，不符合Numba对np.array输入的类型要求，因此找不到匹配的实现。

解决方法 #

方法1：使用np.vstack直接垂直堆叠数组 #

import numpy as np
import numba as nb

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.vstack((a, b))

测试运行：

>>> combine(np.array([1., 2.]), np.array([3., 4.]))
array([[1., 2.],
       [3., 4.]])

方法2：使用np.stack指定轴堆叠 #

如果需要自定义堆叠的轴方向，可使用np.stack：

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    return np.stack((a, b), axis=0)  # axis=0等价于vstack，axis=1会按列堆叠

方法3：手动分配内存并填充 #

适合需要更精细控制数组构建的场景：

@nb.njit
def combine(a: nb.float64[:], b: nb.float64[:]):
    # 确保输入数组长度一致，若不确定可添加长度检查
    result = np.empty((2, len(a)), dtype=np.float64)
    result[0] = a
    result[1] = b
    return result

以上三种方法均能在Numba的nopython模式下正常编译并运行，避免了原写法中数组序列的类型问题。

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