原地操作

主线位置：训练主线的 forward 一步。前向里改写一个 autograd 还要用的张量，PyTorch 在 backward 时报错；burn 里同样的写法怎么执行、代价在哪。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备，Device::flex().autodiff()）

flowchart TB
  fwd["y = x 乘 w"] --> save["autograd 存下 x"]
  save --> pt["PyTorch x.add_ 1"]
  save --> bn["burn x = x + 1"]
  pt --> ver["版本号 0 变 1"]
  ver --> err["backward 报错"]
  bn --> rc["x 的计数是 2"]
  rc --> newbuf["另开缓冲 16 MiB"]
  newbuf --> ok["梯度用旧的 x"]

实验 crate tensorb-own（burn =0.22.1，feature std、flex、autodiff，--release），程序装了一个只计数的全局分配器，记录一段代码向分配器要了多少字节。PyTorch 侧是 own.py。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch：改了 autograd 要用的张量

y = x * w 对 w 求导要用到 x，autograd 把 x 存下来。之后原地改 x：

x = torch.ones(3)
w = torch.full((3,), 2.0, requires_grad=True)
y = x * w
print("x._version before", x._version)
x.add_(1)
print("x._version after", x._version)
y.sum().backward()

x._version before 0
x._version after 1
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3]] is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True, check_nan=False).

PyTorch 的张量和 autograd 存下的是同一块内存。每次原地操作把张量的版本号加一，backward 时对比存下时的版本号，不一致就报错，否则梯度会用被改过的值算出来。直接改一个 requires_grad 的叶子也不行：

RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.
under no_grad: [2.0, 2.0, 2.0] requires_grad True

第二行是包在 torch.no_grad() 里再 add_，这是 PyTorch 手写参数更新的惯用法。改法是不用原地操作，x = x + 1 新建一个张量，backward 得到 w.grad = [1.0, 1.0, 1.0]。

burn：没有 add_

burn 的 Tensor 没有 add_、mul_ 这类方法，也没有实现 AddAssign（+=）。原地改写一个张量只能写成 x = x + 1.0：旧的 x 按值进入运算，新的值绑回同一个名字。同一个场景：

let device = Device::flex().autodiff();
let n = 1 << 22; // 16 MiB of f32
let x = Tensor::<1>::ones([n], &device);
let w = Tensor::<1>::full([n], 2.0, &device).require_grad();
let y = x.clone() * w.clone();
println!("{mode}: after y = x * w, x.can_mut() {}", x.can_mut());
let (x, mib) = measure(|| x + 1.0);
println!("{mode}: x = x + 1 allocated {mib:.3} MiB, x[0] = {}", vals(x.clone().slice(0..1))[0]);
let grads = y.sum().backward();
let g = w.grad(&grads).unwrap();

再在不带 autodiff 的 Device::flex() 上跑一遍（不调 require_grad，不做 backward）：

saved: after y = x * w, x.can_mut() false
saved: x = x + 1 allocated 16.001 MiB, x[0] = 2
saved: dL/dw[0..3] = [1.0, 1.0, 1.0] (x before the add was 1)
plain: after y = x * w, x.can_mut() true
plain: x = x + 1 allocated 0.000 MiB, x[0] = 2

梯度是 1，用的是加 1 之前的 x，和 PyTorch 改法后的结果一样，没有报错。区别在分配：

• 不带 autodiff 时，x.clone() * w 用完那个 clone 就释放了，x 的缓冲又只有它自己持有。x + 1.0 走 Flex 标量运算的原地分支（scalar_op_typed 开头的 tensor.is_unique()），直接写回原缓冲，没有分配。
• 带 autodiff 时，乘法的反向要用 x，float_mul 在登记节点时 prep.checkpoint(&lhs) 留了一个引用。x 的计数因此是 2，can_mut() 是 false，x + 1.0 另开 16 MiB 写结果，存下的那份不受影响。

PyTorch 用版本号检测「存下的张量被改了」，然后报错；burn 用引用计数避免它发生，代价是多一次分配。在 burn 里不存在「改了 autograd 要用的张量」这种错误，只存在「本可以原地写却多分配了一块」这种性能问题。乘法在什么条件下留下哪个操作数，见 burn 卷一个算子的反向。

什么时候真的原地写

不带 autodiff，在 Flex 上对一个 64 MiB 的张量连做三次 x = x * 2.0 + 1.0，再在留着一个 clone 的情况下做一次：

let mut x = Tensor::<2>::random([n, n], Distribution::Default, &device);
for step in 0..3 {
  let (next, mib) = measure(|| x * 2.0 + 1.0);
  x = next;
  println!("step {step}: x = x * 2 + 1 allocated {mib:.3} MiB");
}
let keep = x.clone();
let (next, mib) = measure(|| x * 2.0 + 1.0);
println!("with a live clone: allocated {mib:.3} MiB");

step 0: x = x * 2 + 1 allocated 0.000 MiB
step 1: x = x * 2 + 1 allocated 0.000 MiB
step 2: x = x * 2 + 1 allocated 0.000 MiB
with a live clone: allocated 64.000 MiB

x * 2.0 + 1.0 是两个运算，每个都按值拿到独占的缓冲，都原地写，一个字节也不分配，效果等于 PyTorch 的 x.mul_(2).add_(1)。留一个 clone 以后，* 2.0 只能另开 64 MiB；它的结果是新缓冲、只有自己持有，后面的 + 1.0 又能原地写，所以总共只多了一块。逐元素二元运算的原地分支在 Flex 里怎么写，见 burn 卷 Tensor 与 Kind 和 Flex 的「逐元素：先想办法不分配」。

想写出 PyTorch 那种「原地」代码，在 burn 里要做的只是：不留多余的 clone。切片赋值同理，a = a.slice_assign(s![.., 1..3], v) 在 a 独占时也是原地写（切片与索引）。

Tensor 上有一个 inplace(&mut self, f)，名字像 add_，其实不是：它从 &mut Tensor 里取出值、交给按值的 f、再放回去，源码注释写着「is not an optimization」。它的用途是在只有 &mut 的地方（比如模型的某个字段）套用按值的运算，能不能原地写仍由上面的引用计数决定。

动脑筋答案

程序的 untracked 模式：设备是 Device::flex().autodiff()，w 不调 require_grad()。

untracked: after y = x * w, x.can_mut() true
untracked: x = x + 1 allocated 0.001 MiB, x[0] = 2

几乎是 0：0.001 MiB 是一点零碎分配，没有 16 MiB 的结果缓冲。float_mul 只在另一个操作数需要梯度时才存下这一个：let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs));。w 不需要梯度，x 就不会被存下，计数回到 1，加法原地写。决定能不能原地写的不是设备带不带 autodiff，而是这一刻有没有别人持有这块缓冲。

小结

• PyTorch 的 add_ 和 autograd 存下的张量共用内存，靠版本号在 backward 时报错；叶子张量的原地更新要包在 no_grad() 里。
• burn 没有 add_，也没有 +=；写 x = x + 1.0，旧值按值进入运算。
• Flex 在缓冲只有一个持有者时原地写：x = x * 2.0 + 1.0 连做三次，分配为 0。
• autodiff 存下的张量计数为 2，同样的写法会另开缓冲（16 MiB 的例子多分配 16 MiB），梯度用的是旧值，不会出错。
• Tensor::inplace 只是在 &mut 上套用按值运算，不是优化。