不能直接在 torch.Tensor 上写反向逻辑，因为 autograd 禁止手动修改 tensor.grad 以避免破坏计算图一致性；必须继承 torch.autograd.Function 并实现 forward 与 backward 静态方法，用 ctx.save_for_backward 保存中间变量，backward 返回梯度元组需与 forward 输入参数一一对应，且须使用原生 torch 运算。
为什么不能直接在 torch.Tensor 上写反向逻辑？
PyTorch 的 autograd 不允许你手动修改梯度张量（比如 tensor.grad）来“伪造”反向传播路径——这会破坏计算图的拓扑一致性，导致 backward() 报错或梯度静默丢失。真正可控的方式是继承 torch.autograd.Function，把前向和反向逻辑封装成一个原子操作。
如何定义一个带自定义反向的 Function？
必须同时实现 forward 和 backward 两个静态方法，且 backward 的输入必须与 forward 的输出数量、顺序严格对应（不是输入参数！）。关键点：
-
forward中用ctx.save_for_backward()保存需要在反向中复用的中间变量（如输入张量、mask、临时计算结果），否则反向里拿不到 -
backward返回的梯度元组，长度必须等于forward输入参数个数，每个元素对应相应输入的梯度（None表示该输入不可导） - 所有张量运算必须用
torch原生函数（如torch.sum），不能混用 NumPy 或原地操作（如.mul_()）
示例：实现一个只对正数求导的 ReLU 变体：
class PositiveOnlyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return torch.where(x > 0, x, torch.zeros_like(x))
<pre class="brush:php;toolbar:false;">@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
grad_input = torch.where(x > 0, grad_output, torch.zeros_like(grad_output))
return grad_input
立即学习“Python免费学习笔记（深入）”；
调用自定义 Function 时为什么报 RuntimeError: Trying to backward through the graph a second time？
这是最常见的误用：直接调用 PositiveOnlyReLU.apply(x) 后又对 x 多次调用 backward()，而没清空计算图。根本原因在于 apply 返回的仍是叶子张量（leaf tensor），但它的梯度累加行为受 retain_graph 控制。正确做法：
- 确保每次
backward()都作用在**标量 loss** 上（非向量），除非显式传retain_graph=True - 如果要多次反向（如二阶梯度），必须在第一次
backward()时加retain_graph=True，否则图被释放 - 避免把
apply()结果赋给全局变量或反复复用——它绑定的是某次前向的上下文，重复使用可能触发缓存冲突
什么时候该用 torch.compile 或 torch.jit.script 包裹自定义 Function？
基本不用。PyTorch 2.x 的 torch.compile 目前不支持自定义 Function 的图融合，会直接 fallback 到 eager 模式；torch.jit.script 对 Function 支持极弱，容易因 ctx 或动态控制流报错。真要部署，更稳妥的做法是把自定义逻辑重构成纯 nn.Module + 标准算子（比如用 torch.nn.functional.relu 加条件掩码），仅在必须绕过标准梯度规则时才用 Function。
最易被忽略的一点：自定义 backward 中的数值稳定性。比如你在反向里写了 1 / x，但没检查 x 是否接近零——这不会在前向报错，却会让训练中途梯度爆炸，且难以定位。
