Numpy运算中数组类型保留问题求助

问题根源 #

NumPy的类型存根（用于MyPy等类型检查工具）对自定义Annotated包装的泛型np.ndarray类型支持不足。np.equal这类通用函数的类型提示无法从你的自定义标注类型中正确推导输出数组的形状和布尔类型，导致结果被识别为Any。

可行解决办法 #

1. 更新依赖到最新版

升级NumPy和MyPy，新版对泛型数组类型的推导支持更完善：

pip install --upgrade numpy mypy

2. 使用官方推荐的类型提示

改用numpy.typing.NDArray替代自定义Annotated类型（若不需要额外元数据）：

from numpy.typing import NDArray
import numpy as np

BatchNactionsNpFloatType = NDArray[np.float32]

x: BatchNactionsNpFloatType = np.array([[1.0,2.0],[1.0,3.0],[1.0,2.0],[1.0,2.0]], dtype=np.float32)
y: BatchNactionsNpFloatType = np.array([[1.0,2.0],[1.0,2.0],[3.0,2.0],[1.0,2.0]], dtype=np.float32)

res: NDArray[np.bool_] = np.equal(x, y)

类型检查器能准确识别这种标准类型，自动推导np.equal的输出为布尔数组。

3. 显式指定结果类型

若必须保留Annotated类型，可通过类型注释或断言明确指定结果类型：

# 直接添加类型注释
res = np.equal(x, y)  # type: np.ndarray[tuple[int,int], np.bool_]

# 或使用cast函数
from typing import cast
res = cast(np.ndarray[tuple[int,int], np.bool_], np.equal(x, y))

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