如何在PyTorch中基于列表值定义掩码函数 解决多值张量布尔歧义报错

你遇到的RuntimeError: Boolean value of Tensor with more than one value is ambiguous报错，核心原因如下：

• 原生Python的in运算符不支持PyTorch张量的逐元素匹配判断：当输入x是多元素张量时，x in self.prompt_ids会尝试将整个张量的匹配结果聚合为单个布尔值，PyTorch不允许直接对多布尔值的张量做隐式布尔转换，因此抛出歧义错误。
• 你原有范围判断分支用的是逐元素逻辑运算符&，返回的是和输入x同形状的布尔张量，符合掩码操作的要求，因此可以正常运行。

直接替换列表匹配逻辑为PyTorch内置的逐元素匹配接口torch.isin即可，修改后的代码如下：

import torch

def get_prompt_token_fn(self):
    if self.prompt_ids:
        return lambda x: torch.isin(
            x, 
            torch.tensor(self.prompt_ids, device=x.device, dtype=x.dtype)
        )
    else:
        return lambda x: (x>=self.id_offset) & (x<self.id_offset+self.length)

优化建议 #

如果self.prompt_ids是初始化后固定不变的参数，可以在类的__init__方法中提前将其转换为对应设备的张量，避免每次调用匹配函数时重复生成张量，提升运行效率：

# 类初始化时新增逻辑
self.prompt_ids_tensor = torch.tensor(self.prompt_ids, device=self.device, dtype=torch.long)

# 修改判断分支
def get_prompt_token_fn(self):
    if self.prompt_ids:
        return lambda x: torch.isin(x, self.prompt_ids_tensor)
    else:
        return lambda x: (x>=self.id_offset) & (x<self.id_offset+self.length)

低版本PyTorch兼容方案 #

如果你使用的PyTorch版本低于1.10（该版本才新增torch.isin接口），可以用广播机制手动实现逐元素匹配：

def get_prompt_token_fn(self):
    if self.prompt_ids:
        return lambda x: (x.unsqueeze(-1) == torch.tensor(self.prompt_ids, device=x.device, dtype=x.dtype)).any(dim=-1)
    else:
        return lambda x: (x>=self.id_offset) & (x<self.id_offset+self.length)

以上所有方案返回的都是和输入同形状的布尔张量，和原有范围判断分支的输出格式完全统一，不需要修改后续的掩码调用逻辑。

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