在 Numba 中无法直接用 isinstance 在 @njit 函数内做运行时类型判断；正确做法是使用 @overload 在编译期根据参数类型签名生成专用实现，从而支持对整数标量和整数数组的不同逻辑分支。
在 numba 中无法直接用 `isinstance` 在 `@njit` 函数内做运行时类型判断；正确做法是使用 `@overload` 在编译期根据参数类型签名生成专用实现，从而支持对整数标量和整数数组的不同逻辑分支。
Numba 的 JIT 编译模型要求所有控制流在编译期可静态推导——这意味着 isinstance(indices, nb.int64[:]) 这类运行时类型检查不被支持（即使语法上不报错，也会触发未定义行为或降级到对象模式，且 Numba 会明确发出警告）。真正符合 Numba 设计哲学的类型分发机制是 @overload：它允许你为同一 Python 函数名注册多个底层实现，并由 Numba 在类型解析阶段自动选择最匹配的编译版本。
以下是推荐的、兼容 Numba ≥0.59 的完整实现方案：
import numba as nb
import numpy as np
# 1. 定义两个独立的底层实现（纯 NumPy 风格，无 JIT 装饰）
def test_dispatch_scalar(X, indices):
ref_pos = np.empty(3, dtype=np.float64)
ref_pos[:] = X[:, indices]
return ref_pos
def test_dispatch_vector(X, indices):
ref_pos = np.empty((3, len(indices)), dtype=np.float64)
ref_pos[:, :] = X[:, indices]
return ref_pos
# 2. 提供一个 Python 层 fallback（仅用于类型推导和错误提示）
def test_dispatch_impl(X, indices):
if isinstance(indices, (int, np.integer)):
return test_dispatch_scalar(X, indices)
elif (isinstance(indices, np.ndarray) and
indices.ndim == 1 and
np.issubdtype(indices.dtype, np.integer)):
return test_dispatch_vector(X, indices)
else:
raise TypeError(f"Unsupported type for 'indices': {type(indices)}")
# 3. 使用 @overload 注册类型特化逻辑（关键！）
@nb.extending.overload(test_dispatch_impl)
def test_dispatch_impl_overload(X, indices):
# 编译期类型检查：indices 是整数标量？
if isinstance(indices, nb.types.Integer):
return test_dispatch_scalar
# indices 是一维整数数组？
elif (isinstance(indices, nb.types.Array) and
indices.ndim == 1 and
isinstance(indices.dtype, nb.types.Integer)):
return test_dispatch_vector
else:
# 编译期报错，提升可调试性
raise TypeError("Only scalar integers or 1D integer arrays are supported.")
# 4. 用户调用入口：普通 njit 函数，内部委托给 overload 分发
@nb.njit
def test_dispatch(X, indices):
return test_dispatch_impl(X, indices)
✅ 使用示例：
X = np.random.rand(3, 100) # 标量索引 → 返回 shape=(3,) 的向量 result1 = test_dispatch(X, 42) # ✅ 编译为 test_dispatch_scalar # 数组索引 → 返回 shape=(3, N) 的矩阵 idxs = np.array([10, 20, 30], dtype=np.int64) result2 = test_dispatch(X, idxs) # ✅ 编译为 test_dispatch_vector
⚠️ 重要注意事项：
- 不要在 @njit 函数中使用 isinstance(..., nb.types.X) —— 这些类型对象仅存在于编译期，运行时不可访问；
- @overload 的函数体必须是纯 Python（不能含 @njit），其返回值应为另一个可被 Numba 编译的函数（如 test_dispatch_scalar）；
- 所有类型判断必须基于 nb.types.*（如 nb.types.Integer, nb.types.Array），而非运行时 Python 类型；
- 若需支持更多类型（如 int32、uint64），应优先使用泛型类型（如 nb.types.Integer）而非硬编码 nb.types.int64，以提高兼容性；
- @overload 会为每种输入类型组合生成独立机器码，因此零开销、高性能，且类型安全。
该方案完全替代了已废弃的 @generated_jit，是当前 Numba 官方推荐的多态函数构造方式。
