Numba模板编程：如何避免1~3维数组的重复代码编写？

问题背景 #

Numba对运行时可变维度数组的支持有限，当处理维度仅为固定几种（如1/2/3维）的数组时，手动为每个维度编写重复逻辑的维护成本很高。纯Python下通过构造索引元组实现的通用写法，因Numba不支持传入可迭代对象调用tuple构造器无法直接使用。

原有手动分维度实现示例：

import numba
import numpy as np

def pick1D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    return arr[i]

def pick2D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    j=np.random.randint(0,arr.shape[1])
    return arr[i,j]

def pick3D(arr:np.ndarray):
    i=np.random.randint(0,arr.shape[0])
    j=np.random.randint(0,arr.shape[1])
    k=np.random.randint(0,arr.shape[2])
    return arr[i,j,k]

@numba.generated_jit
def ndpick(arr:np.ndarray):
    if arr.ndim==1:
        return pick1D
    elif arr.ndim==2:
        return pick2D
    elif arr.ndim==3:
        return pick3D

纯Python通用写法（Numba下无法运行）：

def pickND(arr):
    index = tuple(np.random.randint(0,arr.shape))
    return arr[index]

实现方案 #

利用generated_jit在编译阶段执行的特性，根据传入数组的维度动态生成对应实现代码，无需手动编写多份重复逻辑，编译后性能和手写单维度函数完全一致：

import numba
import numpy as np

@numba.generated_jit(nopython=True)
def ndpick(arr):
    # 编译阶段获取当前传入数组的维度
    dim_count = arr.ndim
    # 动态拼接对应维度的索引逻辑代码
    index_code = ",".join(
        [f"np.random.randint(0, arr.shape[{idx}])" for idx in range(dim_count)]
    )
    # 构造完整实现函数
    func_def = f"""
def _impl(arr):
    return arr[{index_code}]
"""
    # 加载生成的函数
    local_ns = {}
    exec(func_def, {"np": np}, local_ns)
    return local_ns["_impl"]

方案说明 #

• 代码生成逻辑仅在首次遇到某一维度的数组类型时执行一次，后续相同维度的调用会直接复用预编译的机器码，无运行时额外开销
• 全程运行在nopython模式下，性能和手写的1D/2D/3D专用函数无差异
• 无需手动维护多份重复代码，后续如需支持更高维度数组无需修改核心逻辑
• 这种实现思路和C++模板的编译期多态逻辑一致，都是在编译阶段根据类型/维度参数生成对应专用实现

测试验证 #

# 1维测试
arr1 = np.arange(10)
print(ndpick(arr1))
# 2维测试
arr2 = np.arange(12).reshape(3,4)
print(ndpick(arr2))
# 3维测试
arr3 = np.arange(24).reshape(2,3,4)
print(ndpick(arr3))

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