使用二维数组索引二维数组时触发Numba TypingError问题

问题原因 #

Numba的njit模式对Numpy的花式索引支持有限，不允许直接用二维整数数组作为索引去访问另一个二维数组，这就是你遇到TypingError的核心原因。

可行的规避方法 #

方法1：展平索引数组后再重塑形状

将二维索引数组展平为一维，完成索引后再把结果重塑为原索引数组的形状（加上原数组的剩余维度），完全匹配Numpy的索引行为：

from numba import njit
import numpy as np

@njit
def index_with_flatten(arr, idx):
    # 展平索引数组
    idx_flat = idx.ravel()
    # 按一维索引取元素，再重塑为目标形状
    result = arr[idx_flat].reshape(idx.shape + arr.shape[1:])
    return result

示例验证：

arr = np.random.rand(3, 2)  # 原二维数组
idx = np.array([[0, 1], [1, 2]], dtype=np.int32)  # 二维索引数组
print(index_with_flatten(arr, idx).shape)  # 输出 (2, 2, 2)，和Numpy行为一致

方法2：显式循环逐个索引（修正版）

你之前的循环报错大概率是因为写法不符合Numba要求，正确的循环需要逐个访问索引数组的元素，并用单个整数索引原数组：

@njit
def index_with_loop(arr, idx):
    # 初始化结果数组，形状匹配Numpy的索引输出
    result_shape = idx.shape + arr.shape[1:]
    result = np.empty(result_shape, dtype=arr.dtype)
    
    # 遍历索引数组的每个位置
    for i in range(idx.shape[0]):
        for j in range(idx.shape[1]):
            result[i, j] = arr[idx[i, j]]
    return result

方法3：针对双坐标索引的特殊处理

如果你的二维索引数组是每行存储一个(i,j)坐标（比如形状为(k,2)），可以把索引拆分为行、列两个一维数组，Numba支持这种索引方式：

@njit
def index_ij_coords(arr, idx):
    rows = idx[:, 0]
    cols = idx[:, 1]
    # 直接用两个一维数组索引，Numba完全支持
    return arr[rows, cols]

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