PyTorch 如何用 torch.where 比较三个张量并生成布尔掩码？

PyTorch 张量不兼容 Python 原生的链式比较规则。直接书写 A == B == C 会被解析为 (A == B) == C：先计算A、B逐元素相等的布尔张量，再将该布尔张量与C做比较，不符合「三个张量同位置值完全一致」的预期逻辑。

通用写法（兼容所有可比较数据类型） #

将连续判等拆分为两个独立比较，通过逐位与&组合条件，直接输出布尔掩码：

import torch

# 测试输入
A = torch.tensor([True, True, False, False])
B = torch.tensor([True, False, False, True])
C = torch.tensor([True, False, False, False])

mask = (A == B) & (B == C)
print(mask)
# 输出: tensor([ True, False,  True, False])

如果需要输出0/1格式的整数张量，直接做类型转换即可，不需要额外调用torch.where：

mask_int = ((A == B) & (B == C)).long()
# 也可通过 .to(torch.int32) 指定具体整数类型
print(mask_int)
# 输出: tensor([1, 0, 1, 0])

如果需要沿用torch.where写法，传入拆分后的组合条件即可：

mask_int = torch.where((A == B) & (B == C), 1, 0)

布尔张量专属简化写法 #

由于布尔值只有True/False两种，三者完全相等仅包含「全为True」「全为False」两种情况，纯布尔输入场景下也可使用：

mask = (A & B & C) | (~A & ~B & ~C)

该写法仅适配布尔类型输入，可读性不如通用判等方案，不推荐通用场景使用。

• 必须为每个相等判断添加括号：Python中&运算符优先级高于==，不加括号会被解析为A == (B & B) == C，逻辑完全错误。
• 逐元素逻辑与必须使用&，不能使用Python原生and关键字：and会尝试将整个张量转为单个布尔值，会直接触发运行错误。

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