写一个迷你自动微分

[dependencies]

# 只有对照测试 tests/burn_compare.rs 用 burn；库本身没有依赖
[dev-dependencies]
burn = { version = "=0.22.1", features = ["flex", "autodiff"] }

/// 行优先存放的 1 维或 2 维 f32 张量。
#[derive(Debug, Clone, PartialEq)]
pub struct Tensor {
  data: Vec<f32>,
  shape: Vec<usize>,
}

  /// 同形状相加；或者 `[m, n] + [n]`，把右边加到每一行上。
  pub fn add(&self, other: &Self) -> Result<Self> {
    match (self.shape.as_slice(), other.shape.as_slice()) {
      ([_, n], [k]) if n == k => {
        let data = self.data.chunks(*n).flat_map(|row| row.iter().zip(&other.data).map(|(a, b)| a + b)).collect();
        Ok(Self { data, shape: self.shape.clone() })
      }
      _ => self.zip(other, "add", |a, b| a + b),
    }
  }

  /// 把梯度求和回 `shape`：`add` 广播的反向。
  pub fn sum_to(&self, shape: &[usize]) -> Result<Self> {
    match (self.shape.as_slice(), shape) {
      (a, b) if a == b => Ok(self.clone()),
      ([_, n], [k]) if n == k => {
        let mut data = vec![0.0; *n];
        for row in self.data.chunks(*n) {
          data.iter_mut().zip(row).for_each(|(acc, x)| *acc += x);
        }
        Ok(Self { data, shape: shape.to_vec() })
      }
      (from, to) => err(format!("sum_to: cannot reduce {from:?} to {to:?}")),
    }
  }

  /// 输入大于 0 的位置梯度照传，其余为 0（和 burn 的 `relu_backward` 一样，0 处取 0）。
  pub fn relu_backward(&self, input: &Self) -> Result<Self> {
    self.zip(input, "relu_backward", |g, x| if x > 0.0 { g } else { 0.0 })
  }

/// 节点 id，也就是它在 tape 里的下标。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Var(usize);

/// 算子和它的输入：反向需要的全部信息。输入的值从 tape 上按 id 取。
#[derive(Debug, Clone, Copy)]
enum Op {
  Leaf,
  Add(Var, Var),
  Mul(Var, Var),
  Matmul(Var, Var),
  Relu(Var),
  Sum(Var),
}

#[derive(Debug)]
struct Node {
  value: Tensor,
  op: Op,
  requires_grad: bool,
}

#[derive(Debug, Default)]
pub struct Tape {
  nodes: Vec<Node>,
}

  fn push_op(&mut self, value: Tensor, op: Op, inputs: &[Var]) -> Var {
    let requires_grad = inputs.iter().any(|&v| self.requires_grad(v));
    self.push(value, op, requires_grad)
  }

  pub fn add(&mut self, a: Var, b: Var) -> Result<Var> {
    let out = self.value(a)?.add(self.value(b)?)?;
    Ok(self.push_op(out, Op::Add(a, b), &[a, b]))
  }

  /// 一个节点的上游梯度 `grad` 分给各个要梯度的输入。
  fn input_grads(&self, op: Op, grad: &Tensor) -> Result<Vec<(Var, Tensor)>> {
    let wants = |v: Var| self.requires_grad(v);
    let mut out = Vec::with_capacity(2);
    match op {
      Op::Leaf => {}
      Op::Add(a, b) => {
        if wants(a) {
          out.push((a, grad.sum_to(self.value(a)?.shape())?));
        }
        if wants(b) {
          out.push((b, grad.sum_to(self.value(b)?.shape())?));
        }
      }
      Op::Mul(a, b) => {
        if wants(a) {
          out.push((a, grad.mul(self.value(b)?)?));
        }
        if wants(b) {
          out.push((b, grad.mul(self.value(a)?)?));
        }
      }
      Op::Matmul(a, b) => {
        // C = A B：dA = dC Bᵀ，dB = Aᵀ dC
        if wants(a) {
          out.push((a, grad.matmul(&self.value(b)?.transpose()?)?));
        }
        if wants(b) {
          out.push((b, self.value(a)?.transpose()?.matmul(grad)?));
        }
      }
      Op::Relu(a) => {
        if wants(a) {
          out.push((a, grad.relu_backward(self.value(a)?)?));
        }
      }
      Op::Sum(a) => {
        if wants(a) {
          let g = grad.data().first().copied().unwrap_or(0.0);
          out.push((a, Tensor::filled(self.value(a)?.shape(), g)));
        }
      }
    }
    Ok(out)
  }

  /// 从 `root` 反向传播。根的梯度是和它同形状的全 1。
  ///
  /// 顺序：节点按创建顺序追加，算子的输入一定早于输出存在，所以下标从大到小
  /// 就是一个逆拓扑序。轮到节点 i 时，所有用到它的节点（下标都比 i 大）已经
  /// 把梯度加进来了。
  pub fn backward(&self, root: Var) -> Result<Gradients> {
    let root_value = self.value(root)?;
    if !self.requires_grad(root) {
      return Err(Error(format!("node {} does not depend on any param", root.0)));
    }
    let mut grads: Vec<Option<Tensor>> = vec![None; root.0 + 1];
    grads[root.0] = Some(Tensor::filled(root_value.shape(), 1.0));

    for (id, node) in self.nodes[..=root.0].iter().enumerate().rev() {
      // 没有梯度：从根走不到它，或者它不要梯度
      let Some(grad) = &grads[id] else { continue };
      for (input, g) in self.input_grads(node.op, grad)? {
        let slot = &mut grads[input.0];
        *slot = Some(match slot.take() {
          Some(acc) => acc.add(&g)?,
          None => g,
        });
      }
    }
    Ok(Gradients { grads })
  }

/// `f` 用 tape 上的输入算出一个标量。对每个输入的每个元素做一次中心差分。
fn check(name: &str, inputs: &[Tensor], f: impl Fn(&mut Tape, &[Var]) -> Result<Var>) {
  let eval = |inputs: &[Tensor]| -> f32 {
    let mut tape = Tape::new();
    let vars: Vec<Var> = inputs.iter().map(|x| tape.param(x.clone())).collect();
    let out = f(&mut tape, &vars).unwrap();
    tape.value(out).unwrap().data()[0]
  };

  let mut tape = Tape::new();
  let vars: Vec<Var> = inputs.iter().map(|x| tape.param(x.clone())).collect();
  let out = f(&mut tape, &vars).unwrap();
  let grads = tape.backward(out).unwrap();

  let h = 1e-2;
  let mut max_err = 0.0f32;
  for (i, input) in inputs.iter().enumerate() {
    let analytic = grads.get(vars[i]).unwrap();
    assert_eq!(analytic.shape(), input.shape());
    for j in 0..input.data().len() {
      let shifted = |delta: f32| {
        let mut data = input.data().to_vec();
        data[j] += delta;
        let mut moved = inputs.to_vec();
        moved[i] = Tensor::new(data, input.shape()).unwrap();
        eval(&moved)
      };
      let numeric = (shifted(h) - shifted(-h)) / (2.0 * h);
      max_err = max_err.max((analytic.data()[j] - numeric).abs());
    }
  }
  println!("{name:<22} max |analytic - numeric| = {max_err:.2e}");
  assert!(max_err < 1e-2, "{name}: gradient mismatch {max_err}");
}

#[test]
fn add_broadcast_row() {
  check("add [3,4]+[4]", &[t(4, &[3, 4]), t(5, &[4])], |tape, v| {
    let y = tape.add(v[0], v[1])?;
    weigh(tape, y, 99)
  });
}

#[test]
fn reused_input() {
  // x 出现在两条路径上：梯度要相加
  check("x*x + x", &[t(11, &[4])], |tape, v| {
    let sq = tape.mul(v[0], v[0])?;
    let y = tape.add(sq, v[0])?;
    weigh(tape, y, 99)
  });
}
// ... sum、add_same_shape、mul、matmul、relu 同样的写法

running 7 tests
test add_broadcast_row ... add [3,4]+[4]          max |analytic - numeric| = 4.17e-6
ok
test add_same_shape ... add [3,4]+[3,4]        max |analytic - numeric| = 2.42e-5
ok
test matmul ... matmul [3,4]x[4,2]     max |analytic - numeric| = 7.15e-6
ok
test mul ... mul                    max |analytic - numeric| = 9.42e-6
ok
test relu ... relu                   max |analytic - numeric| = 3.81e-6
ok
test reused_input ... x*x + x                max |analytic - numeric| = 2.06e-6
ok
test sum ... sum                    max |analytic - numeric| = 9.54e-7
ok

test result: ok. 7 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

test add_broadcast_row ... ok
test add_same_shape ... ok
test matmul ... ok
test mul ... ok
test relu ... FAILED
test reused_input ... ok
test sum ... ok

failures:

---- relu stdout ----
relu                   max |analytic - numeric| = 7.32e-1

thread 'relu' (3475402) panicked at tests/finite_diff.rs:52:3:
relu: gradient mismatch 0.73178875

/// relu(x W1 + b1) W2 + b2，loss = Σ(pred - y)²
fn loss(tape: &mut Tape, params: &[Var; 4], x: Var, neg_y: Var) -> Result<(Var, Var)> {
  let [w1, b1, w2, b2] = *params;
  let h = tape.matmul(x, w1)?;
  let h = tape.add(h, b1)?;
  let h = tape.relu(h)?;
  let out = tape.matmul(h, w2)?;
  let pred = tape.add(out, b2)?;
  let diff = tape.add(pred, neg_y)?;
  let sq = tape.mul(diff, diff)?;
  Ok((tape.sum(sq)?, pred))
}

  for step in 0..=STEPS {
    // 每一步一条新 tape：上一步的图随旧 tape 一起丢掉
    let mut tape = Tape::new();
    let vars = params.clone().map(|p| tape.param(p));
    let x = tape.constant(x.clone());
    let neg_y = tape.constant(neg_y.clone());
    let (l, pred) = loss(&mut tape, &vars, x, neg_y).unwrap();
    let grads = tape.backward(l).unwrap();

    if step % 25 == 0 {
      let l = tape.value(l).unwrap().data()[0];
      println!("step {step:>4}  loss {l:.6}  nodes {}", tape.len());
    }
    // ... step == STEPS 时打印预测、断言 loss < 1e-3、退出
    for (p, v) in params.iter_mut().zip(vars) {
      *p = p.add_scaled(grads.get(v).unwrap(), -lr).unwrap();
    }
  }

running 1 test
test xor ... step    0  loss 2.073196  nodes 14
step   25  loss 0.204576  nodes 14
step   50  loss 0.049814  nodes 14
step   75  loss 0.009232  nodes 14
step  100  loss 0.001496  nodes 14
step  125  loss 0.000226  nodes 14
step  150  loss 0.000034  nodes 14
step  175  loss 0.000005  nodes 14
step  200  loss 0.000001  nodes 14
step  225  loss 0.000000  nodes 14
step  250  loss 0.000000  nodes 14
step  275  loss 0.000000  nodes 14
step  300  loss 0.000000  nodes 14
pred [8.4117055e-6, 0.99999243, 0.9999897, 1.0021031e-5]
ok

fn burn(c: &Case) -> (f32, [Vec<f32>; 4]) {
  let device = Device::flex().autodiff();
  let t2 = |v: &Vec<f32>, shape: [usize; 2]| BTensor::<2>::from_data(TensorData::new(v.clone(), shape), &device);
  let t1 = |v: &Vec<f32>| BTensor::<1>::from_data(TensorData::new(v.clone(), [v.len()]), &device);

  let x = t2(&c.x, [4, 2]);
  let y = t2(&c.y, [4, 1]);
  let w1 = t2(&c.w1, [2, H]).require_grad();
  let b1 = t1(&c.b1).require_grad();
  let w2 = t2(&c.w2, [H, 1]).require_grad();
  let b2 = t1(&c.b2).require_grad();

  let h = activation::relu(x.matmul(w1.clone()) + b1.clone().unsqueeze::<2>());
  let pred = h.matmul(w2.clone()) + b2.clone().unsqueeze::<2>();
  let diff = pred - y;
  let loss = (diff.clone() * diff).sum();

  let grads = loss.backward();
  // ... 用 into_data().try_to_vec::<f32>() 取出 loss 和四个梯度
}

running 1 test
test same_gradients_as_burn ... loss  mini 1.293978
loss  burn 1.2939781
w1.grad mini [0.0, -0.007851243, -0.003602624, 0.0, 0.59956694, 0.27511376]
w1.grad burn [0.0, -0.007851243, -0.003602624, 0.0, 0.59956694, 0.27511376]
b1.grad mini [0.18811604, -0.35759163, -0.16408247]
b1.grad burn [0.18811604, -0.35759163, -0.16408247]
w2.grad mini [0.035068847, 0.38163775, -0.04492247]
w2.grad burn [0.035068847, 0.38163775, -0.04492247]
b2.grad mini [-0.42864764]
b2.grad burn [-0.42864764]
max |mini - burn| = 0e0
ok

  let mut tape = Tape::new();
  let x = tape.param(Tensor::new(vec![1.0, 2.0], &[2]).unwrap());
  let c = tape.constant(Tensor::new(vec![3.0, 3.0], &[2]).unwrap());
  let h = tape.mul(x, x).unwrap();
  let xc = tape.mul(x, c).unwrap();
  let s = tape.add(h, xc).unwrap();
  let loss = tape.sum(s).unwrap();

  let g1 = tape.backward(loss).unwrap();
  let g2 = tape.backward(loss).unwrap();
  // ... 打印 x、h、c 的梯度，对常量 backward，对 1 维张量 matmul

test extras ... x.grad (1st backward) = Some([5.0, 7.0])
x.grad (2nd backward) = Some([5.0, 7.0])
h.grad                = Some([1.0, 1.0])
c.grad                = None
backward(c)           = Err(Error("node 1 does not depend on any param"))
matmul [2] x [2]      = Err(Error("matmul needs 2-D tensors, got [2] and [2]"))
ok

test same_gradients_as_burn ... loss  mini 1.293978
loss  burn 1.2939781
w1.grad mini [0.0, -0.0039256215, -0.001801312, 0.0, 0.29978347, 0.13755688]
w1.grad burn [0.0, -0.007851243, -0.003602624, 0.0, 0.59956694, 0.27511376]
b1.grad mini [0.09405802, -0.17879581, -0.082041234]
b1.grad burn [0.18811604, -0.35759163, -0.16408247]
w2.grad mini [0.017534424, 0.19081888, -0.022461236]
w2.grad burn [0.035068847, 0.38163775, -0.04492247]
b2.grad mini [-0.21432382]
b2.grad burn [-0.42864764]
max |mini - burn| = 2.9978347e-1

thread 'same_gradients_as_burn' (3476128) panicked at tests/burn_compare.rs:102:3:
assertion failed: max_diff < 1e-5
...
test reused_input ... x*x + x                max |analytic - numeric| = 9.32e-1

thread 'reused_input' (3476140) panicked at tests/finite_diff.rs:52:3:
x*x + x: gradient mismatch 0.93168473
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
FAILED
test sum ... sum                    max |analytic - numeric| = 9.54e-7
ok

failures:
    reused_input

test result: FAILED. 6 passed; 1 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
...
test xor ... step    0  loss 2.073196  nodes 14
step   25  loss 0.387116  nodes 14
step   50  loss 0.204980  nodes 14
step   75  loss 0.105003  nodes 14
step  100  loss 0.049699  nodes 14
step  125  loss 0.022374  nodes 14
step  150  loss 0.009434  nodes 14
step  175  loss 0.003916  nodes 14
step  200  loss 0.001581  nodes 14
step  225  loss 0.000625  nodes 14
step  250  loss 0.000249  nodes 14
step  275  loss 0.000097  nodes 14
step  300  loss 0.000037  nodes 14
pred [0.002767101, 0.9974603, 0.996524, 0.0033554286]
ok