Python类变量如何指定np.ndarray/torch类型？解决类型联合报错

错误原因拆解 #

• 用np.ndarray | torch.float报错：torch.float是**torch.dtype的实例**，不是类型对象，而np.ndarray是正经的类型，|运算符要求两边都是类型（或兼容的类型别名），类型不匹配直接触发错误。
• 用np.ndarray | torch.tensor是写法错误：torch.tensor是创建张量的工厂函数，正确的张量类型应该是torch.Tensor（首字母大写）。

正确的多类型标注方案 #

1. 基础需求：numpy数组或任意torch张量 #

如果只需要标注变量可以是np.ndarray或torch.Tensor，直接用类型联合即可：

import numpy as np
import torch
from typing import Union

class MyClass:
    # Python 3.10+支持|语法，更简洁
    data: np.ndarray | torch.Tensor
    # 兼容Python 3.9及以下版本的写法
    # data: Union[np.ndarray, torch.Tensor]

2. 进阶需求：限定torch张量为float dtype #

如果要明确张量必须是float类型，不能直接把torch.float和类型联合，得用Annotated附加元数据，配合类型检查工具（如pyright、mypy）实现约束：

import numpy as np
import torch
from typing import Annotated, Union

# 定义带dtype约束的张量类型别名
FloatTensor = Annotated[torch.Tensor, torch.float]

class MyClass:
    # 变量可以是np.ndarray，或者dtype为float的torch.Tensor
    data: np.ndarray | FloatTensor

提示：这种带约束的标注需要类型检查工具支持，pyright默认兼容Annotated的元数据解析，mypy需要安装mypy-extensions插件才能识别这类约束。

3. 运行时类型校验（可选） #

如果需要在代码运行时验证类型，不能只靠静态标注，得手动写判断逻辑：

class MyClass:
    data: np.ndarray | FloatTensor

    def validate_data(self):
        if isinstance(self.data, np.ndarray):
            return
        elif isinstance(self.data, torch.Tensor):
            if self.data.dtype != torch.float:
                raise TypeError("张量必须是float类型")
        else:
            raise TypeError("数据必须是np.ndarray或float类型的torch张量")

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