如何解读未传入x和y时torch.where()的输出结果？

AI 小北

2026-8-3

问题

理解PyTorch中torch.where()传入x和y时的输出逻辑，但搞不懂不传入x和y时的输出逻辑——明明输入张量形状固定，输出维度却变了。以下是示例代码：

y = torch.ones(3, 2)
x = torch.randn(3, 2)
print(x)
----------------------------
tensor([[-0.0022,  0.4871],
        [ 0.0788,  0.2937],
        [ 0.1909, -2.1636]])
----------------------------
print(torch.where(x > 0, x, y))
----------------------------
tensor([[1.0000, 0.4871],
        [0.0788, 0.2937],
        [0.1909, 1.0000]])
----------------------------
print(torch.where(x > 0))

(tensor([0, 1, 1, 2]), tensor([1, 0, 1, 0]))

解析

torch.where()的行为分两种核心场景，差别在于是否传入x和y参数：

1. 传入x和y参数时

这种场景下，函数会根据第一个条件张量（示例中是x>0）的布尔值，从x和y中选择对应位置的元素，最终输出和输入张量形状完全一致的结果：

条件为True的位置，取x对应位置的值
条件为False的位置，取y对应位置的值
示例中torch.where(x>0, x, y)的输出是3行2列，和x、y的形状完全匹配，逻辑直观。

2. 只传入条件张量（不传入x和y）时

这种场景下，函数的作用是提取所有满足条件的元素的坐标索引，输出格式为一个元组：

元组的长度等于输入张量的维度数（示例中x是2维，所以元组包含2个张量）
每个张量的长度等于条件为True的元素总数（示例中x里有4个正数，所以每个张量都是长度为4的一维张量）
元组里的张量按维度顺序对应坐标：第一个张量存储所有True元素的行索引，第二个存储列索引，一一配对后就是满足条件元素的准确位置。比如示例输出中，(0,1)、(1,0)、(1,1)、(2,0)这四个坐标，正好是x中大于0的元素所在位置。

简言之，不带x/y时，torch.where()不是生成和输入同形状的张量，而是返回满足条件元素的"坐标清单"，所以输出维度会随满足条件的元素数量变化，和输入张量的固定形状无关。