@danielsimonjr/mathts-autograd 0.1.0

package/dist/index.d.ts

@@ -0,0 +1,168 @@
+import { Tensor } from '@danielsimonjr/mathts-tensor';
+
+/**
+ * DualTensor — a Tensor + per-element tangent component for forward-mode AD.
+ * Storage: two Float64Arrays of equal length (primal + tangent), shape from
+ * the wrapped Tensor. Arithmetic follows the dual-number rules:
+ *   (a, a') + (b, b') = (a+b, a'+b')
+ *   (a, a') * (b, b') = (a·b, a·b' + a'·b)
+ *   (a, a') / (b, b') = (a/b, (a'·b − a·b')/b²)
+ *   scale((a, a'), k)  = (k·a, k·a')
+ * Elementwise unless noted. Reductions and contractions add their own rules.
+ *
+ * The dual-number framework: tracking ε such that ε²=0; (a + a'·ε)(b + b'·ε)
+ * = ab + (ab'+a'b)ε, so the tangent is the linearized first-order response.
+ *
+ * @packageDocumentation
+ */
+
+declare class DualTensor {
+    readonly shape: ReadonlyArray<number>;
+    readonly primal: Float64Array;
+    readonly tangent: Float64Array;
+    constructor(shape: ReadonlyArray<number>, primal: Float64Array, tangent: Float64Array);
+    /**
+     * 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
+     * DualTensor flows through them as a structurally-compatible Tensor.
+     * (AD-aware ops branch on `'tangent' in arg` before reaching here.)
+     */
+    get data(): Float64Array;
+    /** Lift a Tensor to a DualTensor with zero tangent. */
+    static fromTensor(t: Tensor): DualTensor;
+    /** Lift a Tensor to a DualTensor with a unit tangent at flat-index `i`. */
+    static unitAt(t: Tensor, i: number): DualTensor;
+    /** Extract the primal as a plain Tensor. */
+    toPrimalTensor(): Tensor;
+    /** Extract the tangent as a plain Tensor. */
+    toTangentTensor(): Tensor;
+    /** Elementwise addition following dual-number rule: (a+b, a'+b'). */
+    add(other: DualTensor): DualTensor;
+    /** Elementwise subtraction following dual-number rule: (a-b, a'-b'). */
+    sub(other: DualTensor): DualTensor;
+    /**
+     * Elementwise multiplication following dual-number rule: (a·b, a·b' + a'·b).
+     *
+     * I3 fix: explicit alias check. When `this === other` (self-multiplication),
+     * use the specialised rule (a·a)' = 2·a·a' directly. The general rule also
+     * degenerates correctly, but the explicit branch is more readable and aligns
+     * the forward-mode implementation with the reverse-mode alias fix.
+     */
+    mul(other: DualTensor): DualTensor;
+    /** Scalar multiplication following dual-number rule: (k·a, k·a'). */
+    scale(k: number): DualTensor;
+    private checkSameShape;
+}
+
+/**
+ * Forward-mode AD via dual numbers. Returns the value and the full Jacobian
+ * of fn at x, shape [...value.shape, ...x.shape] (row-major flatten).
+ *
+ * Implementation: for each flat-index k in x, build a DualTensor with a unit
+ * tangent at position k, run fn (which operates on Tensors — see "interop"
+ * below), extract the tangent of the result. The tangent is the k-th column
+ * of the Jacobian (in row-major-of-output-flattened form).
+ *
+ * Interop with the Tensor type: fn is typed `(x: Tensor) => Tensor`, so the
+ * tracing requires fn's body to be implementable on either Tensor OR
+ * DualTensor. v0.1.0 ships a documented "trace mode" — if fn uses only the
+ * core ops that DualTensor implements (add/sub/mul/scale), the
+ * implementation auto-traces by detecting DualTensor inputs via a thin
+ * wrapper. The actual mechanism in v0.1.0: caller writes fn as if Tensor's
+ * methods accept either Tensor or DualTensor (TypeScript's structural typing
+ * makes this work for the implemented ops). v0.2.0 may broaden via decorator
+ * instrumentation.
+ *
+ * For UPT v0.4.0's call-sites (metric functions in christoffel lowering),
+ * fn uses only add/sub/mul/scale → safe.
+ *
+ * @packageDocumentation
+ */
+
+/**
+ * Compute the value and full Jacobian of `fn` at point `x` using forward-mode
+ * automatic differentiation (dual numbers).
+ *
+ * @param fn - Pure function operating on Tensors via add/sub/mul/scale.
+ *   Must be traceable: at runtime it receives a DualTensor (structurally
+ *   compatible with Tensor) and must return a DualTensor.
+ * @param x - Input tensor at which to differentiate.
+ * @returns `{ value, jacobian }` where:
+ *   - `value` has the output shape of fn applied to x
+ *   - `jacobian` has shape `[...value.shape, ...x.shape]` row-major.
+ *     `jacobian.data[kY * xSize + kX]` = ∂y[kY] / ∂x[kX].
+ *     Verified by Adam+Eve review 2026-05-15 (E13): kY outer, kX inner,
+ *     stride xSize between successive y-rows IS row-major [...y.shape, ...x.shape].
+ */
+declare function forwardGrad(fn: (x: Tensor) => Tensor, x: Tensor): {
+    value: Tensor;
+    jacobian: Tensor;
+};
+
+/**
+ * Tape — records the sequence of ops during a forward pass so we can
+ * replay them in reverse to compute the vector-Jacobian product.
+ *
+ * Each TapedTensor wraps a primal Tensor + a tape-node id. Ops on
+ * TapedTensors record their backward closures into a shared Tape.
+ * After the forward pass, calling Tape.backward(outputGrad) walks the
+ * tape in reverse, accumulating gradients into each input slot.
+ *
+ * v0.1.0 supports the same ops as DualTensor: add, sub, mul, scale.
+ */
+
+type BackwardFn = (outputGrad: Float64Array) => void;
+declare class Tape {
+    private nodes;
+    private inputGradSlots;
+    private nextOpId;
+    private nextInputId;
+    /** Allocate a fresh id for an input or intermediate. */
+    allocate(size: number): {
+        id: number;
+        gradSlot: Float64Array;
+    };
+    record(inputIds: ReadonlyArray<number>, outputSize: number, backward: BackwardFn): {
+        id: number;
+        gradSlot: Float64Array;
+    };
+    /** Seed the final output's gradient slot, then replay in reverse. */
+    backward(outputId: number, outputGrad: Float64Array): void;
+    getInputGrad(id: number): Float64Array | undefined;
+}
+declare class TapedTensor {
+    readonly shape: ReadonlyArray<number>;
+    readonly primal: Float64Array;
+    readonly tape: Tape;
+    readonly id: number;
+    constructor(shape: ReadonlyArray<number>, primal: Float64Array, tape: Tape, id: number);
+    /**
+     * 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
+     * TapedTensor flows through them as a structurally-compatible Tensor.
+     * (AD-aware ops branch on `'tape' in arg` before reaching here.)
+     */
+    get data(): Float64Array;
+    static fromTensorAsInput(t: Tensor, tape: Tape): TapedTensor;
+    toPrimalTensor(): Tensor;
+    add(other: TapedTensor): TapedTensor;
+    sub(other: TapedTensor): TapedTensor;
+    mul(other: TapedTensor): TapedTensor;
+    scale(k: number): TapedTensor;
+    private checkSameShape;
+}
+
+/**
+ * Reverse-mode AD via tape. Returns the value and the vector-Jacobian product
+ * (gradient = ∂(cotangent · value) / ∂x). gradient.shape = x.shape.
+ *
+ * Default cotangent: ones-like(value) — required for scalar outputs, useful
+ * for VJP defaults. For non-scalar value, cotangent's shape must match value.shape.
+ */
+
+declare function reverseGrad(fn: (x: Tensor) => Tensor, x: Tensor, cotangent?: Tensor): {
+    value: Tensor;
+    gradient: Tensor;
+};
+
+export { DualTensor, Tape, TapedTensor, forwardGrad, reverseGrad };

package/dist/index.js

@@ -0,0 +1,305 @@
+// src/dual-tensor.ts
+import { Tensor } from "@danielsimonjr/mathts-tensor";
+var DualTensor = class _DualTensor {
+  shape;
+  primal;
+  tangent;
+  constructor(shape, primal, tangent) {
+    if (primal.length !== tangent.length) {
+      throw new Error(`DualTensor: primal length ${primal.length} != tangent length ${tangent.length}`);
+    }
+    this.shape = shape;
+    this.primal = primal;
+    this.tangent = tangent;
+  }
+  /**
+   * 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
+   * DualTensor flows through them as a structurally-compatible Tensor.
+   * (AD-aware ops branch on `'tangent' in arg` before reaching here.)
+   */
+  get data() {
+    return this.primal;
+  }
+  /** Lift a Tensor to a DualTensor with zero tangent. */
+  static fromTensor(t) {
+    return new _DualTensor(t.shape, new Float64Array(t.data), new Float64Array(t.data.length));
+  }
+  /** Lift a Tensor to a DualTensor with a unit tangent at flat-index `i`. */
+  static unitAt(t, i) {
+    const tan = new Float64Array(t.data.length);
+    tan[i] = 1;
+    return new _DualTensor(t.shape, new Float64Array(t.data), tan);
+  }
+  /** Extract the primal as a plain Tensor. */
+  toPrimalTensor() {
+    return new Tensor(this.shape, new Float64Array(this.primal));
+  }
+  /** Extract the tangent as a plain Tensor. */
+  toTangentTensor() {
+    return new Tensor(this.shape, new Float64Array(this.tangent));
+  }
+  /** Elementwise addition following dual-number rule: (a+b, a'+b'). */
+  add(other) {
+    this.checkSameShape(other, "add");
+    const p = new Float64Array(this.primal.length);
+    const t = new Float64Array(this.tangent.length);
+    for (let i = 0; i < p.length; i++) {
+      p[i] = this.primal[i] + other.primal[i];
+      t[i] = this.tangent[i] + other.tangent[i];
+    }
+    return new _DualTensor(this.shape, p, t);
+  }
+  /** Elementwise subtraction following dual-number rule: (a-b, a'-b'). */
+  sub(other) {
+    this.checkSameShape(other, "sub");
+    const p = new Float64Array(this.primal.length);
+    const t = new Float64Array(this.tangent.length);
+    for (let i = 0; i < p.length; i++) {
+      p[i] = this.primal[i] - other.primal[i];
+      t[i] = this.tangent[i] - other.tangent[i];
+    }
+    return new _DualTensor(this.shape, p, t);
+  }
+  /**
+   * Elementwise multiplication following dual-number rule: (a·b, a·b' + a'·b).
+   *
+   * I3 fix: explicit alias check. When `this === other` (self-multiplication),
+   * use the specialised rule (a·a)' = 2·a·a' directly. The general rule also
+   * degenerates correctly, but the explicit branch is more readable and aligns
+   * the forward-mode implementation with the reverse-mode alias fix.
+   */
+  mul(other) {
+    this.checkSameShape(other, "mul");
+    const p = new Float64Array(this.primal.length);
+    const t = new Float64Array(this.tangent.length);
+    if (this === other) {
+      for (let i = 0; i < p.length; i++) {
+        p[i] = this.primal[i] * this.primal[i];
+        t[i] = 2 * this.primal[i] * this.tangent[i];
+      }
+    } else {
+      for (let i = 0; i < p.length; i++) {
+        p[i] = this.primal[i] * other.primal[i];
+        t[i] = this.tangent[i] * other.primal[i] + this.primal[i] * other.tangent[i];
+      }
+    }
+    return new _DualTensor(this.shape, p, t);
+  }
+  /** Scalar multiplication following dual-number rule: (k·a, k·a'). */
+  scale(k) {
+    const p = new Float64Array(this.primal.length);
+    const t = new Float64Array(this.tangent.length);
+    for (let i = 0; i < p.length; i++) {
+      p[i] = this.primal[i] * k;
+      t[i] = this.tangent[i] * k;
+    }
+    return new _DualTensor(this.shape, p, t);
+  }
+  checkSameShape(other, op) {
+    if (this.shape.length !== other.shape.length || !this.shape.every((v, i) => v === other.shape[i])) {
+      throw new Error(`DualTensor.${op}: shape mismatch [${this.shape}] vs [${other.shape}]`);
+    }
+  }
+};
+
+// src/forward-grad.ts
+import { Tensor as Tensor2 } from "@danielsimonjr/mathts-tensor";
+function forwardGrad(fn, x) {
+  const xDualZero = DualTensor.fromTensor(x);
+  const yProbeRaw = fn(xDualZero);
+  if (!(yProbeRaw instanceof DualTensor)) {
+    throw new Error(
+      "forwardGrad: fn must be AD-traceable \u2014 its return must propagate through DualTensor arithmetic (use add/sub/mul/scale on the argument). A fresh Tensor return loses the tangent and silently corrupts the Jacobian."
+    );
+  }
+  const yPrimal = yProbeRaw.toPrimalTensor();
+  const jacobianShape = [...yPrimal.shape, ...x.shape];
+  const jacobianSize = jacobianShape.reduce((a, b) => a * b, 1);
+  const jacobianData = new Float64Array(jacobianSize);
+  const xSize = x.data.length;
+  const ySize = yPrimal.data.length;
+  for (let kX = 0; kX < xSize; kX++) {
+    const xDualUnit = DualTensor.unitAt(x, kX);
+    const yDualRaw = fn(xDualUnit);
+    if (!(yDualRaw instanceof DualTensor)) {
+      throw new Error("forwardGrad: fn lost AD trace mid-sweep (returned non-DualTensor)");
+    }
+    const yDual = yDualRaw;
+    for (let kY = 0; kY < ySize; kY++) {
+      jacobianData[kY * xSize + kX] = yDual.tangent[kY];
+    }
+  }
+  return {
+    value: yPrimal,
+    jacobian: new Tensor2(jacobianShape, jacobianData)
+  };
+}
+
+// src/tape.ts
+import { Tensor as Tensor3 } from "@danielsimonjr/mathts-tensor";
+var Tape = class {
+  nodes = [];
+  inputGradSlots = /* @__PURE__ */ new Map();
+  // S3 fix: disjoint ID namespaces. Inputs use negative IDs, ops use
+  // non-negative IDs. Eliminates the v0.4.0-review collision where
+  // id = nodes.length + 1000 would collide with op ids past 1000 entries.
+  nextOpId = 0;
+  nextInputId = -1;
+  // negatives = inputs; non-negatives = ops
+  /** Allocate a fresh id for an input or intermediate. */
+  allocate(size) {
+    const id = this.nextInputId--;
+    const gradSlot = new Float64Array(size);
+    this.inputGradSlots.set(id, gradSlot);
+    return { id, gradSlot };
+  }
+  record(inputIds, outputSize, backward) {
+    const outputGradSlot = new Float64Array(outputSize);
+    const id = this.nextOpId++;
+    this.nodes.push({ inputIds, backward, outputGradSlot });
+    this.inputGradSlots.set(id, outputGradSlot);
+    return { id, gradSlot: outputGradSlot };
+  }
+  /** Seed the final output's gradient slot, then replay in reverse. */
+  backward(outputId, outputGrad) {
+    const slot = this.inputGradSlots.get(outputId);
+    if (!slot) throw new Error(`Tape.backward: unknown outputId ${outputId}`);
+    slot.set(outputGrad);
+    for (let n = this.nodes.length - 1; n >= 0; n--) {
+      this.nodes[n].backward(this.nodes[n].outputGradSlot);
+    }
+  }
+  getInputGrad(id) {
+    return this.inputGradSlots.get(id);
+  }
+};
+var TapedTensor = class _TapedTensor {
+  constructor(shape, primal, tape, id) {
+    this.shape = shape;
+    this.primal = primal;
+    this.tape = tape;
+    this.id = id;
+  }
+  /**
+   * 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
+   * TapedTensor flows through them as a structurally-compatible Tensor.
+   * (AD-aware ops branch on `'tape' in arg` before reaching here.)
+   */
+  get data() {
+    return this.primal;
+  }
+  static fromTensorAsInput(t, tape) {
+    const { id } = tape.allocate(t.data.length);
+    return new _TapedTensor(t.shape, new Float64Array(t.data), tape, id);
+  }
+  toPrimalTensor() {
+    return new Tensor3(this.shape, new Float64Array(this.primal));
+  }
+  add(other) {
+    this.checkSameShape(other, "add");
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = this.primal[i] + other.primal[i];
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const otherGradSlot = this.tape.getInputGrad(other.id);
+    const { id } = this.tape.record([this.id, other.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i];
+        otherGradSlot[i] += outputGrad[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  sub(other) {
+    this.checkSameShape(other, "sub");
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = this.primal[i] - other.primal[i];
+    const thisGradSlot = this.tape.getInputGrad(this.id);
+    const otherGradSlot = this.tape.getInputGrad(other.id);
+    const { id } = this.tape.record([this.id, other.id], out.length, (outputGrad) => {
+      for (let i = 0; i < outputGrad.length; i++) {
+        thisGradSlot[i] += outputGrad[i];
+        otherGradSlot[i] -= outputGrad[i];
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  mul(other) {
+    this.checkSameShape(other, "mul");
+    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 = this.primal;
+    const otherPrimal = 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) {
+          thisGradSlot[i] += 2 * outputGrad[i] * thisPrimal[i];
+        } else {
+          thisGradSlot[i] += outputGrad[i] * otherPrimal[i];
+          otherGradSlot[i] += outputGrad[i] * thisPrimal[i];
+        }
+      }
+    });
+    return new _TapedTensor(this.shape, out, this.tape, id);
+  }
+  scale(k) {
+    const out = new Float64Array(this.primal.length);
+    for (let i = 0; i < out.length; i++) out[i] = this.primal[i] * k;
+    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;
+      }
+    });
+    return new _TapedTensor(this.shape, out, 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}]`);
+    }
+  }
+};
+
+// src/reverse-grad.ts
+import { Tensor as Tensor4 } from "@danielsimonjr/mathts-tensor";
+function reverseGrad(fn, x, cotangent) {
+  const tape = new Tape();
+  const xTaped = TapedTensor.fromTensorAsInput(x, tape);
+  const yRaw = fn(xTaped);
+  if (!(yRaw instanceof TapedTensor)) {
+    throw new Error(
+      "reverseGrad: fn must be AD-traceable \u2014 its return must propagate through TapedTensor arithmetic (use add/sub/mul/scale on the argument). A fresh Tensor return loses the tape and silently corrupts the gradient."
+    );
+  }
+  const yTaped = yRaw;
+  const value = yTaped.toPrimalTensor();
+  let ct;
+  if (cotangent === void 0) {
+    const data = new Float64Array(value.data.length).fill(1);
+    ct = new Tensor4(value.shape, data);
+  } else {
+    if (cotangent.shape.length !== value.shape.length || !cotangent.shape.every((v, i) => v === value.shape[i])) {
+      throw new Error(
+        `reverseGrad: cotangent shape [${cotangent.shape}] != value shape [${value.shape}]`
+      );
+    }
+    ct = cotangent;
+  }
+  tape.backward(yTaped.id, new Float64Array(ct.data));
+  const xGradSlot = tape.getInputGrad(xTaped.id);
+  return {
+    value,
+    gradient: new Tensor4(x.shape, new Float64Array(xGradSlot))
+  };
+}
+export {
+  DualTensor,
+  Tape,
+  TapedTensor,
+  forwardGrad,
+  reverseGrad
+};

package/package.json

@@ -0,0 +1,58 @@
+{
+  "name": "@danielsimonjr/mathts-autograd",
+  "version": "0.1.0",
+  "description": "Forward + reverse-mode automatic differentiation for MathTS rank-N Tensor",
+  "author": "Daniel Simon Jr.",
+  "license": "MIT",
+  "type": "module",
+  "main": "./dist/index.js",
+  "module": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "import": "./dist/index.js",
+      "types": "./dist/index.d.ts"
+    }
+  },
+  "files": [
+    "dist",
+    "README.md"
+  ],
+  "scripts": {
+    "build": "tsup src/index.ts --format esm --dts --clean",
+    "dev": "tsup src/index.ts --format esm --dts --watch",
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "test:coverage": "vitest run --coverage",
+    "typecheck": "tsc --noEmit",
+    "lint": "eslint src --ext .ts",
+    "lint:fix": "eslint src --ext .ts --fix",
+    "clean": "rm -rf dist",
+    "build:prod": "tsup src/index.ts --format esm --dts --clean --minify --treeshake"
+  },
+  "dependencies": {
+    "@danielsimonjr/mathts-core": "*",
+    "@danielsimonjr/mathts-tensor": "*"
+  },
+  "devDependencies": {
+    "@types/node": "^25.5.2",
+    "tsup": "^8.0.0",
+    "typescript": "^5.3.0",
+    "vitest": "^4.1.5"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/danielsimonjr/mathts",
+    "directory": "autograd"
+  },
+  "keywords": [
+    "math",
+    "autograd",
+    "automatic-differentiation",
+    "tensor",
+    "typescript"
+  ]
+}