mask_where 照搬 torch.where，结果正好互补
主线位置：张量主线的「计算」一步：按布尔掩码在两个张量之间逐元素选值。
版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）
flowchart TB tw["torch.where c, x, y"] --> sem1["c 真取 x 否则 y"] mw["y.mask_where c, x"] --> sem2["在 y 上把 c 真换成 x"] sem1 --> same["两者相同"] sem2 --> same bad["x.mask_where c, y"] --> inv["结果互补"]
实验 crate pita-t 的 p08，PyTorch 侧 p08_where.py。debug 构建。平台 M2 Max、macOS 27、rustc 1.99.0。
PyTorch 的写法
自己写一个 leaky ReLU（或者任何「正的一半这样算、负的一半那样算」的分段函数）：
x = torch.tensor([-2.0, -0.5, 0.0, 1.0, 3.0])
leaky = torch.where(x > 0, x, 0.01 * x)
torch.where(x > 0, x, 0.01*x) [-0.019999999552965164, -0.004999999888241291, 0.0, 1.0, 3.0]
F.leaky_relu [-0.019999999552965164, -0.004999999888241291, 0.0, 1.0, 3.0]
torch.where(cond, a, b)：条件、条件为真时的值、条件为假时的值。
搬到 burn 的现象
burn 里和 torch.where 对应的是张量方法 mask_where(mask, value)。按参数出现的顺序翻译，x 放到点号前面：
let pos = x.clone().greater_elem(0.0);
let wrong = x.clone().mask_where(pos.clone(), x.clone() * 0.01);
x.mask_where(x > 0, 0.01*x) [-2.0, -0.5, 0.0, 0.01, 0.03]
正数被缩小了 100 倍，负数原样通过，正好是想要的结果的补集。形状、dtype 都对，不会报错，训练也能跑，只是激活函数反了。
原因
mask_where 的文档是「Update the given tensor with the value tensor where the mask is true」：
pub fn mask_where(self, mask: Tensor<D, Bool>, value: Self) -> Self {
Self::new(K::mask_where(
self.primitive,
mask.primitive,
value.primitive,
))
}
它是一次「改写」：从 self 出发，掩码为真的位置换成 value。所以 self 扮演的是 torch.where 的第三个参数（条件为假时的值），value 是第二个参数。逐字翻译把两者对调了。它的标量版本 mask_fill 也是同样的读法（a.mask_fill(mask, v) 对应 a.masked_fill(mask, v)，这一个和 PyTorch 的方法同序）。burn 0.22.1 里没有 torch.where 那样「条件在前、两个分支并列」的自由函数，origin/main（2026-10-28）上也没有。各种掩码操作的对照表见切片与索引的最后一节。
改法
接收者是条件为假时的值：
// torch.where(c, a, b) == b.mask_where(c, a)
let right = (x.clone() * 0.01).mask_where(pos.clone(), x.clone());
// 或者保留 x 在前面，把条件取反
let flipped = x.clone().mask_where(pos.bool_not(), x.clone() * 0.01);
println!("activation::leaky_relu {:?}", fv(activation::leaky_relu(x, 0.01)));
(0.01*x).mask_where(x > 0, x) [-0.02, -0.005, 0.0, 1.0, 3.0]
x.mask_where(!(x > 0), 0.01*x) [-0.02, -0.005, 0.0, 1.0, 3.0]
activation::leaky_relu [-0.02, -0.005, 0.0, 1.0, 3.0]
两边各把 torch.where 的结果写成 JSON（burn 侧先把 f32 转成 f64 再打印，避免最短表示和完整位数的差别），compare.py 容差 0：
leaky n= 5 max|diff|=0.00e+00
worst=0.00e+00 tol=0e+00 OK
误差为 0：两边做的是同样的 f32 乘法 0.01 * x，再逐元素选值，没有求和顺序的差别。反向验证：把写反的结果写进 JSON，同一条命令报 leaky ... max|diff|=2.97e+00、MISMATCH，退出码 1（3.0 对 0.03）。
能用现成激活函数的地方直接用：burn 的 activation::leaky_relu、activation::relu 等不需要自己写掩码。
怎么检查
搬代码时把
torch.where(c, a, b)机械地改写成b.mask_where(c, a)，不要按「哪个张量是主角」去放接收者。单元测试的输入要同时有满足和不满足条件的元素，并且两个分支的值要能区分（上面
x和0.01 * x在x = 0处相同，只看 0 是看不出来的）。如果代码里有很多处，写一个和
torch.where同序的辅助函数，搬代码时只改函数名：/// torch.where(c, a, b), same argument order fn where_<const D: usize>(c: Tensor<D, Bool>, a: Tensor<D>, b: Tensor<D>) -> Tensor<D> { b.mask_where(c, a) }where_(x > 0, x, 0.01*x) [-0.02, -0.005, 0.0, 1.0, 3.0]
小结
torch.where(c, a, b)在 burn 里是b.mask_where(c, a)：接收者是条件为假时的值。- 照参数顺序写成
a.mask_where(c, b)结果正好互补，不报错。 - 改对之后和 PyTorch 逐位相同（误差 0）；测试输入要两个分支都覆盖。