如何使用torch.where多条件替换张量值为指定ID序列？

首先明确核心要求：torch.where要求条件张量、x、y三者的形状必须广播兼容。你的问题出在y的形状和原张量不匹配，下面分场景给出解决方案：

场景1：替换不满足条件的元素为1-10 ID的对应值 #

假设你的2dtensor是(N, 10)形状的张量（N是行数），1-10的ID张量是(10,)的一维张量，操作步骤如下：

1. 构造多条件时，将行级条件扩展为可广播到整个张量的形状：

   # 构造行级条件（形状为(N,)）
   condition = (2dtensor[:,1] < 4.2) & (2dtensor[:,1] > 3.8) & (2dtensor[:,0] < 3.6)
   # 扩展为(N,1)，这样能广播到(N,10)的张量维度
   condition = condition.unsqueeze(1)
   

2. 创建1-10的ID张量，直接传入torch.where（PyTorch会自动广播形状）：

   import torch
   
   # 生成1-10的ID张量
   id_tensor = torch.arange(1, 11)
   # 执行替换
   result = torch.where(condition, 2dtensor, id_tensor)
   

场景2：替换不满足条件的整行为1-10 ID张量 #

如果需要把不满足条件的整行全部替换为ID张量的对应值，上面的代码完全适用——因为扩展后的条件会对整行的所有元素应用同一个判断。

失败原因分析 #

• 直接传标量1：会被广播为和2dtensor相同的形状，但无法实现1-10的ID替换；若2dtensor是(N,2)形状，传10长度张量则因形状不兼容无法广播。
• 直接传10长度张量：若2dtensor是(N,2)形状，(10,)和(N,2)无法广播，导致报错。此时需确认你的张量维度是否符合预期，或调整ID张量的形状（比如将ID张量转为(1,10)，同时调整原张量形状，或重新定义问题需求）。

额外检查项 #

• 用.shape属性查看2dtensor、condition、id_tensor的形状，确保三者广播后维度一致。
• 多条件组合时，确保每个条件的形状一致，避免因维度不匹配导致逻辑错误。

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