自动微分

这一部分把求导主线拆开讲。burn 的自动微分是反向模式（reverse mode），只支持一阶导数，实现集中在 burn-autodiff 这一个 crate（约 1.1 万行，其中 ops/tensor.rs 一个文件就有 4830 行，是每个浮点算子的反向）。它是一个后端装饰器 Autodiff<B, C>：套在任何后端外面，前向照常调 B 的算子，同时把「怎么求导」记成图上的一个步骤；backward() 时再把这些步骤倒着执行一遍。

0.22 里用户不再写 Autodiff<B> 这个类型。设备上打一个标记 Device::flex().autodiff()，Dispatch 在运行期看到标记，就把算子交给 Autodiff<Flex, C>；C 是梯度检查点策略，同样由设备上的运行期值换成类型参数。分派本身怎么生成 match 见 Dispatch，本部分只讲 autodiff 那一层。

1. Autodiff 装饰器：Autodiff<B, C> 只替换浮点张量的类型，AutodiffTensor 是「原语 + 节点 + 引用计数」；Requirement 的三种取值决定一个张量是叶子、被追踪的中间结果还是常量；Device::autodiff() 经 DispatchAutodiffContext 和生成的 match 落到 Autodiff<Flex, NoCheckpointing> 或 Autodiff<Flex, BalancedCheckpointing>。
2. 一个算子的反向：以 float_mul、float_matmul 为例走一遍 Backward::prepare 的类型状态链、OpsKind::Tracked / UnTracked、保存的状态（只存节点 id，张量放在检查点里）；广播时梯度怎么求和回原形状；四组算子的梯度和有限差分对上。
3. backward 与 Gradients：步骤存在全局的 GraphLocator 里，按节点 id 找图；backward 先校验祖先、再按深度分桶倒序执行；Gradients 是按节点 id 存的类型擦除容器，非叶子的梯度用完即删；只前向不 backward 时保存的状态会一直留到下一次 backward。
4. 梯度检查点：NoCheckpointing 把每个算子都当成计算密集、保存输出；BalancedCheckpointing 对逐元素算子只记「怎么重算」，backward 时按拓扑序补算；在 Metal 上同一个 8 层网络前向后占用从 92 MiB 降到 32 MiB。

flowchart TB
  dev["Device autodiff"] --> ctx["DispatchAutodiffContext"]
  ctx --> ad["Autodiff B C 算子"]
  ad --> fwd["B 的前向算子"]
  ad --> reg["登记 Step 和检查点"]
  reg --> srv["AutodiffServer"]
  srv --> bw["backward 按深度倒序"]
  bw --> grads["Gradients 按节点 id"]
  grads --> gp["GradientsParams 给优化器"]

前两页在 backward() 之前，管「前向时记下了什么」；后两页在 backward() 里面，管「怎么用记下的东西算梯度、用完怎么释放」。Gradients 怎么按参数 id 重新组织成 GradientsParams 交给优化器，见优化器；不依赖 burn 自己写一个三百行的反向模式自动微分，见 mini-autodiff。