如何用Pythonic方式检查PyTorch张量的数据类型（如ComplexFloat）？

方法1：直接对比dtype属性 #

最直观的方式是将张量的dtype属性与目标类型直接比较：

import torch

# 示例张量
tensor = torch.tensor([1+2j, 3+4j], dtype=torch.complex64)

if tensor.dtype == torch.complex64:  # 对应ComplexFloat64类型
    print("张量是ComplexFloat64类型")
elif tensor.dtype == torch.complex32:  # 对应ComplexFloat32类型
    print("张量是ComplexFloat32类型")

方法2：用isinstance检查dtype对象 #

PyTorch的dtype本身是类实例，因此可以直接对tensor.dtype使用isinstance，贴合你想要的Pythonic写法：

if isinstance(tensor.dtype, torch.complex64):
    # 处理ComplexFloat64类型的逻辑
elif isinstance(tensor.dtype, torch.complex128):
    # 处理ComplexFloat128类型的逻辑

方法3：封装通用判断函数 #

如果需要复用判断逻辑，可以封装一个简单函数：

def matches_dtype(tensor, target_dtype):
    return isinstance(tensor.dtype, target_dtype)

# 使用示例
if matches_dtype(tensor, torch.complex64):
    print("张量匹配目标数据类型")

注意：PyTorch中的ComplexFloat系列类型对应torch.complex32（ComplexFloat32）、torch.complex64（ComplexFloat64）、torch.complex128（ComplexFloat128），可根据精度需求选择对应常量。

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