4 · 自动微分：Tape、梯度函数与反向路径

梯度不是“重新跑一遍 forward”

GradientTape 的核心职责是保存可用于反向传播的最小信息：被 watch 的源、执行过的 op、输入输出 Tensor 的标识、以及每个 op 对应的 gradient function。求导时，tape 从目标反向遍历记录，按拓扑依赖累加每个源的梯度。

watch(x)
  ↓
forward: y = matmul(x, w) → loss
  ↓ record op + ids + gradient recipe
gradient(loss, x)
  ↓ reverse traversal
MatMul gradient → accumulate → x.grad

Tape 的两层边界

Python tensorflow/python/eager/backprop.py 负责用户语义：persistent/non-persistent、watch_accessed_variables、gradient / jacobian / batch_jacobian 的输入输出结构。实际的 op 记录和反向遍历由低层 tape 实现，Python 通过 pywrap_tfe 交互。

这种分层让 GradientTape 可以与 eager Tensor 紧密结合，同时让 tf.function 内的梯度被转换成图中的 forward/backward 函数，而不是强迫 Python 在运行时逐个反向调用。

梯度注册表

每个可微 op 都需要一份梯度规则，通常在 Python ops / math_grad 等模块中注册。规则接收 forward 的 inputs、attrs、outputs 和 upstream gradients，返回对输入的 gradients。没有规则、输入不可微或中间值不再可用时，结果可能是 None 或报错。

从源码阅读角度，梯度系统是“算子定义”的第二面：

REGISTER_OP / op wrapper  → forward signature
gradient registration     → reverse signature
kernel implementation     → forward device code

非持久化与持久化

非持久化 tape 在一次 gradient() 后释放记录，降低内存；persistent=True 保留记录，可多次求导，但必须显式释放引用。高阶梯度需要把一次求导本身放在另一个 tape 中，因此第一阶梯度的构造过程也必须被记录。

tf.function 中的梯度

在 tracing 阶段，GradientTape 看到的是符号 Tensor。concrete_function.py 里有 first-order/higher-order tape gradient functions 的封装：forward 图、backward 图、JVP/tangent 相关逻辑会被组合成可调用函数，并通过 FunctionDef / Graph 库注册。

这解释了一个常见现象：eager 下的梯度问题常表现为 Python 对象或 tape 生命周期错误；tf.function 下则可能表现为 graph captures、签名、控制流或未注册 gradient 的图构造错误。

资源变量与副作用

变量更新不是普通 Tensor 的纯函数。梯度逻辑需要区分读取变量的 value、写入资源的 assign，以及控制依赖。tf.function、Keras optimizer 和分布式 strategy 都会围绕资源 handle 建立额外的顺序约束。

源码锚点

• tensorflow/python/eager/backprop.py：GradientTape 公共语义与输入输出结构。
• tensorflow/python/eager/pywrap_tfe_src.cc：Python 与 C API tape 边界。
• tensorflow/python/ops/math_grad.py、array_grad.py：典型梯度注册规则。
• tensorflow/python/eager/polymorphic_function/concrete_function.py：forward/backward function 组合。
• tensorflow/core/graph/while_context.h：控制流图中的循环梯度上下文。

一个适合源码阅读的实验

x = tf.Variable(3.0)
with tf.GradientTape(persistent=True) as tape:
  y = x * x
  z = y * x
print(tape.gradient(y, x))  # 2x
print(tape.gradient(z, x))  # 3x²

第一次结果之后仍能求第二次，是因为 persistent tape 没有释放 forward 记录；把它改成默认 tape，会看到生命周期约束。