Automatic Differentiation (Autograd)

1/**2 * Example 14: Automatic Differentiation (Autograd)3 *4 * Deepbox's autograd system tracks operations on GradTensors to build a5 * computation graph, then computes gradients via reverse-mode differentiation.6 */78import { GradTensor, noGrad, parameter, tensor } from "deepbox/ndarray";910console.log("=== Automatic Differentiation ===\n");1112// ---------------------------------------------------------------------------13// Part 1: Basic gradient computation14// ---------------------------------------------------------------------------15console.log("--- Part 1: Basic Gradients ---");1617// f(x) = x^2  =>  df/dx = 2x18const x = parameter([2, 3, 4]);19const y = x.mul(x).sum();20y.backward();2122console.log("x       :", x.tensor.toString());23console.log("f(x)=x²  sum:", y.tensor.toString());24console.log("grad    :", x.grad?.toString() ?? "null");25// Expected gradients: [4, 6, 8]2627// ---------------------------------------------------------------------------28// Part 2: Multi-variable gradients29// ---------------------------------------------------------------------------30console.log("\n--- Part 2: Multi-Variable Gradients ---");3132const a = parameter([33  [1, 2],34  [3, 4],35]);36const w = parameter([[0.5], [0.5]]);3738// y = sum(a @ w)39const z = a.matmul(w).sum();40z.backward();4142console.log("a =", a.tensor.toString());43console.log("w =", w.tensor.toString());44console.log("z = sum(a @ w) =", z.tensor.toString());45console.log("dz/da =", a.grad?.toString() ?? "null");46console.log("dz/dw =", w.grad?.toString() ?? "null");4748// ---------------------------------------------------------------------------49// Part 3: Chained operations50// ---------------------------------------------------------------------------51console.log("\n--- Part 3: Chained Operations ---");5253const p = parameter([1, 2, 3, 4]);5455// f(p) = sum(relu(p * 2 - 3))56const scaled = p.mul(GradTensor.fromTensor(tensor([2, 2, 2, 2]), { requiresGrad: false }));57const shifted = scaled.sub(GradTensor.fromTensor(tensor([3, 3, 3, 3]), { requiresGrad: false }));58const activated = shifted.relu();59const loss = activated.sum();60loss.backward();6162console.log("p       :", p.tensor.toString());63console.log("2p - 3  :", shifted.tensor.toString());64console.log("relu    :", activated.tensor.toString());65console.log("grad    :", p.grad?.toString() ?? "null");6667// ---------------------------------------------------------------------------68// Part 4: noGrad for inference69// ---------------------------------------------------------------------------70console.log("\n--- Part 4: noGrad for Inference ---");7172const q = parameter([1, 2, 3]);73noGrad(() => {74  // Operations inside noGrad do not track gradients75  const result = q.mul(q);76  console.log("noGrad result:", result.tensor.toString());77  console.log("requiresGrad:", result.requiresGrad);78});7980// ---------------------------------------------------------------------------81// Part 5: Zero gradients and re-compute82// ---------------------------------------------------------------------------83console.log("\n--- Part 5: Gradient Accumulation ---");8485const v = parameter([1, 2, 3]);8687// First backward88const loss1 = v.mul(v).sum();89loss1.backward();90console.log("After first backward, grad:", v.grad?.toString() ?? "null");9192// Zero gradients before second pass93v.zeroGrad();94console.log("After zeroGrad, grad:", v.grad?.toString() ?? "null");9596// Second backward with different computation97const loss2 = v.mul(GradTensor.fromTensor(tensor([3, 3, 3]), { requiresGrad: false })).sum();98loss2.backward();99console.log("After second backward, grad:", v.grad?.toString() ?? "null");100101console.log("\n=== Autograd Complete ===");102