@danielsimonjr/mathts-autograd 0.1.0 → 0.1.2

package/dist/index.js

@@ -6,7 +6,9 @@ var DualTensor = class _DualTensor {
   tangent;
   constructor(shape, primal, tangent) {
     if (primal.length !== tangent.length) {
-      throw new Error(`DualTensor: primal length ${primal.length} != tangent length ${tangent.length}`);
+      throw new Error(
+        `DualTensor: primal length ${primal.length} != tangent length ${tangent.length}`
+      );
     }
     this.shape = shape;
     this.primal = primal;
@@ -137,7 +139,7 @@ function forwardGrad(fn, x) {
 }
 
 // src/tape.ts
-import { Tensor as Tensor3 } from "@danielsimonjr/mathts-tensor";
+import { Tensor as Tensor3, tensorSvd, tensorEig } from "@danielsimonjr/mathts-tensor";
 var Tape = class {
   nodes = [];
   inputGradSlots = /* @__PURE__ */ new Map();
@@ -175,12 +177,24 @@ var Tape = class {
   }
 };
 var TapedTensor = class _TapedTensor {
-  constructor(shape, primal, tape, id) {
+  constructor(shape, primal, tape, id, axisLabels) {
     this.shape = shape;
     this.primal = primal;
     this.tape = tape;
     this.id = id;
+    if (axisLabels !== void 0) {
+      this.axisLabels = [...axisLabels];
+    }
   }
+  shape;
+  primal;
+  tape;
+  id;
+  /**
+   * Optional per-axis Index labels. When set, enables `contract`.
+   * Must have the same length as `shape` when present.
+   */
+  axisLabels;
   /**
    * S5 fix: existing engine ops (e.g. lower, pderiv, contract) reach into
    * `.data`. The getter returns the primal so those ops still work when a
@@ -192,10 +206,17 @@ var TapedTensor = class _TapedTensor {
   }
   static fromTensorAsInput(t, tape) {
     const { id } = tape.allocate(t.data.length);
-    return new _TapedTensor(t.shape, new Float64Array(t.data), tape, id);
+    return new _TapedTensor(t.shape, new Float64Array(t.data), tape, id, t.axisLabels);
   }
   toPrimalTensor() {
-    return new Tensor3(this.shape, new Float64Array(this.primal));
+    return new Tensor3(this.shape, new Float64Array(this.primal), this.axisLabels);
+  }
+  /**
+   * Reconstruct the Tensor primal with the given axisLabels (used internally
+   * when the primal was computed from an op that produces labelled output).
+   */
+  toPrimalTensorWith(axisLabels) {
+    return new Tensor3(this.shape, new Float64Array(this.primal), axisLabels);
   }
   add(other) {
     this.checkSameShape(other, "add");
@@ -225,6 +246,39 @@ var TapedTensor = class _TapedTensor {
     });
     return new _TapedTensor(this.shape, out, this.tape, id);
   }
+  /**
+   * Elementwise division: this / other.
+   *
+   * Adjoints (quotient rule):
+   *   dA = dY / b
+   *   dB = −dY · a / b²
+   *
+   * The alias case (a.divide(a)) is handled explicitly: the combined gradient
+   * is dA + dB = dY/a − dY·a/a² = dY/a − dY/a = 0. This is correct since
+   * a/a = 1 everywhere and d(1)/da = 0.
+   */
+  divide(other) {
+    this.checkSameShape(other, "divide");
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = this.primal[i] / other.primal[i];
+    const thisPrimal = new Float64Array(this.primal);
+    const otherPrimal = new Float64Array(other.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const otherGradSlot = this.tape.getInputGrad(other.id);
+    const isAliased = this === other;
+    const { id } = this.tape.record([this.id, other.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        if (isAliased) {
+        } else {
+          const b = otherPrimal[i];
+          const a = thisPrimal[i];
+          thisGradSlot[i] += outputGrad[i] / b;
+          otherGradSlot[i] += -outputGrad[i] * a / (b * b);
+        }
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
   mul(other) {
     this.checkSameShape(other, "mul");
     const out = new Float64Array(this.primal.length);
@@ -257,12 +311,1457 @@ var TapedTensor = class _TapedTensor {
     });
     return new _TapedTensor(this.shape, out, this.tape, id);
   }
+  /**
+   * Reverse-mode AD over `Tensor.contract`.
+   *
+   * Both operands must carry `axisLabels`; the resulting TapedTensor inherits
+   * the contracted-output axisLabels (non-shared axis concatenation).
+   *
+   * Adjoint derivation (T-notation):
+   *   Let Y = A.contract(B)  (contraction over shared indices S).
+   *   dA = dY.contract(B')  where B' = B with its free axes re-labelled to
+   *        match the shared positions in A. Equivalently: for each element of A,
+   *        dA[...a_free, ...s] = Σ_{...b_free} dY[...a_free, ...b_free] · B[...s, ...b_free]
+   *   dB = A'.contract(dY)  symmetrically.
+   *
+   * Implementation: because `Tensor.contract` matches by Index identity, we
+   * build the backward contraction by re-labelling the free axes of dY with
+   * the Index objects from the other operand — so `Tensor.contract` automatically
+   * finds the right shared axes.
+   */
+  contract(other) {
+    const aPrimal = this.toPrimalTensorWith(this.axisLabels);
+    const bPrimal = other.toPrimalTensorWith(other.axisLabels);
+    const yPrimal = aPrimal.contract(bPrimal);
+    const aLabels = this.axisLabels;
+    const bLabels = other.axisLabels;
+    const yLabels = yPrimal.axisLabels;
+    const aFreeAxes = [];
+    const aSharedAxes = [];
+    for (let i = 0; i < aLabels.length; i++) {
+      const inOutput = yLabels.some((yl) => yl.matches(aLabels[i]));
+      if (inOutput) aFreeAxes.push(i);
+      else aSharedAxes.push(i);
+    }
+    const bFreeAxes = [];
+    const bSharedAxes = [];
+    for (let i = 0; i < bLabels.length; i++) {
+      const inOutput = yLabels.some((yl) => yl.matches(bLabels[i]));
+      if (inOutput) bFreeAxes.push(i);
+      else bSharedAxes.push(i);
+    }
+    const aGradSlot = this.tape.getInputGrad(this.id);
+    const bGradSlot = this.tape.getInputGrad(other.id);
+    const aPrimalData = new Float64Array(this.primal);
+    const bPrimalData = new Float64Array(other.primal);
+    const aShape = [...this.shape];
+    const bShape = [...other.shape];
+    const { id } = this.tape.record(
+      [this.id, other.id],
+      yPrimal.data.length,
+      (outputGrad) => {
+        const dY = new Tensor3([...yPrimal.shape], new Float64Array(outputGrad), yLabels);
+        let bRelabelled = new Tensor3(bShape, bPrimalData, bLabels);
+        for (let k = 0; k < bSharedAxes.length; k++) {
+          const bAxis = bSharedAxes[k];
+          const aAxis = aSharedAxes[k];
+          bRelabelled = bRelabelled.replaceIndex(bLabels[bAxis], aLabels[aAxis]);
+        }
+        const dAFull = dY.contract(bRelabelled);
+        const dAData = _scatterToOriginalAxes(dAFull, aLabels);
+        for (let i = 0; i < aGradSlot.length; i++) {
+          aGradSlot[i] += dAData[i];
+        }
+        let aRelabelled = new Tensor3(aShape, aPrimalData, aLabels);
+        for (let k = 0; k < aSharedAxes.length; k++) {
+          const aAxis = aSharedAxes[k];
+          const bAxis = bSharedAxes[k];
+          aRelabelled = aRelabelled.replaceIndex(aLabels[aAxis], bLabels[bAxis]);
+        }
+        const dBFull = aRelabelled.contract(dY);
+        const dBData = _scatterToOriginalAxes(dBFull, bLabels);
+        for (let i = 0; i < bGradSlot.length; i++) {
+          bGradSlot[i] += dBData[i];
+        }
+      }
+    );
+    return new _TapedTensor(
+      [...yPrimal.shape],
+      new Float64Array(yPrimal.data),
+      this.tape,
+      id,
+      yLabels
+    );
+  }
+  /**
+   * Reverse-mode AD over batched matmul.
+   *
+   * Accepts operands of rank ≥ 2. Convention: trailing 2 axes are the matrix
+   * dims (rows × cols); all leading axes are batch dims (must be broadcast-
+   * compatible). For rank-2 inputs this is classical matrix multiplication.
+   *
+   * Adjoint derivation (classical matrix calc, extended to batched):
+   *   Y = A @ B  (A: ...×m×k, B: ...×k×n  →  Y: ...×m×n)
+   *   dA = dY @ Bᵀ       (dA: ...×m×k)
+   *   dB = Aᵀ @ dY       (dB: ...×k×n)
+   *
+   * Implemented via `Tensor.einsum` with a dynamically-built spec:
+   *   forward:  '...ik,...kj->...ij'
+   *   dA:       '...ij,...kj->...ik'   (contract dY with B on j)
+   *   dB:       '...ki,...kj->...ij'   (contract A on k with dY on k)
+   */
+  matmul(other) {
+    const rankA = this.shape.length;
+    const rankB = other.shape.length;
+    if (rankA < 2) {
+      throw new Error(`TapedTensor.matmul: 'this' operand must have rank \u2265 2, got rank ${rankA}`);
+    }
+    if (rankB < 2) {
+      throw new Error(`TapedTensor.matmul: 'other' operand must have rank \u2265 2, got rank ${rankB}`);
+    }
+    const batchRankA = rankA - 2;
+    const batchRankB = rankB - 2;
+    if (batchRankA !== batchRankB) {
+      throw new Error(
+        `TapedTensor.matmul: batch rank mismatch (${batchRankA} vs ${batchRankB}); both operands must have the same rank`
+      );
+    }
+    for (let b = 0; b < batchRankA; b++) {
+      if (this.shape[b] !== other.shape[b]) {
+        throw new Error(
+          `TapedTensor.matmul: batch dimension ${b} mismatch (${this.shape[b]} vs ${other.shape[b]})`
+        );
+      }
+    }
+    const m = this.shape[rankA - 2];
+    const k = this.shape[rankA - 1];
+    const k2 = other.shape[rankB - 2];
+    const n = other.shape[rankB - 1];
+    if (k !== k2) {
+      throw new Error(`TapedTensor.matmul: inner dimension mismatch (${k} vs ${k2})`);
+    }
+    const aShape = [...this.shape];
+    const bShape = [...other.shape];
+    const batchShape = aShape.slice(0, batchRankA);
+    const yShape = [...batchShape, m, n];
+    const aPrimalData = new Float64Array(this.primal);
+    const bPrimalData = new Float64Array(other.primal);
+    const yData = _batchedMatmul(aPrimalData, bPrimalData, aShape, bShape);
+    const aGradSlot = this.tape.getInputGrad(this.id);
+    const bGradSlot = this.tape.getInputGrad(other.id);
+    const { id } = this.tape.record(
+      [this.id, other.id],
+      yData.length,
+      (outputGrad) => {
+        const dA = _batchedMatmulGradA(outputGrad, bPrimalData, yShape, bShape);
+        for (let i = 0; i < aGradSlot.length; i++) {
+          aGradSlot[i] += dA[i];
+        }
+        const dB = _batchedMatmulGradB(aPrimalData, outputGrad, aShape, yShape);
+        for (let i = 0; i < bGradSlot.length; i++) {
+          bGradSlot[i] += dB[i];
+        }
+      }
+    );
+    return new _TapedTensor(yShape, new Float64Array(yData), this.tape, id);
+  }
   checkSameShape(other, op) {
     if (this.shape.length !== other.shape.length || !this.shape.every((v, i) => v === other.shape[i])) {
       throw new Error(`TapedTensor.${op}: shape mismatch [${this.shape}] vs [${other.shape}]`);
     }
   }
+  // ---------------------------------------------------------------------------
+  // Reductions
+  // ---------------------------------------------------------------------------
+  /**
+   * Sum elements along the given axis/axes (or all axes if omitted).
+   *
+   * Adjoint: dX[...] = dY[reduced(idx)] broadcast back to input shape.
+   * Each input element receives the output-gradient entry from its reduced
+   * counterpart (the non-reduced coordinates select the dY element; the
+   * reduced coordinates are collapsed to 0 in the keepDims=false case).
+   */
+  sum(axis, opts) {
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.sum(axis, opts);
+    const inputShape = [...this.shape];
+    const inputSize = this.primal.length;
+    const keepDims = opts?.keepDims ?? false;
+    const axes = _resolveAxes(axis, this.shape.length);
+    const axisSet = new Set(axes);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+      const outShape = [...outT.shape];
+      const outStrides = _rowMajorStrides(outShape);
+      const inStrides = _rowMajorStrides(inputShape);
+      const inCoords = new Array(inputShape.length).fill(0);
+      for (let n = 0; n < inputSize; n++) {
+        let outFlat = 0;
+        let outDim = 0;
+        for (let k = 0; k < inputShape.length; k++) {
+          if (axisSet.has(k)) {
+            if (keepDims) {
+              outDim++;
+            }
+          } else {
+            outFlat += inCoords[k] * outStrides[outDim];
+            outDim++;
+          }
+        }
+        thisGradSlot[n] += outputGrad[outFlat];
+        for (let k = inputShape.length - 1; k >= 0; k--) {
+          if (++inCoords[k] < inputShape[k]) break;
+          inCoords[k] = 0;
+        }
+      }
+      void inStrides;
+    });
+    return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+  }
+  /**
+   * Arithmetic mean along the given axis/axes (or all axes if omitted).
+   *
+   * Adjoint: dX[...] = dY[reduced(idx)] / N, broadcast back to input shape.
+   * N = product of reduced-axis dimensions.
+   */
+  mean(axis, opts) {
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.mean(axis, opts);
+    const inputShape = [...this.shape];
+    const inputSize = this.primal.length;
+    const keepDims = opts?.keepDims ?? false;
+    const axes = _resolveAxes(axis, this.shape.length);
+    const axisSet = new Set(axes);
+    const N = axes.reduce((prod, ax) => prod * this.shape[ax], 1);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+      const outShape = [...outT.shape];
+      const outStrides = _rowMajorStrides(outShape);
+      const inCoords = new Array(inputShape.length).fill(0);
+      for (let n = 0; n < inputSize; n++) {
+        let outFlat = 0;
+        let outDim = 0;
+        for (let k = 0; k < inputShape.length; k++) {
+          if (axisSet.has(k)) {
+            if (keepDims) outDim++;
+          } else {
+            outFlat += inCoords[k] * outStrides[outDim];
+            outDim++;
+          }
+        }
+        thisGradSlot[n] += outputGrad[outFlat] / N;
+        for (let k = inputShape.length - 1; k >= 0; k--) {
+          if (++inCoords[k] < inputShape[k]) break;
+          inCoords[k] = 0;
+        }
+      }
+    });
+    return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+  }
+  /**
+   * Product of elements along the given axis/axes (or all axes if omitted).
+   *
+   * Adjoint: dX_i = dY * (prod_over_axes(x) / x_i) per element.
+   *
+   * Zero-element corners:
+   * - Exactly one x_i = 0 in the reduced group: d/dx_i = product of all others
+   *   (which is the full product / x_i evaluated via alternate product),
+   *   and d/dx_j = 0 for all j ≠ i where x_j ≠ 0.
+   * - Two or more zeros in the reduced group: gradient is 0 everywhere for that
+   *   group (because changing any single zero cannot change a product that is
+   *   zero due to another zero).
+   *
+   * Implementation: uses prefix/suffix products to handle zeros robustly.
+   */
+  prod(axis, opts) {
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.prod(axis, opts);
+    const inputShape = [...this.shape];
+    const keepDims = opts?.keepDims ?? false;
+    const axes = _resolveAxes(axis, this.shape.length);
+    const axisSet = new Set(axes);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+      const outShape = [...outT.shape];
+      const outStrides = _rowMajorStrides(outShape);
+      const inStrides = _rowMajorStrides(inputShape);
+      const groups = /* @__PURE__ */ new Map();
+      const inCoords = new Array(inputShape.length).fill(0);
+      for (let n = 0; n < primalData.length; n++) {
+        let outFlat = 0;
+        let outDim = 0;
+        for (let k = 0; k < inputShape.length; k++) {
+          if (axisSet.has(k)) {
+            if (keepDims) outDim++;
+          } else {
+            outFlat += inCoords[k] * outStrides[outDim];
+            outDim++;
+          }
+        }
+        const group = groups.get(outFlat);
+        if (group === void 0) {
+          groups.set(outFlat, [n]);
+        } else {
+          group.push(n);
+        }
+        for (let k = inputShape.length - 1; k >= 0; k--) {
+          if (++inCoords[k] < inputShape[k]) break;
+          inCoords[k] = 0;
+        }
+      }
+      void inStrides;
+      for (const [outFlat, group] of groups) {
+        const dY = outputGrad[outFlat];
+        const m = group.length;
+        const prefix = new Float64Array(m);
+        prefix[0] = 1;
+        for (let i = 1; i < m; i++) {
+          prefix[i] = prefix[i - 1] * primalData[group[i - 1]];
+        }
+        const suffix = new Float64Array(m);
+        suffix[m - 1] = 1;
+        for (let i = m - 2; i >= 0; i--) {
+          suffix[i] = suffix[i + 1] * primalData[group[i + 1]];
+        }
+        for (let i = 0; i < m; i++) {
+          thisGradSlot[group[i]] += dY * prefix[i] * suffix[i];
+        }
+      }
+    });
+    return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+  }
+  /**
+   * Maximum value along the given axis/axes (or all axes if omitted).
+   *
+   * Adjoint: dY is scattered to the argmax position(s).
+   * Tie-breaking: "first-wins" — the gradient flows to the first (smallest
+   * flat-index) element among those that attain the maximum.
+   */
+  max(axis, opts) {
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.max(axis, opts);
+    const inputShape = [...this.shape];
+    const keepDims = opts?.keepDims ?? false;
+    const axes = _resolveAxes(axis, this.shape.length);
+    const axisSet = new Set(axes);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+      const outShape = [...outT.shape];
+      const outStrides = _rowMajorStrides(outShape);
+      const argmaxSlot = new Int32Array(outT.data.length).fill(-1);
+      const inCoords = new Array(inputShape.length).fill(0);
+      for (let n = 0; n < primalData.length; n++) {
+        let outFlat = 0;
+        let outDim = 0;
+        for (let k = 0; k < inputShape.length; k++) {
+          if (axisSet.has(k)) {
+            if (keepDims) outDim++;
+          } else {
+            outFlat += inCoords[k] * outStrides[outDim];
+            outDim++;
+          }
+        }
+        if (argmaxSlot[outFlat] === -1 || primalData[n] > primalData[argmaxSlot[outFlat]]) {
+          argmaxSlot[outFlat] = n;
+        }
+        for (let k = inputShape.length - 1; k >= 0; k--) {
+          if (++inCoords[k] < inputShape[k]) break;
+          inCoords[k] = 0;
+        }
+      }
+      for (let outFlat = 0; outFlat < outT.data.length; outFlat++) {
+        const winner = argmaxSlot[outFlat];
+        if (winner !== -1) {
+          thisGradSlot[winner] += outputGrad[outFlat];
+        }
+      }
+    });
+    return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+  }
+  /**
+   * Minimum value along the given axis/axes (or all axes if omitted).
+   *
+   * Adjoint: dY is scattered to the argmin position(s).
+   * Tie-breaking: "first-wins" — gradient flows to the first (smallest
+   * flat-index) element that attains the minimum.
+   */
+  min(axis, opts) {
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.min(axis, opts);
+    const inputShape = [...this.shape];
+    const keepDims = opts?.keepDims ?? false;
+    const axes = _resolveAxes(axis, this.shape.length);
+    const axisSet = new Set(axes);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+      const outShape = [...outT.shape];
+      const outStrides = _rowMajorStrides(outShape);
+      const argminSlot = new Int32Array(outT.data.length).fill(-1);
+      const inCoords = new Array(inputShape.length).fill(0);
+      for (let n = 0; n < primalData.length; n++) {
+        let outFlat = 0;
+        let outDim = 0;
+        for (let k = 0; k < inputShape.length; k++) {
+          if (axisSet.has(k)) {
+            if (keepDims) outDim++;
+          } else {
+            outFlat += inCoords[k] * outStrides[outDim];
+            outDim++;
+          }
+        }
+        if (argminSlot[outFlat] === -1 || primalData[n] < primalData[argminSlot[outFlat]]) {
+          argminSlot[outFlat] = n;
+        }
+        for (let k = inputShape.length - 1; k >= 0; k--) {
+          if (++inCoords[k] < inputShape[k]) break;
+          inCoords[k] = 0;
+        }
+      }
+      for (let outFlat = 0; outFlat < outT.data.length; outFlat++) {
+        const winner = argminSlot[outFlat];
+        if (winner !== -1) {
+          thisGradSlot[winner] += outputGrad[outFlat];
+        }
+      }
+    });
+    return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+  }
+  /**
+   * p-norm of the tensor.
+   *
+   * Supported p values: 1, 2, 'fro', 'inf'. Default p = 2.
+   * When `opts.axis` is given, reduces along that axis; otherwise reduces all axes.
+   *
+   * Adjoints:
+   * - p=2 / p='fro': dX = dY · x / ‖x‖₂  (Frobenius is the 2-norm of the flattened tensor)
+   * - p=1:           dX = dY · sign(x)  (subgradient = 0 at exact zero)
+   * - p='inf':       dX scattered to the element(s) of max absolute value;
+   *                  tie-breaking: first-wins. Sign of the scattered gradient
+   *                  matches sign(x_max).
+   */
+  norm(opts) {
+    const p = opts?.p ?? 2;
+    const keepDims = opts?.keepDims ?? false;
+    const primalT = this.toPrimalTensor();
+    const outT = primalT.norm({ p, axis: opts?.axis, keepDims });
+    const inputShape = [...this.shape];
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const axes = opts?.axis !== void 0 ? [opts.axis] : _resolveAxes(void 0, this.shape.length);
+    const axisSet = new Set(axes);
+    if (p === 2 || p === "fro") {
+      const normData = new Float64Array(outT.data);
+      const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+        const outShape = [...outT.shape];
+        const outStrides = _rowMajorStrides(outShape);
+        const inCoords = new Array(inputShape.length).fill(0);
+        for (let n = 0; n < primalData.length; n++) {
+          let outFlat = 0;
+          let outDim = 0;
+          for (let k = 0; k < inputShape.length; k++) {
+            if (axisSet.has(k)) {
+              if (keepDims) outDim++;
+            } else {
+              outFlat += inCoords[k] * outStrides[outDim];
+              outDim++;
+            }
+          }
+          const normVal = normData[outFlat];
+          if (normVal !== 0) {
+            thisGradSlot[n] += outputGrad[outFlat] * primalData[n] / normVal;
+          }
+          for (let k = inputShape.length - 1; k >= 0; k--) {
+            if (++inCoords[k] < inputShape[k]) break;
+            inCoords[k] = 0;
+          }
+        }
+      });
+      return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+    }
+    if (p === 1) {
+      const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+        const outShape = [...outT.shape];
+        const outStrides = _rowMajorStrides(outShape);
+        const inCoords = new Array(inputShape.length).fill(0);
+        for (let n = 0; n < primalData.length; n++) {
+          let outFlat = 0;
+          let outDim = 0;
+          for (let k = 0; k < inputShape.length; k++) {
+            if (axisSet.has(k)) {
+              if (keepDims) outDim++;
+            } else {
+              outFlat += inCoords[k] * outStrides[outDim];
+              outDim++;
+            }
+          }
+          const xi = primalData[n];
+          const sign = xi > 0 ? 1 : xi < 0 ? -1 : 0;
+          thisGradSlot[n] += outputGrad[outFlat] * sign;
+          for (let k = inputShape.length - 1; k >= 0; k--) {
+            if (++inCoords[k] < inputShape[k]) break;
+            inCoords[k] = 0;
+          }
+        }
+      });
+      return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+    }
+    {
+      const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
+        const outShape = [...outT.shape];
+        const outStrides = _rowMajorStrides(outShape);
+        const argmaxAbsSlot = new Int32Array(outT.data.length).fill(-1);
+        const inCoords = new Array(inputShape.length).fill(0);
+        for (let n = 0; n < primalData.length; n++) {
+          let outFlat = 0;
+          let outDim = 0;
+          for (let k = 0; k < inputShape.length; k++) {
+            if (axisSet.has(k)) {
+              if (keepDims) outDim++;
+            } else {
+              outFlat += inCoords[k] * outStrides[outDim];
+              outDim++;
+            }
+          }
+          const prevWinner = argmaxAbsSlot[outFlat];
+          if (prevWinner === -1 || Math.abs(primalData[n]) > Math.abs(primalData[prevWinner])) {
+            argmaxAbsSlot[outFlat] = n;
+          }
+          for (let k = inputShape.length - 1; k >= 0; k--) {
+            if (++inCoords[k] < inputShape[k]) break;
+            inCoords[k] = 0;
+          }
+        }
+        for (let outFlat = 0; outFlat < outT.data.length; outFlat++) {
+          const winner = argmaxAbsSlot[outFlat];
+          if (winner !== -1) {
+            const xi = primalData[winner];
+            const sign = xi > 0 ? 1 : xi < 0 ? -1 : 0;
+            thisGradSlot[winner] += outputGrad[outFlat] * sign;
+          }
+        }
+      });
+      return new _TapedTensor([...outT.shape], new Float64Array(outT.data), this.tape, id);
+    }
+  }
+  // ---------------------------------------------------------------------------
+  // Elementwise transcendentals
+  // ---------------------------------------------------------------------------
+  /**
+   * Elementwise natural logarithm.
+   *
+   * Adjoint: dX = dY / x
+   */
+  log() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = Math.log(this.primal[i]);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] / primalData[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise exponential.
+   *
+   * Adjoint: dX = dY · y  where y = exp(x). Primal output is cached.
+   */
+  exp() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = Math.exp(this.primal[i]);
+    const expData = new Float64Array(out);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] * expData[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise sine.
+   *
+   * Adjoint: dX = dY · cos(x)
+   */
+  sin() {
+    const out = new Float64Array(this.primal.length);
+    const cosData = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) {
+      out[i] = Math.sin(this.primal[i]);
+      cosData[i] = Math.cos(this.primal[i]);
+    }
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] * cosData[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise cosine.
+   *
+   * Adjoint: dX = −dY · sin(x)
+   */
+  cos() {
+    const out = new Float64Array(this.primal.length);
+    const sinData = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) {
+      out[i] = Math.cos(this.primal[i]);
+      sinData[i] = Math.sin(this.primal[i]);
+    }
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += -outputGrad[i] * sinData[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise tangent.
+   *
+   * Adjoint: dX = dY / cos²(x)  (= dY · sec²(x))
+   */
+  tan() {
+    const out = new Float64Array(this.primal.length);
+    const cosData = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) {
+      out[i] = Math.tan(this.primal[i]);
+      cosData[i] = Math.cos(this.primal[i]);
+    }
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        const c = cosData[i];
+        thisGradSlot[i] += outputGrad[i] / (c * c);
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise square root.
+   *
+   * Adjoint: dX = dY / (2·y)  where y = sqrt(x). Primal output is cached.
+   */
+  sqrt() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = Math.sqrt(this.primal[i]);
+    const sqrtData = new Float64Array(out);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] / (2 * sqrtData[i]);
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise square (x²).
+   *
+   * Adjoint: dX = dY · 2x
+   */
+  square() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = this.primal[i] * this.primal[i];
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] * 2 * primalData[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise fixed-exponent power: x^k.
+   *
+   * Only fixed (non-TapedTensor) exponents are supported. Variable-exponent
+   * pow(taped, taped) is a follow-up slice.
+   *
+   * Adjoint: dX = dY · k · x^(k−1)
+   */
+  pow(k) {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = Math.pow(this.primal[i], k);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i] * k * Math.pow(primalData[i], k - 1);
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise reciprocal: 1 / x.
+   *
+   * Adjoint: dX = −dY / x²
+   */
+  reciprocal() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = 1 / this.primal[i];
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        const xi = primalData[i];
+        thisGradSlot[i] += -outputGrad[i] / (xi * xi);
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  /**
+   * Elementwise absolute value: |x|.
+   *
+   * Adjoint: dX = dY · sign(x)
+   * Subgradient at exact zero is defined as 0 (rather than undefined).
+   */
+  abs() {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = Math.abs(this.primal[i]);
+    const primalData = new Float64Array(this.primal);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        const xi = primalData[i];
+        const sign = xi > 0 ? 1 : xi < 0 ? -1 : 0;
+        thisGradSlot[i] += outputGrad[i] * sign;
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  // ---------------------------------------------------------------------------
+  // tensordot — generalised dot product over explicit axis pairs.
+  // ---------------------------------------------------------------------------
+  /**
+   * Reverse-mode AD over `Tensor.tensordot`.
+   *
+   * `axes[i] = [a, b]` contracts axis `a` of `this` with axis `b` of `other`.
+   * The result shape is `this`'s non-contracted axes (in original order)
+   * followed by `other`'s non-contracted axes (in original order).
+   *
+   * Adjoint derivation (NumPy/PyTorch tensordot backward, see Townsend 2016 §6,
+   * and the canonical PyTorch implementation `TensorDotBackward0` in
+   * `torch/csrc/autograd/generated/Functions.cpp`):
+   *
+   *   Z = tensordot(A, B, axes)
+   *   dA = tensordot(dZ, B, [axes_of_dZ_corresponding_to_B's_free, B's_free])
+   *        then permute back into A's original axis order.
+   *   dB = tensordot(A, dZ, [A's_free, axes_of_dZ_corresponding_to_A's_free])
+   *        then permute back into B's original axis order.
+   *
+   * The axis-permutation bookkeeping is the trickiest part: Tensor.tensordot
+   * produces output axes in the order [A's free axes (original A order),
+   * B's free axes (original B order)] — and after the backward contractions
+   * the survivors come out in B's original (resp. A's original) axis order
+   * for the contracted side, which then needs to be permuted into pair order
+   * (so axis k of the contracted block matches the kth pair) and finally
+   * scattered back into A's (resp. B's) full original axis order.
+   *
+   * For the rank-2 × rank-2 single-axis case (i.e. ordinary matmul A·B with
+   * axes = [[1, 0]]), this reduces to dA = dZ · Bᵀ, dB = Aᵀ · dZ — the same
+   * adjoint as `TapedTensor.matmul`.
+   */
+  tensordot(other, axes) {
+    const aShape = [...this.shape];
+    const bShape = [...other.shape];
+    const aRank = aShape.length;
+    const bRank = bShape.length;
+    const aContractedInPairOrder = [];
+    const bContractedInPairOrder = [];
+    for (const [a, b] of axes) {
+      aContractedInPairOrder.push(a);
+      bContractedInPairOrder.push(b);
+    }
+    const aContractedSet = new Set(aContractedInPairOrder);
+    const bContractedSet = new Set(bContractedInPairOrder);
+    const aFree = [];
+    for (let i = 0; i < aRank; i++) if (!aContractedSet.has(i)) aFree.push(i);
+    const bFree = [];
+    for (let i = 0; i < bRank; i++) if (!bContractedSet.has(i)) bFree.push(i);
+    const aPrimalT = new Tensor3(aShape, new Float64Array(this.primal));
+    const bPrimalT = new Tensor3(bShape, new Float64Array(other.primal));
+    const zPrimal = aPrimalT.tensordot(bPrimalT, axes);
+    const zShape = [...zPrimal.shape];
+    const naFree = aFree.length;
+    const nbFree = bFree.length;
+    const aPrimalData = new Float64Array(this.primal);
+    const bPrimalData = new Float64Array(other.primal);
+    const aGradSlot = this.tape.getInputGrad(this.id);
+    const bGradSlot = this.tape.getInputGrad(other.id);
+    const { id } = this.tape.record(
+      [this.id, other.id],
+      zPrimal.data.length,
+      (outputGrad) => {
+        const dZ = new Tensor3(zShape, new Float64Array(outputGrad));
+        const aPrimalTbk = new Tensor3(aShape, aPrimalData);
+        const bPrimalTbk = new Tensor3(bShape, bPrimalData);
+        const dAaxes = [];
+        for (let i = 0; i < nbFree; i++) {
+          dAaxes.push([naFree + i, bFree[i]]);
+        }
+        const dAPartial = dZ.tensordot(bPrimalTbk, dAaxes);
+        const bContractedSortedByOriginal = [...bContractedInPairOrder].sort((x, y) => x - y);
+        const nPairs = aContractedInPairOrder.length;
+        const tailPerm = [];
+        for (let k = 0; k < nPairs; k++) {
+          const bOrigAxis = bContractedInPairOrder[k];
+          const idx = bContractedSortedByOriginal.indexOf(bOrigAxis);
+          tailPerm.push(idx);
+        }
+        const dAFirstPerm = [];
+        for (let i = 0; i < naFree; i++) dAFirstPerm.push(i);
+        for (let k = 0; k < nPairs; k++) dAFirstPerm.push(naFree + tailPerm[k]);
+        const isFirstPermIdentity = dAFirstPerm.every((v, idx) => v === idx);
+        const dAReordered = isFirstPermIdentity ? dAPartial : dAPartial.transpose(dAFirstPerm);
+        const scatterA = new Array(aRank);
+        for (let a = 0; a < aRank; a++) {
+          const freeIdx = aFree.indexOf(a);
+          if (freeIdx !== -1) {
+            scatterA[a] = freeIdx;
+          } else {
+            const contIdx = aContractedInPairOrder.indexOf(a);
+            scatterA[a] = naFree + contIdx;
+          }
+        }
+        const isScatterAIdentity = scatterA.every((v, idx) => v === idx);
+        const dAFinal = isScatterAIdentity ? dAReordered : dAReordered.transpose(scatterA);
+        for (let i = 0; i < aGradSlot.length; i++) {
+          aGradSlot[i] += dAFinal.data[i];
+        }
+        const dBaxes = [];
+        for (let i = 0; i < naFree; i++) {
+          dBaxes.push([aFree[i], i]);
+        }
+        const dBPartial = aPrimalTbk.tensordot(dZ, dBaxes);
+        const aContractedSortedByOriginal = [...aContractedInPairOrder].sort((x, y) => x - y);
+        const dBHeadPerm = [];
+        for (let k = 0; k < nPairs; k++) {
+          const aOrigAxis = aContractedInPairOrder[k];
+          const idx = aContractedSortedByOriginal.indexOf(aOrigAxis);
+          dBHeadPerm.push(idx);
+        }
+        const dBFirstPerm = [];
+        for (let k = 0; k < nPairs; k++) dBFirstPerm.push(dBHeadPerm[k]);
+        for (let i = 0; i < nbFree; i++) dBFirstPerm.push(nPairs + i);
+        const isDBFirstPermIdentity = dBFirstPerm.every((v, idx) => v === idx);
+        const dBReordered = isDBFirstPermIdentity ? dBPartial : dBPartial.transpose(dBFirstPerm);
+        const scatterB = new Array(bRank);
+        for (let b = 0; b < bRank; b++) {
+          const freeIdx = bFree.indexOf(b);
+          if (freeIdx !== -1) {
+            scatterB[b] = nPairs + freeIdx;
+          } else {
+            const contIdx = bContractedInPairOrder.indexOf(b);
+            scatterB[b] = contIdx;
+          }
+        }
+        const isScatterBIdentity = scatterB.every((v, idx) => v === idx);
+        const dBFinal = isScatterBIdentity ? dBReordered : dBReordered.transpose(scatterB);
+        for (let i = 0; i < bGradSlot.length; i++) {
+          bGradSlot[i] += dBFinal.data[i];
+        }
+      }
+    );
+    return new _TapedTensor(zShape, new Float64Array(zPrimal.data), this.tape, id);
+  }
+  // ---------------------------------------------------------------------------
+  // SVD — full singular value decomposition on rank-2 input.
+  // ---------------------------------------------------------------------------
+  /**
+   * Reverse-mode AD over the full SVD of a rank-2 matrix.
+   *
+   * Forward: `A = U · diag(S) · Vt`, where for input shape [m, n], k = min(m, n):
+   *   - U  has shape [m, k]
+   *   - S  has shape [k]
+   *   - Vt has shape [k, n]    (Vt is V^T in the standard A = U Σ V^T convention,
+   *                              i.e. its rows are right-singular-vector components)
+   *
+   * Returned TapedTensors share a single backward closure. When backward()
+   * runs, it pulls dU, dS, dV from each output's gradient slot, assembles
+   * dA, and writes to the input's gradient slot.
+   *
+   * Adjoint (real, distinct nonzero singular values, m = n square case;
+   * extended to rectangular below). Derived directly from the forward
+   * Jacobian; equivalent to PyTorch's `svd_backward`
+   * (`aten/src/ATen/native/BatchLinearAlgebra.cpp`) and Townsend (2016)
+   * "Differentiating the Singular Value Decomposition" §3:
+   *
+   *   Let α = skew(U^T · dU),  β = skew(V^T · dV)        (k×k, antisymmetric)
+   *   Build C (k×k):
+   *     C[i,i] = dS[i]
+   *     C[i,j] = (α[i,j] + β[i,j]) / (s_j − s_i)
+   *            + (α[i,j] − β[i,j]) / (s_j + s_i)         for i ≠ j
+   *   dA_in = U · C · V^T                                 (m×n in-subspace part)
+   *
+   * Rectangular correction (when m > k, i.e. m > n):
+   *   dA += (I − U U^T) · dU · diag(1/s) · V^T
+   * Rectangular correction (when n > k, i.e. n > m):
+   *   dA += U · diag(1/s) · dV^T · (I − V V^T)
+   *
+   * Regularisation at repeated/near-zero singular values (PyTorch-equivalent
+   * subgradient choice): the (i,j) entry of C is masked to 0 whenever
+   * `|s_j − s_i| < REL_TOL · max(|s|)` (the "difference" denominator) or
+   * `|s_j + s_i| < REL_TOL · max(|s|)` (the "sum" denominator, only relevant
+   * when both are ~0). REL_TOL = 1e-10. This makes the gradient a subgradient
+   * at exact degeneracy — finite, but not the unique true derivative (which
+   * does not exist at degeneracies). The rectangular correction also masks
+   * 1/s_i when |s_i| < REL_TOL · max(|s|).
+   *
+   * Throws if input is not rank-2. For rank > 2 inputs the user should
+   * reshape first.
+   */
+  svd() {
+    if (this.shape.length !== 2) {
+      throw new Error(
+        `TapedTensor.svd: expected rank-2 input, got rank ${this.shape.length}. Reshape the tensor to 2-D first.`
+      );
+    }
+    const m = this.shape[0];
+    const n = this.shape[1];
+    const k = Math.min(m, n);
+    const inputT = new Tensor3([m, n], new Float64Array(this.primal));
+    const svdResult = tensorSvd(inputT, [0]);
+    if (svdResult.truncatedDim !== k) {
+      throw new Error(
+        `TapedTensor.svd: internal \u2014 tensorSvd returned truncatedDim=${svdResult.truncatedDim}, expected ${k}. (A zero singular value was dropped by the default cutoff?)`
+      );
+    }
+    const uPrimal = new Float64Array(svdResult.U.data);
+    const sPrimal = new Float64Array(svdResult.S.data);
+    const vtPrimal = new Float64Array(svdResult.V.data);
+    const dUBuf = new Float64Array(m * k);
+    const dSBuf = new Float64Array(k);
+    const dVtBuf = new Float64Array(k * n);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const sharedBackward = () => {
+      const dA = _svdBackward(uPrimal, sPrimal, vtPrimal, dUBuf, dSBuf, dVtBuf, m, n, k);
+      for (let i = 0; i < thisGradSlot.length; i++) {
+        thisGradSlot[i] += dA[i];
+      }
+    };
+    const { id: uId } = this.tape.record([this.id], m * k, (outputGrad) => {
+      for (let i = 0; i < dUBuf.length; i++) dUBuf[i] += outputGrad[i];
+      sharedBackward();
+    });
+    const { id: sId } = this.tape.record([this.id], k, (outputGrad) => {
+      for (let i = 0; i < dSBuf.length; i++) dSBuf[i] += outputGrad[i];
+    });
+    const { id: vId } = this.tape.record([this.id], k * n, (outputGrad) => {
+      for (let i = 0; i < dVtBuf.length; i++) dVtBuf[i] += outputGrad[i];
+    });
+    const uTape = new _TapedTensor([m, k], uPrimal, this.tape, uId);
+    const sTape = new _TapedTensor([k], sPrimal, this.tape, sId);
+    const vTape = new _TapedTensor([k, n], vtPrimal, this.tape, vId);
+    return { U: uTape, S: sTape, V: vTape };
+  }
+  // ---------------------------------------------------------------------------
+  // Symmetric eigendecomposition.
+  // ---------------------------------------------------------------------------
+  /**
+   * Reverse-mode AD over the eigendecomposition of a rank-2 matrix.
+   *
+   * Symmetric path (`symmetric: true`):
+   *   Forward: `A = U · diag(Λ) · U^T` for symmetric A (n×n).
+   *   Adjoint (Magnus & Neudecker 1999 §10.6.6; PyTorch `linalg_eigh_backward`):
+   *     F[i,j] = 1 / (λ_j − λ_i)   for i ≠ j, 0 otherwise (with degeneracy mask)
+   *     dA_raw = U · (diag(dΛ) + F ∘ (U^T · dU)) · U^T
+   *     dA     = (dA_raw + dA_raw^T) / 2          (symmetrise)
+   *
+   * Non-symmetric path (`symmetric: false`):
+   *   Forward: `A = V · diag(λ) · V^{-1}` (V columns are right eigenvectors).
+   *   Adjoint (Magnus & Neudecker 1999 §10.6 / Giles 2008 §3.2 / Townsend 2016 §4;
+   *   cross-check: PyTorch `linalg_eig_backward`):
+   *     E[i,j] = 1 / (λ_j − λ_i)  for i ≠ j, 0 otherwise (with degeneracy mask)
+   *     dA = V^{-T} · ( E ∘ (V^T · dV) + diag(dλ) ) · V^T
+   *
+   *   Restrictions (all enforced — throw a clear error otherwise):
+   *   1. Eigenvalues must be real. The underlying matrix-eig primitive returns
+   *      placeholder eigenvectors (not actual complex vectors) when complex
+   *      eigenvalues arise, so the adjoint formula cannot be evaluated. Real-
+   *      Schur differentiation would require complex arithmetic infrastructure
+   *      throughout the Tape/TapedTensor stack, which is out of scope.
+   *   2. A must be diagonalisable (non-defective). The adjoint assumes V is
+   *      invertible; defective inputs have algebraic > geometric multiplicity
+   *      so V is rank-deficient. Detected by cond_∞(V) > 1e14.
+   *
+   * Regularisation at repeated eigenvalues (subgradient choice): mask
+   * F/E[i,j] = 0 when `|λ_i − λ_j| < REL_TOL · max(|λ|)`. REL_TOL = 1e-10.
+   *
+   * Throws if input is not rank-2 or square, if `symmetric` is missing, or
+   * (non-symmetric path) on complex eigenvalues / defective input.
+   */
+  eig(opts) {
+    if (!opts || typeof opts.symmetric !== "boolean") {
+      throw new Error(
+        "TapedTensor.eig: opts.symmetric must be set to true or false. Pass { symmetric: true } for symmetric A, { symmetric: false } for general A."
+      );
+    }
+    if (this.shape.length !== 2 || this.shape[0] !== this.shape[1]) {
+      throw new Error(`TapedTensor.eig: expected rank-2 square input, got shape [${this.shape}].`);
+    }
+    const n = this.shape[0];
+    if (opts.symmetric === true) {
+      return this._eigSymmetric(n);
+    }
+    return this._eigGeneral(n);
+  }
+  _eigSymmetric(n) {
+    const inputT = new Tensor3([n, n], new Float64Array(this.primal));
+    const eigResult = tensorEig(inputT, [0], { symmetric: true, computeVectors: true });
+    const lambdaPrimal = new Float64Array(eigResult.eigenvalues.data);
+    if (!eigResult.eigenvectors) {
+      throw new Error("TapedTensor.eig: tensorEig did not return eigenvectors");
+    }
+    const uPrimal = new Float64Array(eigResult.eigenvectors.data);
+    const dLambdaBuf = new Float64Array(n);
+    const dUBuf = new Float64Array(n * n);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const sharedBackward = () => {
+      const dA = _eigSymBackward(uPrimal, lambdaPrimal, dUBuf, dLambdaBuf, n);
+      for (let i = 0; i < thisGradSlot.length; i++) {
+        thisGradSlot[i] += dA[i];
+      }
+    };
+    const { id: lId } = this.tape.record([this.id], n, (outputGrad) => {
+      for (let i = 0; i < dLambdaBuf.length; i++) dLambdaBuf[i] += outputGrad[i];
+      sharedBackward();
+    });
+    const { id: uId } = this.tape.record([this.id], n * n, (outputGrad) => {
+      for (let i = 0; i < dUBuf.length; i++) dUBuf[i] += outputGrad[i];
+    });
+    const lTape = new _TapedTensor([n], lambdaPrimal, this.tape, lId);
+    const uTape = new _TapedTensor([n, n], uPrimal, this.tape, uId);
+    return { eigvals: lTape, eigvecs: uTape };
+  }
+  _eigGeneral(n) {
+    const inputT = new Tensor3([n, n], new Float64Array(this.primal));
+    const eigResult = tensorEig(inputT, [0], { symmetric: false, computeVectors: true });
+    if (eigResult.eigenvaluesImaginary) {
+      const imData = eigResult.eigenvaluesImaginary.data;
+      for (let i = 0; i < n; i++) {
+        if (imData[i] !== 0) {
+          throw new Error(
+            "TapedTensor.eig({ symmetric: false }): input has complex eigenvalues. Reverse-mode AD over the general non-symmetric eig is restricted to real eigenvalues (the underlying matrix primitive does not return complex eigenvectors). Either symmetrise the input, or restrict to inputs with a real spectrum."
+          );
+        }
+      }
+    }
+    if (!eigResult.eigenvectors) {
+      throw new Error("TapedTensor.eig: tensorEig did not return eigenvectors");
+    }
+    const lambdaPrimal = new Float64Array(eigResult.eigenvalues.data);
+    const vPrimal = new Float64Array(eigResult.eigenvectors.data);
+    const vInvOrNull = _invertWithCondCheck(vPrimal, n, 1e14);
+    if (vInvOrNull === null) {
+      throw new Error(
+        "TapedTensor.eig({ symmetric: false }): input appears defective (eigenvector matrix V has cond_\u221E(V) > 1e14). The general-case adjoint requires diagonalisability. Consider perturbing the input or using the symmetric path if A is symmetric."
+      );
+    }
+    const vInv = vInvOrNull;
+    const dLambdaBuf = new Float64Array(n);
+    const dVBuf = new Float64Array(n * n);
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const sharedBackward = () => {
+      const dA = _eigGeneralBackward(vPrimal, vInv, lambdaPrimal, dVBuf, dLambdaBuf, n);
+      for (let i = 0; i < thisGradSlot.length; i++) {
+        thisGradSlot[i] += dA[i];
+      }
+    };
+    const { id: lId } = this.tape.record([this.id], n, (outputGrad) => {
+      for (let i = 0; i < dLambdaBuf.length; i++) dLambdaBuf[i] += outputGrad[i];
+      sharedBackward();
+    });
+    const { id: vId } = this.tape.record([this.id], n * n, (outputGrad) => {
+      for (let i = 0; i < dVBuf.length; i++) dVBuf[i] += outputGrad[i];
+    });
+    const lTape = new _TapedTensor([n], lambdaPrimal, this.tape, lId);
+    const vTape = new _TapedTensor([n, n], vPrimal, this.tape, vId);
+    return { eigvals: lTape, eigvecs: vTape };
+  }
 };
+function _svdBackward(U, s, Vt, dU, dS, dVt, m, n, k) {
+  const REL_TOL = 1e-10;
+  let sMax = 0;
+  for (let i = 0; i < k; i++) {
+    const a = Math.abs(s[i]);
+    if (a > sMax) sMax = a;
+  }
+  const tolAbs = REL_TOL * (sMax > 0 ? sMax : 1);
+  const V = new Float64Array(n * k);
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < k; j++) {
+      V[i * k + j] = Vt[j * n + i];
+    }
+  }
+  const dV = new Float64Array(n * k);
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < k; j++) {
+      dV[i * k + j] = dVt[j * n + i];
+    }
+  }
+  const UtdU = _matmulMNK(
+    U,
+    dU,
+    m,
+    k,
+    k,
+    /*aT=*/
+    true,
+    /*bT=*/
+    false
+  );
+  const VtdV = _matmulMNK(
+    V,
+    dV,
+    n,
+    k,
+    k,
+    /*aT=*/
+    true,
+    /*bT=*/
+    false
+  );
+  const C = new Float64Array(k * k);
+  for (let i = 0; i < k; i++) {
+    for (let j = 0; j < k; j++) {
+      if (i === j) {
+        C[i * k + j] = dS[i];
+      } else {
+        const alpha = (UtdU[i * k + j] - UtdU[j * k + i]) * 0.5;
+        const beta = (VtdV[i * k + j] - VtdV[j * k + i]) * 0.5;
+        const denomDiff = s[j] - s[i];
+        const denomSum = s[j] + s[i];
+        let term = 0;
+        if (Math.abs(denomDiff) > tolAbs) {
+          term += (alpha + beta) / denomDiff;
+        }
+        if (Math.abs(denomSum) > tolAbs) {
+          term += (alpha - beta) / denomSum;
+        }
+        C[i * k + j] = term;
+      }
+    }
+  }
+  const UC = _matmulMNK(U, C, m, k, k, false, false);
+  const dA = _matmulMNK(UC, Vt, m, k, n, false, false);
+  if (m > k) {
+    const UUtdU = _matmulMNK(U, UtdU, m, k, k, false, false);
+    const dUperp = new Float64Array(m * k);
+    for (let i = 0; i < m * k; i++) dUperp[i] = dU[i] - UUtdU[i];
+    for (let i = 0; i < m; i++) {
+      for (let j = 0; j < k; j++) {
+        const sj = s[j];
+        if (Math.abs(sj) > tolAbs) {
+          dUperp[i * k + j] /= sj;
+        } else {
+          dUperp[i * k + j] = 0;
+        }
+      }
+    }
+    const corr = _matmulMNK(dUperp, Vt, m, k, n, false, false);
+    for (let i = 0; i < dA.length; i++) dA[i] += corr[i];
+  }
+  if (n > k) {
+    const VVtdV = _matmulMNK(V, VtdV, n, k, k, false, false);
+    const dVperp = new Float64Array(n * k);
+    for (let i = 0; i < n * k; i++) dVperp[i] = dV[i] - VVtdV[i];
+    for (let i = 0; i < n; i++) {
+      for (let j = 0; j < k; j++) {
+        const sj = s[j];
+        if (Math.abs(sj) > tolAbs) {
+          dVperp[i * k + j] /= sj;
+        } else {
+          dVperp[i * k + j] = 0;
+        }
+      }
+    }
+    const corr = _matmulMNK(U, dVperp, m, k, n, false, true);
+    for (let i = 0; i < dA.length; i++) dA[i] += corr[i];
+  }
+  return dA;
+}
+function _eigSymBackward(U, lam, dU, dLam, n) {
+  const REL_TOL = 1e-10;
+  let lamMax = 0;
+  for (let i = 0; i < n; i++) {
+    const a = Math.abs(lam[i]);
+    if (a > lamMax) lamMax = a;
+  }
+  const tolAbs = REL_TOL * (lamMax > 0 ? lamMax : 1);
+  const UtdU = _matmulMNK(U, dU, n, n, n, true, false);
+  const M = new Float64Array(n * n);
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      if (i === j) {
+        M[i * n + j] = dLam[i];
+      } else {
+        const denom = lam[j] - lam[i];
+        if (Math.abs(denom) > tolAbs) {
+          M[i * n + j] = UtdU[i * n + j] / denom;
+        } else {
+          M[i * n + j] = 0;
+        }
+      }
+    }
+  }
+  const UM = _matmulMNK(U, M, n, n, n, false, false);
+  const dAraw = _matmulMNK(UM, U, n, n, n, false, true);
+  const dA = new Float64Array(n * n);
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      dA[i * n + j] = (dAraw[i * n + j] + dAraw[j * n + i]) * 0.5;
+    }
+  }
+  return dA;
+}
+function _eigGeneralBackward(V, vInv, lam, dV, dLam, n) {
+  const REL_TOL = 1e-10;
+  let lamMax = 0;
+  for (let i = 0; i < n; i++) {
+    const a = Math.abs(lam[i]);
+    if (a > lamMax) lamMax = a;
+  }
+  const tolAbs = REL_TOL * (lamMax > 0 ? lamMax : 1);
+  const VtdV = _matmulMNK(V, dV, n, n, n, true, false);
+  const VtV = _matmulMNK(V, V, n, n, n, true, false);
+  const diagVtdV = new Float64Array(n);
+  for (let i = 0; i < n; i++) diagVtdV[i] = VtdV[i * n + i];
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      VtdV[i * n + j] -= VtV[i * n + j] * diagVtdV[j];
+    }
+  }
+  const M = new Float64Array(n * n);
+  for (let i = 0; i < n; i++) {
+    for (let j = 0; j < n; j++) {
+      if (i === j) {
+        M[i * n + j] = dLam[i];
+      } else {
+        const denom = lam[j] - lam[i];
+        if (Math.abs(denom) > tolAbs) {
+          M[i * n + j] = VtdV[i * n + j] / denom;
+        } else {
+          M[i * n + j] = 0;
+        }
+      }
+    }
+  }
+  const VinvT_M = _matmulMNK(vInv, M, n, n, n, true, false);
+  const dA = _matmulMNK(VinvT_M, V, n, n, n, false, true);
+  return dA;
+}
+function _invertWithCondCheck(V, n, maxCond) {
+  const LU = new Float64Array(V);
+  const perm = new Int32Array(n);
+  for (let i = 0; i < n; i++) perm[i] = i;
+  let pivotEpsScale = 0;
+  for (let i = 0; i < n * n; i++) {
+    const a = Math.abs(V[i]);
+    if (a > pivotEpsScale) pivotEpsScale = a;
+  }
+  const singularEps = pivotEpsScale === 0 ? 0 : pivotEpsScale * 1e-300;
+  for (let k = 0; k < n; k++) {
+    let maxVal = Math.abs(LU[k * n + k]);
+    let maxIdx = k;
+    for (let i = k + 1; i < n; i++) {
+      const v = Math.abs(LU[i * n + k]);
+      if (v > maxVal) {
+        maxVal = v;
+        maxIdx = i;
+      }
+    }
+    if (maxIdx !== k) {
+      for (let j = 0; j < n; j++) {
+        const tmp = LU[k * n + j];
+        LU[k * n + j] = LU[maxIdx * n + j];
+        LU[maxIdx * n + j] = tmp;
+      }
+      const t = perm[k];
+      perm[k] = perm[maxIdx];
+      perm[maxIdx] = t;
+    }
+    const pivot = LU[k * n + k];
+    if (!Number.isFinite(pivot) || Math.abs(pivot) <= singularEps) {
+      return null;
+    }
+    for (let i = k + 1; i < n; i++) {
+      LU[i * n + k] /= pivot;
+      const lik = LU[i * n + k];
+      for (let j = k + 1; j < n; j++) {
+        LU[i * n + j] -= lik * LU[k * n + j];
+      }
+    }
+  }
+  const inv = new Float64Array(n * n);
+  const y = new Float64Array(n);
+  for (let col = 0; col < n; col++) {
+    for (let i = 0; i < n; i++) {
+      y[i] = perm[i] === col ? 1 : 0;
+    }
+    for (let i = 0; i < n; i++) {
+      let s = y[i];
+      for (let j = 0; j < i; j++) s -= LU[i * n + j] * y[j];
+      y[i] = s;
+    }
+    for (let i = n - 1; i >= 0; i--) {
+      let s = y[i];
+      for (let j = i + 1; j < n; j++) s -= LU[i * n + j] * inv[j * n + col];
+      inv[i * n + col] = s / LU[i * n + i];
+    }
+  }
+  let normV = 0;
+  let normInv = 0;
+  for (let i = 0; i < n; i++) {
+    let rowV = 0;
+    let rowI = 0;
+    for (let j = 0; j < n; j++) {
+      rowV += Math.abs(V[i * n + j]);
+      const inv_ij = inv[i * n + j];
+      if (!Number.isFinite(inv_ij)) return null;
+      rowI += Math.abs(inv_ij);
+    }
+    if (rowV > normV) normV = rowV;
+    if (rowI > normInv) normInv = rowI;
+  }
+  const cond = normV * normInv;
+  if (!Number.isFinite(cond) || cond > maxCond) return null;
+  return inv;
+}
+function _matmulMNK(A, B, M, K, N, aT, bT) {
+  const C = new Float64Array(M * N);
+  for (let i = 0; i < M; i++) {
+    for (let j = 0; j < N; j++) {
+      let sum = 0;
+      for (let q = 0; q < K; q++) {
+        const aVal = aT ? A[q * M + i] : A[i * K + q];
+        const bVal = bT ? B[j * K + q] : B[q * N + j];
+        sum += aVal * bVal;
+      }
+      C[i * N + j] = sum;
+    }
+  }
+  return C;
+}
+function _resolveAxes(axis, rank) {
+  if (axis === void 0) {
+    return Array.from({ length: rank }, (_, i) => i);
+  }
+  if (typeof axis === "number") {
+    return [axis];
+  }
+  return [...axis].sort((a, b) => a - b);
+}
+function _rowMajorStrides(shape) {
+  const strides = new Array(shape.length);
+  let acc = 1;
+  for (let k = shape.length - 1; k >= 0; k--) {
+    strides[k] = acc;
+    acc *= shape[k];
+  }
+  return strides;
+}
+function _scatterToOriginalAxes(tensor, targetLabels) {
+  const currentLabels = tensor.axisLabels;
+  const rank = targetLabels.length;
+  const perm = new Array(rank);
+  for (let i = 0; i < rank; i++) {
+    const j = currentLabels.findIndex((lbl) => lbl.matches(targetLabels[i]));
+    if (j === -1) {
+      throw new Error(
+        `_scatterToOriginalAxes: target label ${targetLabels[i].toString()} not found in tensor axisLabels`
+      );
+    }
+    perm[i] = j;
+  }
+  const isIdentity = perm.every((j, i) => j === i);
+  if (isIdentity) {
+    return new Float64Array(tensor.data);
+  }
+  const permuted = tensor.transpose(perm);
+  return new Float64Array(permuted.data);
+}
+function _batchedMatmul(aData, bData, aShape, bShape) {
+  const rank = aShape.length;
+  const batchRank = rank - 2;
+  const m = aShape[rank - 2];
+  const k = aShape[rank - 1];
+  const n = bShape[rank - 1];
+  let batchSize = 1;
+  for (let b = 0; b < batchRank; b++) batchSize *= aShape[b];
+  const yData = new Float64Array(batchSize * m * n);
+  for (let batch = 0; batch < batchSize; batch++) {
+    const aOff = batch * m * k;
+    const bOff = batch * k * n;
+    const yOff = batch * m * n;
+    for (let i = 0; i < m; i++) {
+      for (let j = 0; j < n; j++) {
+        let sum = 0;
+        for (let q = 0; q < k; q++) {
+          sum += aData[aOff + i * k + q] * bData[bOff + q * n + j];
+        }
+        yData[yOff + i * n + j] = sum;
+      }
+    }
+  }
+  return yData;
+}
+function _batchedMatmulGradA(dYData, bData, yShape, bShape) {
+  const rank = yShape.length;
+  const batchRank = rank - 2;
+  const m = yShape[rank - 2];
+  const n = yShape[rank - 1];
+  const k = bShape[rank - 2];
+  let batchSize = 1;
+  for (let b = 0; b < batchRank; b++) batchSize *= yShape[b];
+  const dAData = new Float64Array(batchSize * m * k);
+  for (let batch = 0; batch < batchSize; batch++) {
+    const dYOff = batch * m * n;
+    const bOff = batch * k * n;
+    const dAOff = batch * m * k;
+    for (let i = 0; i < m; i++) {
+      for (let q = 0; q < k; q++) {
+        let sum = 0;
+        for (let j = 0; j < n; j++) {
+          sum += dYData[dYOff + i * n + j] * bData[bOff + q * n + j];
+        }
+        dAData[dAOff + i * k + q] = sum;
+      }
+    }
+  }
+  return dAData;
+}
+function _batchedMatmulGradB(aData, dYData, aShape, yShape) {
+  const rank = aShape.length;
+  const batchRank = rank - 2;
+  const m = aShape[rank - 2];
+  const k = aShape[rank - 1];
+  const n = yShape[rank - 1];
+  let batchSize = 1;
+  for (let b = 0; b < batchRank; b++) batchSize *= aShape[b];
+  const dBData = new Float64Array(batchSize * k * n);
+  for (let batch = 0; batch < batchSize; batch++) {
+    const aOff = batch * m * k;
+    const dYOff = batch * m * n;
+    const dBOff = batch * k * n;
+    for (let q = 0; q < k; q++) {
+      for (let j = 0; j < n; j++) {
+        let sum = 0;
+        for (let i = 0; i < m; i++) {
+          sum += aData[aOff + i * k + q] * dYData[dYOff + i * n + j];
+        }
+        dBData[dBOff + q * n + j] = sum;
+      }
+    }
+  }
+  return dBData;
+}
 
 // src/reverse-grad.ts
 import { Tensor as Tensor4 } from "@danielsimonjr/mathts-tensor";