@danielsimonjr/mathts-autograd 0.2.0 → 0.3.1

package/dist/index.d.ts

@@ -391,14 +391,21 @@ declare class TapedTensor {
      */
     square(): TapedTensor;
     /**
-     * Elementwise fixed-exponent power: x^k.
+     * Elementwise power: x^k.
      *
-     * Only fixed (non-TapedTensor) exponents are supported. Variable-exponent
-     * pow(taped, taped) is a follow-up slice.
+     * `k` may be a fixed `number` (constant exponent) or a `TapedTensor`
+     * (variable exponent, both inputs on the tape).
      *
-     * Adjoint: dX = dY · k · x^(k−1)
+     * Adjoints, for y = a^b:
+     *   - fixed exponent k:   dA = dY · k · a^(k−1)
+     *   - variable exponent:  dA = dY · b · a^(b−1),   dB = dY · a^b · ln(a)
+     *   - aliased a^a:         dA = dY · a^a · (ln a + 1)
+     *
+     * The dB term needs ln(a); for a ≤ 0 it is mathematically undefined and
+     * yields NaN/−Inf, honestly signalling non-differentiability w.r.t. the
+     * exponent rather than masking it.
      */
-    pow(k: number): TapedTensor;
+    pow(k: number | TapedTensor): TapedTensor;
     /**
      * Elementwise reciprocal: 1 / x.
      *
@@ -594,4 +601,51 @@ declare function reverseGrad(fn: (x: Tensor) => Tensor, x: Tensor, cotangent?: T
     gradient: Tensor;
 };
 
-export { DualTensor, Tape, TapedTensor, forwardGrad, reverseGrad };
+/**
+ * GC15 — ergonomic, number-friendly autodiff bridge.
+ *
+ * The low-level building blocks (`Tape`, `TapedTensor`, `reverseGrad`) require
+ * manual tape management and `Tensor` marshalling. This is the JAX-style
+ * convenience layer: write a function of a single `TapedTensor` (the AD-aware
+ * surface — `.sin()`, `.mul()`, `.sum()`, `.matmul()`, … all record adjoints)
+ * and differentiate it against plain JS numbers / arrays.
+ *
+ *   grad((x) => x.sin())(1)              // [cos(1)]
+ *   valueAndGrad((x) => x.pow(2).sum())([1, 2, 3])  // { value: 14, grad: [2,4,6] }
+ *   derivative((x) => x.sin().mul(x), 2) // sin(2) + 2cos(2)
+ *
+ * `f` must use TapedTensor operations (not the plain `functions/` ops, which are
+ * not AD-aware). Differentiating arbitrary `functions/` ops would require
+ * overloading them onto dual numbers / the tape — a separate, larger effort.
+ */
+
+type ScalarFn = (x: TapedTensor) => TapedTensor;
+type NumericInput = number | readonly number[] | Float64Array;
+/**
+ * Forward value and reverse-mode gradient of a scalar-valued function.
+ *
+ * @param fn  maps a TapedTensor (shape `[n]`) to a scalar TapedTensor (length 1).
+ * @param x   the point — a number, number[], or Float64Array.
+ * @returns   `{ value, grad }` where `grad` has the same length as `x`.
+ * @throws    if `fn` does not return a length-1 (scalar) result.
+ */
+declare function valueAndGrad(fn: ScalarFn, x: NumericInput): {
+    value: number;
+    grad: Float64Array;
+};
+/**
+ * Gradient function of a scalar-valued function: `grad(f)(x)` returns ∇f(x).
+ */
+declare function grad(fn: ScalarFn): (x: NumericInput) => Float64Array;
+/**
+ * Scalar convenience: derivative of `f: ℝ → ℝ` at `x0`.
+ */
+declare function derivative(fn: ScalarFn, x0: number): number;
+/**
+ * Jacobian of a vector-valued function `f: ℝⁿ → ℝᵐ`, one reverse pass per output
+ * row. Returns an `m × n` array. (For tall outputs forward-mode is cheaper; this
+ * favours wide-input / few-output cases, the common ML gradient shape.)
+ */
+declare function jacobian(fn: ScalarFn, x: NumericInput): number[][];
+
+export { DualTensor, type NumericInput, type ScalarFn, Tape, TapedTensor, derivative, forwardGrad, grad, jacobian, reverseGrad, valueAndGrad };

package/dist/index.js

@@ -658,8 +658,8 @@ var TapedTensor = class _TapedTensor {
     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 outStrides = Tensor3.rowMajorStrides(outShape);
+      const inStrides = Tensor3.rowMajorStrides(inputShape);
       const inCoords = new Array(inputShape.length).fill(0);
       for (let n = 0; n < inputSize; n++) {
         let outFlat = 0;
@@ -702,7 +702,7 @@ var TapedTensor = class _TapedTensor {
     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 outStrides = Tensor3.rowMajorStrides(outShape);
       const inCoords = new Array(inputShape.length).fill(0);
       for (let n = 0; n < inputSize; n++) {
         let outFlat = 0;
@@ -750,8 +750,8 @@ var TapedTensor = class _TapedTensor {
     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 outStrides = Tensor3.rowMajorStrides(outShape);
+      const inStrides = Tensor3.rowMajorStrides(inputShape);
       const groups = /* @__PURE__ */ new Map();
       const inCoords = new Array(inputShape.length).fill(0);
       for (let n = 0; n < primalData.length; n++) {
@@ -815,7 +815,7 @@ var TapedTensor = class _TapedTensor {
     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 outStrides = Tensor3.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++) {
@@ -864,7 +864,7 @@ var TapedTensor = class _TapedTensor {
     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 outStrides = Tensor3.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++) {
@@ -922,7 +922,7 @@ var TapedTensor = class _TapedTensor {
       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 outStrides = Tensor3.rowMajorStrides(outShape);
         const inCoords = new Array(inputShape.length).fill(0);
         for (let n = 0; n < primalData.length; n++) {
           let outFlat = 0;
@@ -950,7 +950,7 @@ var TapedTensor = class _TapedTensor {
     if (p === 1) {
       const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
         const outShape = [...outT.shape];
-        const outStrides = _rowMajorStrides(outShape);
+        const outStrides = Tensor3.rowMajorStrides(outShape);
         const inCoords = new Array(inputShape.length).fill(0);
         for (let n = 0; n < primalData.length; n++) {
           let outFlat = 0;
@@ -977,7 +977,7 @@ var TapedTensor = class _TapedTensor {
     {
       const { id } = this.tape.record([this.id], outT.data.length, (outputGrad) => {
         const outShape = [...outT.shape];
-        const outStrides = _rowMajorStrides(outShape);
+        const outStrides = Tensor3.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++) {
@@ -1145,21 +1145,55 @@ var TapedTensor = class _TapedTensor {
     return new _TapedTensor(this.shape, out, this.tape, id);
   }
   /**
-   * Elementwise fixed-exponent power: x^k.
+   * Elementwise power: x^k.
    *
-   * Only fixed (non-TapedTensor) exponents are supported. Variable-exponent
-   * pow(taped, taped) is a follow-up slice.
+   * `k` may be a fixed `number` (constant exponent) or a `TapedTensor`
+   * (variable exponent, both inputs on the tape).
    *
-   * Adjoint: dX = dY · k · x^(k−1)
+   * Adjoints, for y = a^b:
+   *   - fixed exponent k:   dA = dY · k · a^(k−1)
+   *   - variable exponent:  dA = dY · b · a^(b−1),   dB = dY · a^b · ln(a)
+   *   - aliased a^a:         dA = dY · a^a · (ln a + 1)
+   *
+   * The dB term needs ln(a); for a ≤ 0 it is mathematically undefined and
+   * yields NaN/−Inf, honestly signalling non-differentiability w.r.t. the
+   * exponent rather than masking it.
    */
   pow(k) {
+    if (typeof k === "number") {
+      const out2 = new Float64Array(this.primal.length);
+      for (let i = 0; i < out2.length; i++) out2[i] = Math.pow(this.primal[i], k);
+      const primalData = new Float64Array(this.primal);
+      const thisGradSlot2 = this.tape.getInputGrad(this.id);
+      const { id: id2 } = this.tape.record([this.id], out2.length, (outputGrad) => {
+        for (let i = 0; i < outputGrad.length; i++) {
+          thisGradSlot2[i] += outputGrad[i] * k * Math.pow(primalData[i], k - 1);
+        }
+      });
+      return new _TapedTensor(this.shape, out2, this.tape, id2);
+    }
+    const other = k;
+    this.checkSameShape(other, "pow");
     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);
+    for (let i = 0; i < out.length; i++) out[i] = Math.pow(this.primal[i], other.primal[i]);
+    const aPrimal = new Float64Array(this.primal);
+    const bPrimal = new Float64Array(other.primal);
+    const yPrimal = new Float64Array(out);
     const thisGradSlot = this.tape.getInputGrad(this.id);
-    const { id } = this.tape.record([this.id], out.length, (outputGrad) => {
+    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++) {
-        thisGradSlot[i] += outputGrad[i] * k * Math.pow(primalData[i], k - 1);
+        const a = aPrimal[i];
+        const b = bPrimal[i];
+        const y = yPrimal[i];
+        const g = outputGrad[i];
+        if (isAliased) {
+          thisGradSlot[i] += g * y * (Math.log(a) + 1);
+        } else {
+          thisGradSlot[i] += g * b * Math.pow(a, b - 1);
+          otherGradSlot[i] += g * y * Math.log(a);
+        }
       }
     });
     return new _TapedTensor(this.shape, out, this.tape, id);
@@ -1939,15 +1973,6 @@ function _resolveAxes(axis, rank) {
   }
   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;
@@ -2076,10 +2101,58 @@ function reverseGrad(fn, x, cotangent) {
     gradient: new Tensor4(x.shape, new Float64Array(xGradSlot))
   };
 }
+
+// src/grad.ts
+function toFloat64(x) {
+  if (typeof x === "number") return Float64Array.of(x);
+  return x instanceof Float64Array ? new Float64Array(x) : Float64Array.from(x);
+}
+function valueAndGrad(fn, x) {
+  const xs = toFloat64(x);
+  const tape = new Tape();
+  const { id } = tape.allocate(xs.length);
+  const leaf = new TapedTensor([xs.length], xs, tape, id);
+  const out = fn(leaf);
+  if (out.primal.length !== 1) {
+    throw new Error(
+      `valueAndGrad: fn must return a scalar (length 1); got length ${out.primal.length}. Reduce with .sum()/.mean(), or use jacobian() for vector outputs.`
+    );
+  }
+  tape.backward(out.id, Float64Array.of(1));
+  return { value: out.primal[0], grad: new Float64Array(tape.getInputGrad(id)) };
+}
+function grad(fn) {
+  return (x) => valueAndGrad(fn, x).grad;
+}
+function derivative(fn, x0) {
+  return valueAndGrad(fn, x0).grad[0];
+}
+function jacobian(fn, x) {
+  const xs = toFloat64(x);
+  const probeTape = new Tape();
+  const { id: probeId } = probeTape.allocate(xs.length);
+  const probeOut = fn(new TapedTensor([xs.length], new Float64Array(xs), probeTape, probeId));
+  const m = probeOut.primal.length;
+  const rows = [];
+  for (let r = 0; r < m; r++) {
+    const tape = new Tape();
+    const { id } = tape.allocate(xs.length);
+    const out = fn(new TapedTensor([xs.length], new Float64Array(xs), tape, id));
+    const seed = new Float64Array(m);
+    seed[r] = 1;
+    tape.backward(out.id, seed);
+    rows.push(Array.from(tape.getInputGrad(id)));
+  }
+  return rows;
+}
 export {
   DualTensor,
   Tape,
   TapedTensor,
+  derivative,
   forwardGrad,
-  reverseGrad
+  grad,
+  jacobian,
+  reverseGrad,
+  valueAndGrad
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@danielsimonjr/mathts-autograd",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "description": "Forward + reverse-mode automatic differentiation for MathTS rank-N Tensor",
   "author": "Daniel Simon Jr.",
   "license": "MIT",
@@ -31,8 +31,8 @@
     "build:prod": "tsup src/index.ts --format esm --dts --clean --minify --treeshake"
   },
   "dependencies": {
-    "@danielsimonjr/mathts-core": "^0.1.4",
-    "@danielsimonjr/mathts-tensor": "^0.1.3"
+    "@danielsimonjr/mathts-core": "^0.2.0",
+    "@danielsimonjr/mathts-tensor": "^0.2.1"
   },
   "devDependencies": {
     "@types/node": "^25.5.2",