PyTorch中是否有theano.tensor.switch的替代方案？需支持梯度保留置零

当然有啦！PyTorch里刚好有和Theano的theano.tensor.switch功能几乎完全匹配的API，而且完美支持你需要的非零元素梯度传播需求。下面给你两种靠谱的解决方案：

方法1：直接用torch.where()（对应Theano switch的官方替代） #

torch.where()就是PyTorch里实现条件选择的核心函数，用法和Theano的switch几乎一模一样。针对你的需求，代码可以这么写：

import torch

# 示例张量，开启梯度追踪
A = torch.tensor([0.05, 0.15, 0.08, 0.2], requires_grad=True)
threshold = 0.1

# 对应Theano的switch逻辑：满足条件(A < 0.1)时取0，否则取A本身
B = torch.where(A < threshold, torch.tensor(0.0, device=A.device), A)

这个操作是完全可微分的：当元素大于等于阈值时，梯度会直接传递给原张量A的对应位置；当元素低于阈值时，因为我们传入的是常数0，这部分的梯度会被置为0，完全符合你的要求。

方法2：布尔掩码实现（更直观的写法） #

如果你觉得torch.where()不够直观，也可以用布尔掩码来实现相同的效果，代码更简洁：

mask = A >= threshold  # 生成布尔掩码，标记需要保留的元素
B = A * mask.to(A.dtype)  # 掩码转为和A相同的 dtype 后相乘，低于阈值的元素会被置零

同样，这个操作也支持梯度传播：掩码的计算不会阻断梯度，相乘后只有被保留的非零元素会把梯度回传到A，完全满足你的需求。

验证梯度传播效果 #

你可以通过简单的反向传播来验证梯度是否正常工作：

# 对B求和后反向传播
B.sum().backward()
# 打印A的梯度，会看到低于0.1的元素梯度为0，其余为1
print(A.grad)  # 输出: tensor([0., 1., 0., 1.])

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