Automatic Differentiation (Autograd)

1/**2 * Example 14: Automatic Differentiation (Autograd)3 *4 * Deepbox records the operations done on tracked tensors (GradTensor) as a5 * computation graph. Calling backward() walks the graph in reverse and fills6 * in the gradient of every tracked input.7 *8 * Use parameter(...) only for tensors you want gradients for, such as weights9 * or a value you differentiate with respect to. Training data stays a plain10 * tensor.11 */12 13import { noGrad, parameter, tensor } from "deepbox/ndarray";14import { Linear } from "deepbox/nn";15 16console.log("=== Automatic Differentiation ===\n");17 18// ---------------------------------------------------------------------------19// Part 1: Basic gradient computation20// ---------------------------------------------------------------------------21console.log("--- Part 1: Basic Gradients ---");22 23// f(x) = sum(x^2)  =>  df/dx = 2x24const x = parameter([2, 3, 4]);25const y = x.mul(x).sum();26y.backward();27 28console.log("x       :", x.tensor.toString());29console.log("f(x)    :", y.item());30console.log("grad    :", x.grad?.toString() ?? "null");31// Expected gradients: [4, 6, 8]32 33// ---------------------------------------------------------------------------34// Part 2: Multi-variable gradients35// ---------------------------------------------------------------------------36console.log("\n--- Part 2: Multi-Variable Gradients ---");37 38const a = parameter([39  [1, 2],40  [3, 4],41]);42const w = parameter([[0.5], [0.5]]);43 44// z = sum(a @ w)45const z = a.matmul(w).sum();46z.backward();47 48console.log("a =", a.tensor.toString());49console.log("w =", w.tensor.toString());50console.log("z = sum(a @ w) =", z.item());51console.log("dz/da =", a.grad?.toString() ?? "null");52console.log("dz/dw =", w.grad?.toString() ?? "null");53 54// ---------------------------------------------------------------------------55// Part 3: Chained operations56// ---------------------------------------------------------------------------57console.log("\n--- Part 3: Chained Operations ---");58 59const p = parameter([1, 2, 3, 4]);60 61// f(p) = sum(relu(2p - 3)). Plain numbers can be used directly in the chain.62const shifted = p.mul(2).sub(3);63const activated = shifted.relu();64const loss = activated.sum();65loss.backward();66 67console.log("p       :", p.tensor.toString());68console.log("2p - 3  :", shifted.tensor.toString());69console.log("relu    :", activated.tensor.toString());70console.log("grad    :", p.grad?.toString() ?? "null");71// relu passes the gradient only where 2p - 3 is positive, so p = 1 gets 0.72 73// ---------------------------------------------------------------------------74// Part 4: noGrad for inference75// ---------------------------------------------------------------------------76console.log("\n--- Part 4: noGrad for Inference ---");77 78const q = parameter([1, 2, 3]);79noGrad(() => {80  // Operations inside noGrad are not recorded, so the result has no graph.81  const result = q.mul(q);82  console.log("noGrad result:", result.tensor.toString());83  console.log("requiresGrad:", result.requiresGrad);84});85 86// ---------------------------------------------------------------------------87// Part 5: Gradient accumulation and zeroGrad88// ---------------------------------------------------------------------------89console.log("\n--- Part 5: Gradient Accumulation ---");90 91const v = parameter([1, 2, 3]);92 93// First backward94v.mul(v).sum().backward();95console.log("After first backward, grad:", v.grad?.toString() ?? "null");96 97// Without zeroGrad, a second backward adds to the stored gradient.98v.mul(3).sum().backward();99console.log("After second backward without zeroGrad:", v.grad?.toString() ?? "null");100 101// Reset the gradient before computing a new one.102v.zeroGrad();103v.mul(3).sum().backward();104console.log("After zeroGrad and a new backward:", v.grad?.toString() ?? "null");105 106// ---------------------------------------------------------------------------107// Part 6: Gradients of a layer with plain-tensor input108// ---------------------------------------------------------------------------109console.log("\n--- Part 6: Layer Gradients from a Plain Tensor ---");110 111// The input is a plain tensor. The layer's weight and bias are the tracked112// parameters, so forward() returns a tracked result.113const layer = new Linear(2, 1);114const input = tensor([115  [1, 2],116  [3, 4],117]);118const out = layer.forward(input);119console.log("Output tracks gradients:", out.requiresGrad);120 121out.sum().backward();122// namedParameters() yields [name, GradTensor] pairs, here "weight" and "bias".123for (const [name, param] of layer.namedParameters()) {124  console.log(`${name}.grad:`, param.grad?.toString() ?? "null");125}126// d(sum(out))/d(weight) is the column sums of the input, [1 + 3, 2 + 4].127// d(sum(out))/d(bias) is the number of input rows, 2.128 129console.log("\n=== Autograd Complete ===");