Numba嵌套列表推导编译失败求助：数组维度重排报错

问题场景 #

需要将形状为(a, b, c, c)的数组重排为(b, a, c, c)，即生成包含b个按a维度顺序堆叠的新数组，使用Numba 0.55.1编写的代码如下：

@njit(locals={'out': float64[:, :, :, :], 'write_list': float64[:, :, :, :]})
def prepare_write(prep: ndarray):

    bands, years, width, height = prep.shape

    # 根据文档，嵌套列表推导是支持的...

    out = np.array([[prep[b][i] for b in range(bands)] for i in range(years)])
    # for i in range(years):
    #     out[i] = np.array([write_list[b][i] for b in range(bands)])
    for v in out:
        vote(v)
    return out

其中vote是已jit编译的原地修改数组函数，无额外问题。

报错信息 #

编译时触发如下错误：

>>> setitem(array(undefined, 1d, C), int64, array(float64, 2d, C))
 
There are 16 candidate implementations:
  - Of which 16 did not match due to:
  Overload of function 'setitem': File: <numerous>: Line N/A.
    With argument(s): '(array(undefined, 1d, C), int64, array(float64, 2d, C))':
   No match.

报错原因 #

Numba 0.55.1对嵌套列表推导转换为多维numpy数组的类型推导支持存在局限：

• 嵌套列表推导生成的是二维数组列表，在转换为4维numpy数组时，Numba无法正确推导中间步骤的数组类型与最终目标float64[:, :, :, :]的匹配关系
• 内部的setitem操作试图将2维数组赋值到类型未明确的1维数组容器中，导致找不到匹配的重载实现

解决方法 #

方法1：使用numpy transpose直接转置维度（最优解） #

利用numpy原生的维度转置操作，简洁高效且Numba支持良好：

@njit(locals={'out': float64[:, :, :, :]})
def prepare_write(prep: ndarray):
    bands, years, width, height = prep.shape
    # 调整维度顺序：原(bands, years, width, height) → (years, bands, width, height)
    out = prep.transpose(1, 0, 2, 3)
    for v in out:
        vote(v)
    return out

方法2：手动初始化数组并循环填充 #

明确初始化目标形状的数组，避免列表推导的类型推导问题：

@njit(locals={'out': float64[:, :, :, :]})
def prepare_write(prep: ndarray):
    bands, years, width, height = prep.shape
    # 预先初始化指定形状和类型的空数组
    out = np.empty((years, bands, width, height), dtype=np.float64)
    # 循环填充对应切片
    for i in range(years):
        for b in range(bands):
            out[i, b] = prep[b, i]
    for v in out:
        vote(v)
    return out

方法3：升级Numba版本（可选） #

Numba后续版本（如0.57及以上）对列表推导和多维数组类型推导的支持有优化，升级后可能直接兼容原列表推导写法，但前两种方法在效率和兼容性上更可靠。

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