tensorgrad 0.0.1

package/dist/shape.js

@@ -0,0 +1,259 @@
+// Shape inference and validation for each op kind.
+//
+// Every op in src/ops.ts validates its inputs and computes its output shape
+// through helpers here. Errors throw with the captured call-site so the
+// stack trace points at the user's line, not into the library.
+//
+// Broadcasting rules (deliberately limited):
+//   * For element-wise binops (add/sub/mul/div), we support trailing-axis
+//     broadcasting: the smaller operand's shape must be a suffix of the
+//     larger's, with axes of size 1 broadcasting to any size. Examples
+//     ALLOWED:  [B, T, D] op [D]  →  [B, T, D]
+//               [B, T, D] op [1, D]  → [B, T, D]
+//               [B, T, D] op [B, T, D]  → [B, T, D]
+//     Examples REJECTED:  [B, T, D] op [B]   (suffix mismatch)
+//                         [B, T, D] op [T, D] when T != B (legal numpy, banned here)
+//   The restriction makes codegen and autograd much simpler and covers every
+//   broadcast pattern in our transformer (biases, layernorm gain/bias, masks).
+import { formatSite } from './ir.js';
+// ============================================================================
+// Errors
+// ============================================================================
+export class ShapeError extends Error {
+    constructor(message, site) {
+        const formatted = site ? `${message}\n  at ${formatSite(site)}` : message;
+        super(formatted);
+        this.name = 'ShapeError';
+    }
+}
+function fail(message, site) {
+    throw new ShapeError(message, site);
+}
+// ============================================================================
+// Shape utilities
+// ============================================================================
+export function shapesEqual(a, b) {
+    if (a.length !== b.length)
+        return false;
+    for (let i = 0; i < a.length; i++)
+        if (a[i] !== b[i])
+            return false;
+    return true;
+}
+export function shapeSize(shape) {
+    let n = 1;
+    for (const d of shape)
+        n *= d;
+    return n;
+}
+export function showShape(shape) {
+    return `[${shape.join(', ')}]`;
+}
+// Standard right-aligned NumPy-style broadcasting. Pad the shorter shape with
+// leading 1s, then per-axis: equal dims unify, size-1 dims broadcast on either
+// side, otherwise incompatible. Returns the resulting shape or null.
+export function broadcastTrailing(a, b) {
+    const rank = Math.max(a.length, b.length);
+    const out = new Array(rank);
+    for (let i = 0; i < rank; i++) {
+        const ai = i - (rank - a.length);
+        const bi = i - (rank - b.length);
+        const av = ai < 0 ? 1 : a[ai];
+        const bv = bi < 0 ? 1 : b[bi];
+        if (av === bv)
+            out[i] = av;
+        else if (av === 1)
+            out[i] = bv;
+        else if (bv === 1)
+            out[i] = av;
+        else
+            return null;
+    }
+    return out;
+}
+// ============================================================================
+// Per-op shape rules
+// ============================================================================
+//
+// Each rule takes the input shapes and returns the output shape, or throws.
+// All rules accept a `site` for error attribution.
+export function inferElementwiseBinop(opName, aShape, bShape, site) {
+    const result = broadcastTrailing(aShape, bShape);
+    if (!result) {
+        fail(`${opName}: incompatible shapes ${showShape(aShape)} and ${showShape(bShape)}. ` +
+            `Trailing-suffix broadcasting only — the smaller shape must be a suffix of the larger, ` +
+            `with size-1 axes broadcasting to any size.`, site);
+    }
+    return result;
+}
+export function inferUnary(_opName, aShape, _site) {
+    return aShape;
+}
+export function inferMeanLast(opName, aShape, site) {
+    if (aShape.length === 0)
+        fail(`${opName}: cannot reduce a 0-d tensor`, site);
+    // keepdims=true: replace last axis with 1.
+    return [...aShape.slice(0, -1), 1];
+}
+export function inferSumLast(opName, aShape, site) {
+    if (aShape.length === 0)
+        fail(`${opName}: cannot reduce a 0-d tensor`, site);
+    // keepdims=false: drop the last axis.
+    return aShape.slice(0, -1);
+}
+export function inferReshape(opName, aShape, newShape, site) {
+    // Validate -1 placeholder (at most one allowed) and total size match.
+    let inferIdx = -1;
+    let knownSize = 1;
+    for (let i = 0; i < newShape.length; i++) {
+        const d = newShape[i];
+        if (d === -1) {
+            if (inferIdx !== -1)
+                fail(`${opName}: at most one -1 dim allowed in newShape ${showShape(newShape)}`, site);
+            inferIdx = i;
+        }
+        else if (d <= 0) {
+            fail(`${opName}: invalid dim ${d} in newShape ${showShape(newShape)}`, site);
+        }
+        else {
+            knownSize *= d;
+        }
+    }
+    const totalIn = shapeSize(aShape);
+    const out = [...newShape];
+    if (inferIdx !== -1) {
+        if (totalIn % knownSize !== 0) {
+            fail(`${opName}: cannot reshape ${showShape(aShape)} (size ${totalIn}) to ${showShape(newShape)} — known dims multiply to ${knownSize}`, site);
+        }
+        out[inferIdx] = totalIn / knownSize;
+    }
+    else if (knownSize !== totalIn) {
+        fail(`${opName}: size mismatch — input ${showShape(aShape)} has ${totalIn} elements but newShape ${showShape(newShape)} has ${knownSize}`, site);
+    }
+    return out;
+}
+export function inferTranspose(opName, aShape, perm, site) {
+    if (perm.length !== aShape.length) {
+        fail(`${opName}: perm length ${perm.length} must equal input rank ${aShape.length}`, site);
+    }
+    const seen = new Set();
+    for (const p of perm) {
+        if (p < 0 || p >= aShape.length)
+            fail(`${opName}: perm index ${p} out of range for rank ${aShape.length}`, site);
+        if (seen.has(p))
+            fail(`${opName}: perm has duplicate index ${p}`, site);
+        seen.add(p);
+    }
+    return perm.map(p => aShape[p]);
+}
+// matmul: a [..., M, K] · b [K, N]  →  [..., M, N].  b is unbatched.
+export function inferMatmul(opName, aShape, bShape, site) {
+    if (aShape.length < 2)
+        fail(`${opName}: lhs must have rank >= 2, got ${showShape(aShape)}`, site);
+    if (bShape.length !== 2)
+        fail(`${opName}: rhs must have rank 2, got ${showShape(bShape)} — use matmulBatched for batched rhs`, site);
+    const M = aShape[aShape.length - 2];
+    const Ka = aShape[aShape.length - 1];
+    const Kb = bShape[0];
+    const N = bShape[1];
+    if (Ka !== Kb)
+        fail(`${opName}: inner dims don't match — ${showShape(aShape)} · ${showShape(bShape)} (last axis of lhs = ${Ka}, first axis of rhs = ${Kb})`, site);
+    return [...aShape.slice(0, -2), M, N];
+}
+// matmul_batched: a [..., M, K] · b [..., K, N]  →  [..., M, N].  Both have leading batch dims.
+export function inferMatmulBatched(opName, aShape, bShape, site) {
+    if (aShape.length < 2 || bShape.length < 2) {
+        fail(`${opName}: both inputs must have rank >= 2, got ${showShape(aShape)} and ${showShape(bShape)}`, site);
+    }
+    if (aShape.length !== bShape.length) {
+        fail(`${opName}: ranks must match (got ${aShape.length} vs ${bShape.length}). Reshape if you need different batch dims.`, site);
+    }
+    const aBatch = aShape.slice(0, -2);
+    const bBatch = bShape.slice(0, -2);
+    for (let i = 0; i < aBatch.length; i++) {
+        if (aBatch[i] !== bBatch[i]) {
+            fail(`${opName}: batch dims must match — ${showShape(aShape)} vs ${showShape(bShape)}`, site);
+        }
+    }
+    const M = aShape[aShape.length - 2];
+    const Ka = aShape[aShape.length - 1];
+    const Kb = bShape[bShape.length - 2];
+    const N = bShape[bShape.length - 1];
+    if (Ka !== Kb)
+        fail(`${opName}: inner dims don't match — last axis of lhs = ${Ka}, second-to-last of rhs = ${Kb}`, site);
+    return [...aBatch, M, N];
+}
+export function inferOneHot(opName, indicesShape, depth, site) {
+    if (depth <= 0)
+        fail(`${opName}: depth must be positive, got ${depth}`, site);
+    return [...indicesShape, depth];
+}
+// where_causal preserves shape but requires the last two axes to be square.
+export function inferWhereCausal(opName, aShape, site) {
+    if (aShape.length < 2)
+        fail(`${opName}: requires rank >= 2, got ${showShape(aShape)}`, site);
+    const m = aShape[aShape.length - 2];
+    const n = aShape[aShape.length - 1];
+    if (m !== n)
+        fail(`${opName}: last two axes must be equal (square mask), got ${showShape(aShape)}`, site);
+    return aShape;
+}
+export function inferSliceLastRange(opName, aShape, start, end, site) {
+    if (aShape.length === 0)
+        fail(`${opName}: cannot slice 0-d tensor`, site);
+    const last = aShape[aShape.length - 1];
+    if (start < 0 || end > last || start >= end) {
+        fail(`${opName}: invalid range [${start}, ${end}) for last axis of size ${last}`, site);
+    }
+    return [...aShape.slice(0, -1), end - start];
+}
+// broadcast_to: validate that `aShape` can broadcast to `targetShape` under
+// right-aligned NumPy rules. Returns targetShape on success.
+export function inferBroadcastTo(opName, aShape, targetShape, site) {
+    if (aShape.length > targetShape.length) {
+        fail(`${opName}: source rank ${aShape.length} > target rank ${targetShape.length}`, site);
+    }
+    const offset = targetShape.length - aShape.length;
+    for (let i = 0; i < aShape.length; i++) {
+        const av = aShape[i];
+        const tv = targetShape[offset + i];
+        if (av !== tv && av !== 1) {
+            fail(`${opName}: cannot broadcast ${showShape(aShape)} to ${showShape(targetShape)} — axis ${i} (size ${av}) doesn't match target axis ${offset + i} (size ${tv}) and isn't 1`, site);
+        }
+    }
+    return targetShape;
+}
+// sum_to_shape: validate that `targetShape` is a valid right-aligned reduction
+// of `aShape` (i.e., aShape can have been produced by broadcasting targetShape).
+export function inferSumToShape(opName, aShape, targetShape, site) {
+    if (targetShape.length > aShape.length) {
+        fail(`${opName}: target rank ${targetShape.length} > source rank ${aShape.length}`, site);
+    }
+    const offset = aShape.length - targetShape.length;
+    for (let i = 0; i < targetShape.length; i++) {
+        const av = aShape[offset + i];
+        const tv = targetShape[i];
+        if (av !== tv && tv !== 1) {
+            fail(`${opName}: cannot sum-reduce ${showShape(aShape)} to ${showShape(targetShape)} — target axis ${i} (size ${tv}) must be 1 or match source`, site);
+        }
+    }
+    return targetShape;
+}
+// Three-way broadcast for `where(cond, a, b)`. All three shapes must broadcast
+// to a common shape under standard NumPy rules.
+export function inferWhere(opName, condShape, aShape, bShape, site) {
+    const ab = broadcastTrailing(aShape, bShape);
+    if (!ab)
+        fail(`${opName}: a/b incompatible: ${showShape(aShape)} vs ${showShape(bShape)}`, site);
+    const result = broadcastTrailing(condShape, ab);
+    if (!result)
+        fail(`${opName}: cond ${showShape(condShape)} incompatible with broadcast(a, b) ${showShape(ab)}`, site);
+    return result;
+}
+export function inferReluGrad(opName, xShape, dyShape, site) {
+    if (!shapesEqual(xShape, dyShape)) {
+        fail(`${opName}: x and dy must have matching shapes, got ${showShape(xShape)} and ${showShape(dyShape)}`, site);
+    }
+    return xShape;
+}
+//# sourceMappingURL=shape.js.map

package/dist/shape.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"shape.js","sourceRoot":"","sources":["../src/shape.ts"],"names":[],"mappings":"AAAA,mDAAmD;AACnD,EAAE;AACF,4EAA4E;AAC5E,wEAAwE;AACxE,+DAA+D;AAC/D,EAAE;AACF,6CAA6C;AAC7C,0EAA0E;AAC1E,wEAAwE;AACxE,uEAAuE;AACvE,+CAA+C;AAC/C,iDAAiD;AACjD,oDAAoD;AACpD,+DAA+D;AAC/D,qFAAqF;AACrF,6EAA6E;AAC7E,+EAA+E;AAG/E,OAAO,EAAE,UAAU,EAAE,MAAM,SAAS,CAAA;AAEpC,+EAA+E;AAC/E,SAAS;AACT,+EAA+E;AAE/E,MAAM,OAAO,UAAW,SAAQ,KAAK;IACnC,YAAY,OAAe,EAAE,IAAqB;QAChD,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,GAAG,OAAO,UAAU,UAAU,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,CAAA;QACzE,KAAK,CAAC,SAAS,CAAC,CAAA;QAChB,IAAI,CAAC,IAAI,GAAG,YAAY,CAAA;IAC1B,CAAC;CACF;AAED,SAAS,IAAI,CAAC,OAAe,EAAE,IAAqB;IAClD,MAAM,IAAI,UAAU,CAAC,OAAO,EAAE,IAAI,CAAC,CAAA;AACrC,CAAC;AAED,+EAA+E;AAC/E,kBAAkB;AAClB,+EAA+E;AAE/E,MAAM,UAAU,WAAW,CAAC,CAAQ,EAAE,CAAQ;IAC5C,IAAI,CAAC,CAAC,MAAM,KAAK,CAAC,CAAC,MAAM;QAAE,OAAO,KAAK,CAAA;IACvC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,EAAE,CAAC,EAAE;QAAE,IAAI,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YAAE,OAAO,KAAK,CAAA;IAClE,OAAO,IAAI,CAAA;AACb,CAAC;AAED,MAAM,UAAU,SAAS,CAAC,KAAY;IACpC,IAAI,CAAC,GAAG,CAAC,CAAA;IACT,KAAK,MAAM,CAAC,IAAI,KAAK;QAAE,CAAC,IAAI,CAAC,CAAA;IAC7B,OAAO,CAAC,CAAA;AACV,CAAC;AAED,MAAM,UAAU,SAAS,CAAC,KAAY;IACpC,OAAO,IAAI,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,GAAG,CAAA;AAChC,CAAC;AAED,8EAA8E;AAC9E,+EAA+E;AAC/E,qEAAqE;AACrE,MAAM,UAAU,iBAAiB,CAAC,CAAQ,EAAE,CAAQ;IAClD,MAAM,IAAI,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,MAAM,EAAE,CAAC,CAAC,MAAM,CAAC,CAAA;IACzC,MAAM,GAAG,GAAa,IAAI,KAAK,CAAC,IAAI,CAAC,CAAA;IACrC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,EAAE,CAAC,EAAE,EAAE,CAAC;QAC9B,MAAM,EAAE,GAAG,CAAC,GAAG,CAAC,IAAI,GAAG,CAAC,CAAC,MAAM,CAAC,CAAA;QAChC,MAAM,EAAE,GAAG,CAAC,GAAG,CAAC,IAAI,GAAG,CAAC,CAAC,MAAM,CAAC,CAAA;QAChC,MAAM,EAAE,GAAG,EAAE,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,EAAE,CAAE,CAAA;QAC9B,MAAM,EAAE,GAAG,EAAE,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,EAAE,CAAE,CAAA;QAC9B,IAAI,EAAE,KAAK,EAAE;YAAE,GAAG,CAAC,CAAC,CAAC,GAAG,EAAE,CAAA;aACrB,IAAI,EAAE,KAAK,CAAC;YAAE,GAAG,CAAC,CAAC,CAAC,GAAG,EAAE,CAAA;aACzB,IAAI,EAAE,KAAK,CAAC;YAAE,GAAG,CAAC,CAAC,CAAC,GAAG,EAAE,CAAA;;YACzB,OAAO,IAAI,CAAA;IAClB,CAAC;IACD,OAAO,GAAG,CAAA;AACZ,CAAC;AAED,+EAA+E;AAC/E,qBAAqB;AACrB,+EAA+E;AAC/E,EAAE;AACF,4EAA4E;AAC5E,mDAAmD;AAEnD,MAAM,UAAU,qBAAqB,CACnC,MAAc,EAAE,MAAa,EAAE,MAAa,EAAE,IAAqB;IAEnE,MAAM,MAAM,GAAG,iBAAiB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IAChD,IAAI,CAAC,MAAM,EAAE,CAAC;QACZ,IAAI,CACF,GAAG,MAAM,yBAAyB,SAAS,CAAC,MAAM,CAAC,QAAQ,SAAS,CAAC,MAAM,CAAC,IAAI;YAChF,wFAAwF;YACxF,4CAA4C,EAC5C,IAAI,CACL,CAAA;IACH,CAAC;IACD,OAAO,MAAM,CAAA;AACf,CAAC;AAED,MAAM,UAAU,UAAU,CAAC,OAAe,EAAE,MAAa,EAAE,KAAsB;IAC/E,OAAO,MAAM,CAAA;AACf,CAAC;AAED,MAAM,UAAU,aAAa,CAAC,MAAc,EAAE,MAAa,EAAE,IAAqB;IAChF,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,8BAA8B,EAAE,IAAI,CAAC,CAAA;IAC5E,2CAA2C;IAC3C,OAAO,CAAC,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAA;AACpC,CAAC;AAED,MAAM,UAAU,YAAY,CAAC,MAAc,EAAE,MAAa,EAAE,IAAqB;IAC/E,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,8BAA8B,EAAE,IAAI,CAAC,CAAA;IAC5E,sCAAsC;IACtC,OAAO,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAA;AAC5B,CAAC;AAED,MAAM,UAAU,YAAY,CAAC,MAAc,EAAE,MAAa,EAAE,QAAe,EAAE,IAAqB;IAChG,sEAAsE;IACtE,IAAI,QAAQ,GAAG,CAAC,CAAC,CAAA;IACjB,IAAI,SAAS,GAAG,CAAC,CAAA;IACjB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACzC,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAE,CAAA;QACtB,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC;YACb,IAAI,QAAQ,KAAK,CAAC,CAAC;gBAAE,IAAI,CAAC,GAAG,MAAM,4CAA4C,SAAS,CAAC,QAAQ,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;YAC3G,QAAQ,GAAG,CAAC,CAAA;QACd,CAAC;aAAM,IAAI,CAAC,IAAI,CAAC,EAAE,CAAC;YAClB,IAAI,CAAC,GAAG,MAAM,iBAAiB,CAAC,gBAAgB,SAAS,CAAC,QAAQ,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;QAC9E,CAAC;aAAM,CAAC;YACN,SAAS,IAAI,CAAC,CAAA;QAChB,CAAC;IACH,CAAC;IACD,MAAM,OAAO,GAAG,SAAS,CAAC,MAAM,CAAC,CAAA;IACjC,MAAM,GAAG,GAAG,CAAC,GAAG,QAAQ,CAAC,CAAA;IACzB,IAAI,QAAQ,KAAK,CAAC,CAAC,EAAE,CAAC;QACpB,IAAI,OAAO,GAAG,SAAS,KAAK,CAAC,EAAE,CAAC;YAC9B,IAAI,CAAC,GAAG,MAAM,oBAAoB,SAAS,CAAC,MAAM,CAAC,UAAU,OAAO,QAAQ,SAAS,CAAC,QAAQ,CAAC,6BAA6B,SAAS,EAAE,EAAE,IAAI,CAAC,CAAA;QAChJ,CAAC;QACD,GAAG,CAAC,QAAQ,CAAC,GAAG,OAAO,GAAG,SAAS,CAAA;IACrC,CAAC;SAAM,IAAI,SAAS,KAAK,OAAO,EAAE,CAAC;QACjC,IAAI,CAAC,GAAG,MAAM,2BAA2B,SAAS,CAAC,MAAM,CAAC,QAAQ,OAAO,0BAA0B,SAAS,CAAC,QAAQ,CAAC,QAAQ,SAAS,EAAE,EAAE,IAAI,CAAC,CAAA;IAClJ,CAAC;IACD,OAAO,GAAG,CAAA;AACZ,CAAC;AAED,MAAM,UAAU,cAAc,CAAC,MAAc,EAAE,MAAa,EAAE,IAAuB,EAAE,IAAqB;IAC1G,IAAI,IAAI,CAAC,MAAM,KAAK,MAAM,CAAC,MAAM,EAAE,CAAC;QAClC,IAAI,CAAC,GAAG,MAAM,iBAAiB,IAAI,CAAC,MAAM,0BAA0B,MAAM,CAAC,MAAM,EAAE,EAAE,IAAI,CAAC,CAAA;IAC5F,CAAC;IACD,MAAM,IAAI,GAAG,IAAI,GAAG,EAAU,CAAA;IAC9B,KAAK,MAAM,CAAC,IAAI,IAAI,EAAE,CAAC;QACrB,IAAI,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,MAAM,CAAC,MAAM;YAAE,IAAI,CAAC,GAAG,MAAM,gBAAgB,CAAC,0BAA0B,MAAM,CAAC,MAAM,EAAE,EAAE,IAAI,CAAC,CAAA;QAChH,IAAI,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC;YAAE,IAAI,CAAC,GAAG,MAAM,8BAA8B,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;QACvE,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,CAAA;IACb,CAAC;IACD,OAAO,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,MAAM,CAAC,CAAC,CAAE,CAAC,CAAA;AAClC,CAAC;AAED,qEAAqE;AACrE,MAAM,UAAU,WAAW,CAAC,MAAc,EAAE,MAAa,EAAE,MAAa,EAAE,IAAqB;IAC7F,IAAI,MAAM,CAAC,MAAM,GAAG,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,kCAAkC,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACjG,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,+BAA+B,SAAS,CAAC,MAAM,CAAC,sCAAsC,EAAE,IAAI,CAAC,CAAA;IACpI,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACpC,MAAM,EAAE,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACrC,MAAM,EAAE,GAAG,MAAM,CAAC,CAAC,CAAE,CAAA;IACrB,MAAM,CAAC,GAAG,MAAM,CAAC,CAAC,CAAE,CAAA;IACpB,IAAI,EAAE,KAAK,EAAE;QAAE,IAAI,CAAC,GAAG,MAAM,8BAA8B,SAAS,CAAC,MAAM,CAAC,MAAM,SAAS,CAAC,MAAM,CAAC,wBAAwB,EAAE,yBAAyB,EAAE,GAAG,EAAE,IAAI,CAAC,CAAA;IAClK,OAAO,CAAC,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AACvC,CAAC;AAED,gGAAgG;AAChG,MAAM,UAAU,kBAAkB,CAAC,MAAc,EAAE,MAAa,EAAE,MAAa,EAAE,IAAqB;IACpG,IAAI,MAAM,CAAC,MAAM,GAAG,CAAC,IAAI,MAAM,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QAC3C,IAAI,CAAC,GAAG,MAAM,0CAA0C,SAAS,CAAC,MAAM,CAAC,QAAQ,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IAC7G,CAAC;IACD,IAAI,MAAM,CAAC,MAAM,KAAK,MAAM,CAAC,MAAM,EAAE,CAAC;QACpC,IAAI,CAAC,GAAG,MAAM,2BAA2B,MAAM,CAAC,MAAM,OAAO,MAAM,CAAC,MAAM,8CAA8C,EAAE,IAAI,CAAC,CAAA;IACjI,CAAC;IACD,MAAM,MAAM,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAA;IAClC,MAAM,MAAM,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAA;IAClC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACvC,IAAI,MAAM,CAAC,CAAC,CAAC,KAAK,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC;YAC5B,IAAI,CAAC,GAAG,MAAM,6BAA6B,SAAS,CAAC,MAAM,CAAC,OAAO,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;QAC/F,CAAC;IACH,CAAC;IACD,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACpC,MAAM,EAAE,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACrC,MAAM,EAAE,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACrC,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACpC,IAAI,EAAE,KAAK,EAAE;QAAE,IAAI,CAAC,GAAG,MAAM,iDAAiD,EAAE,6BAA6B,EAAE,EAAE,EAAE,IAAI,CAAC,CAAA;IACxH,OAAO,CAAC,GAAG,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AAC1B,CAAC;AAED,MAAM,UAAU,WAAW,CAAC,MAAc,EAAE,YAAmB,EAAE,KAAa,EAAE,IAAqB;IACnG,IAAI,KAAK,IAAI,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,iCAAiC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC7E,OAAO,CAAC,GAAG,YAAY,EAAE,KAAK,CAAC,CAAA;AACjC,CAAC;AAED,4EAA4E;AAC5E,MAAM,UAAU,gBAAgB,CAAC,MAAc,EAAE,MAAa,EAAE,IAAqB;IACnF,IAAI,MAAM,CAAC,MAAM,GAAG,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,6BAA6B,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IAC5F,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACpC,MAAM,CAAC,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACpC,IAAI,CAAC,KAAK,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,oDAAoD,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACzG,OAAO,MAAM,CAAA;AACf,CAAC;AAED,MAAM,UAAU,mBAAmB,CAAC,MAAc,EAAE,MAAa,EAAE,KAAa,EAAE,GAAW,EAAE,IAAqB;IAClH,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC;QAAE,IAAI,CAAC,GAAG,MAAM,2BAA2B,EAAE,IAAI,CAAC,CAAA;IACzE,MAAM,IAAI,GAAG,MAAM,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;IACvC,IAAI,KAAK,GAAG,CAAC,IAAI,GAAG,GAAG,IAAI,IAAI,KAAK,IAAI,GAAG,EAAE,CAAC;QAC5C,IAAI,CAAC,GAAG,MAAM,oBAAoB,KAAK,KAAK,GAAG,2BAA2B,IAAI,EAAE,EAAE,IAAI,CAAC,CAAA;IACzF,CAAC;IACD,OAAO,CAAC,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,GAAG,GAAG,KAAK,CAAC,CAAA;AAC9C,CAAC;AAED,4EAA4E;AAC5E,6DAA6D;AAC7D,MAAM,UAAU,gBAAgB,CAAC,MAAc,EAAE,MAAa,EAAE,WAAkB,EAAE,IAAqB;IACvG,IAAI,MAAM,CAAC,MAAM,GAAG,WAAW,CAAC,MAAM,EAAE,CAAC;QACvC,IAAI,CAAC,GAAG,MAAM,iBAAiB,MAAM,CAAC,MAAM,kBAAkB,WAAW,CAAC,MAAM,EAAE,EAAE,IAAI,CAAC,CAAA;IAC3F,CAAC;IACD,MAAM,MAAM,GAAG,WAAW,CAAC,MAAM,GAAG,MAAM,CAAC,MAAM,CAAA;IACjD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACvC,MAAM,EAAE,GAAG,MAAM,CAAC,CAAC,CAAE,CAAA;QACrB,MAAM,EAAE,GAAG,WAAW,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;QACnC,IAAI,EAAE,KAAK,EAAE,IAAI,EAAE,KAAK,CAAC,EAAE,CAAC;YAC1B,IAAI,CAAC,GAAG,MAAM,sBAAsB,SAAS,CAAC,MAAM,CAAC,OAAO,SAAS,CAAC,WAAW,CAAC,WAAW,CAAC,UAAU,EAAE,+BAA+B,MAAM,GAAG,CAAC,UAAU,EAAE,eAAe,EAAE,IAAI,CAAC,CAAA;QACvL,CAAC;IACH,CAAC;IACD,OAAO,WAAW,CAAA;AACpB,CAAC;AAED,+EAA+E;AAC/E,iFAAiF;AACjF,MAAM,UAAU,eAAe,CAAC,MAAc,EAAE,MAAa,EAAE,WAAkB,EAAE,IAAqB;IACtG,IAAI,WAAW,CAAC,MAAM,GAAG,MAAM,CAAC,MAAM,EAAE,CAAC;QACvC,IAAI,CAAC,GAAG,MAAM,iBAAiB,WAAW,CAAC,MAAM,kBAAkB,MAAM,CAAC,MAAM,EAAE,EAAE,IAAI,CAAC,CAAA;IAC3F,CAAC;IACD,MAAM,MAAM,GAAG,MAAM,CAAC,MAAM,GAAG,WAAW,CAAC,MAAM,CAAA;IACjD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QAC5C,MAAM,EAAE,GAAG,MAAM,CAAC,MAAM,GAAG,CAAC,CAAE,CAAA;QAC9B,MAAM,EAAE,GAAG,WAAW,CAAC,CAAC,CAAE,CAAA;QAC1B,IAAI,EAAE,KAAK,EAAE,IAAI,EAAE,KAAK,CAAC,EAAE,CAAC;YAC1B,IAAI,CAAC,GAAG,MAAM,uBAAuB,SAAS,CAAC,MAAM,CAAC,OAAO,SAAS,CAAC,WAAW,CAAC,kBAAkB,CAAC,UAAU,EAAE,6BAA6B,EAAE,IAAI,CAAC,CAAA;QACxJ,CAAC;IACH,CAAC;IACD,OAAO,WAAW,CAAA;AACpB,CAAC;AAED,+EAA+E;AAC/E,gDAAgD;AAChD,MAAM,UAAU,UAAU,CAAC,MAAc,EAAE,SAAgB,EAAE,MAAa,EAAE,MAAa,EAAE,IAAqB;IAC9G,MAAM,EAAE,GAAG,iBAAiB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IAC5C,IAAI,CAAC,EAAE;QAAE,IAAI,CAAC,GAAG,MAAM,uBAAuB,SAAS,CAAC,MAAM,CAAC,OAAO,SAAS,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IAChG,MAAM,MAAM,GAAG,iBAAiB,CAAC,SAAS,EAAE,EAAE,CAAC,CAAA;IAC/C,IAAI,CAAC,MAAM;QAAE,IAAI,CAAC,GAAG,MAAM,UAAU,SAAS,CAAC,SAAS,CAAC,sCAAsC,SAAS,CAAC,EAAE,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACrH,OAAO,MAAM,CAAA;AACf,CAAC;AAED,MAAM,UAAU,aAAa,CAAC,MAAc,EAAE,MAAa,EAAE,OAAc,EAAE,IAAqB;IAChG,IAAI,CAAC,WAAW,CAAC,MAAM,EAAE,OAAO,CAAC,EAAE,CAAC;QAClC,IAAI,CAAC,GAAG,MAAM,6CAA6C,SAAS,CAAC,MAAM,CAAC,QAAQ,SAAS,CAAC,OAAO,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACjH,CAAC;IACD,OAAO,MAAM,CAAA;AACf,CAAC"}

package/dist/trace.d.ts

@@ -0,0 +1,8 @@
+import type { Graph, Tensor, Shape, Dtype } from './ir.js';
+export declare function currentGraph(): Graph;
+export declare function trace(fn: () => Tensor | Tensor[]): Graph;
+export declare function traceInto<T>(g: Graph, fn: () => T): T;
+export declare function paramInput(name: string, shape: Shape, dtype?: Dtype): Tensor;
+export declare function tensorInput(name: string, shape: Shape, dtype?: Dtype): Tensor;
+export declare function stateInput(name: string, shape: Shape, dtype?: Dtype, initValue?: number): Tensor;
+//# sourceMappingURL=trace.d.ts.map

package/dist/trace.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"trace.d.ts","sourceRoot":"","sources":["../src/trace.ts"],"names":[],"mappings":"AAgBA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAM1D,wBAAgB,YAAY,IAAI,KAAK,CAQpC;AAID,wBAAgB,KAAK,CAAC,EAAE,EAAE,MAAM,MAAM,GAAG,MAAM,EAAE,GAAG,KAAK,CAgBxD;AAMD,wBAAgB,SAAS,CAAC,CAAC,EAAE,CAAC,EAAE,KAAK,EAAE,EAAE,EAAE,MAAM,CAAC,GAAG,CAAC,CAUrD;AAOD,wBAAgB,UAAU,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAOnF;AAED,wBAAgB,WAAW,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAOpF;AAID,wBAAgB,UAAU,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,GAAE,KAAa,EAAE,SAAS,SAAI,GAAG,MAAM,CAOlG"}

package/dist/trace.js

@@ -0,0 +1,93 @@
+// Trace driver. Holds the "current graph" in module-local state so user code
+// can call ops without threading a graph parameter through every function.
+//
+// Usage:
+//
+//   const graph = trace(() => {
+//     const x = tensorInput('x', [B, T], 'i32')
+//     const w = paramInput('w', [V, D], 'f32')
+//     // ... user computation building tensors ...
+//     return finalLossTensor
+//   })
+//
+// `trace` is single-threaded and re-entrant only via nested calls (which share
+// the outer graph — but we don't currently have a use for nesting). Calling an
+// op outside a `trace(...)` block is an error.
+import { makeGraph, addOp, captureSite } from './ir.js';
+// Module-local: the graph being built right now, or null if no trace is active.
+let _current = null;
+export function currentGraph() {
+    if (!_current) {
+        throw new Error('tensorgrad: ops can only be called inside trace(). ' +
+            'Did you forget to wrap your forward pass?');
+    }
+    return _current;
+}
+// Run `fn` with a fresh graph as the current one; capture and return the graph.
+// `fn` must return the tensor (or array of tensors) to mark as graph outputs.
+export function trace(fn) {
+    if (_current) {
+        throw new Error('tensorgrad: nested trace() is not supported');
+    }
+    const g = makeGraph();
+    _current = g;
+    try {
+        const result = fn();
+        const outputs = Array.isArray(result) ? result : [result];
+        for (const t of outputs) {
+            ;
+            g.outputs.push(t.id);
+        }
+    }
+    finally {
+        _current = null;
+    }
+    return g;
+}
+// Re-enter an existing graph to append more ops. Used by autograd to add
+// backward ops to a graph that's already been traced. `fn` runs with the
+// supplied graph as the current one; any ops it calls append to that graph.
+// Returns whatever `fn` returns.
+export function traceInto(g, fn) {
+    if (_current) {
+        throw new Error('tensorgrad: traceInto() called while another trace is active');
+    }
+    _current = g;
+    try {
+        return fn();
+    }
+    finally {
+        _current = null;
+    }
+}
+// ---- Leaf tensor builders --------------------------------------------------
+// Inputs are added to the graph as `param_input` or `tensor_input` op nodes.
+// Their .source on the Tensor points at that node so codegen knows where to
+// bind external data.
+export function paramInput(name, shape, dtype = 'f32') {
+    const g = currentGraph();
+    if (g.ops.some(op => (op.kind === 'param_input' || op.kind === 'tensor_input') && op.name === name)) {
+        throw new Error(`tensorgrad: input name '${name}' already used in this trace`);
+    }
+    const site = captureSite('paramInput');
+    return addOp(g, 'param_input', shape, dtype, site, { name });
+}
+export function tensorInput(name, shape, dtype = 'f32') {
+    const g = currentGraph();
+    if (g.ops.some(op => (op.kind === 'param_input' || op.kind === 'tensor_input') && op.name === name)) {
+        throw new Error(`tensorgrad: input name '${name}' already used in this trace`);
+    }
+    const site = captureSite('tensorInput');
+    return addOp(g, 'tensor_input', shape, dtype, site, { name });
+}
+// Persistent state buffer. Allocated at compile time, zero-(or initValue-)initialized,
+// and updated across step() calls via writebacks declared by the optimizer helper.
+export function stateInput(name, shape, dtype = 'f32', initValue = 0) {
+    const g = currentGraph();
+    if (g.ops.some(op => op.kind === 'state_input' && op.name === name)) {
+        throw new Error(`tensorgrad: state name '${name}' already used in this trace`);
+    }
+    const site = captureSite('stateInput');
+    return addOp(g, 'state_input', shape, dtype, site, { name, initValue });
+}
+//# sourceMappingURL=trace.js.map

package/dist/trace.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"trace.js","sourceRoot":"","sources":["../src/trace.ts"],"names":[],"mappings":"AAAA,6EAA6E;AAC7E,2EAA2E;AAC3E,EAAE;AACF,SAAS;AACT,EAAE;AACF,gCAAgC;AAChC,gDAAgD;AAChD,+CAA+C;AAC/C,mDAAmD;AACnD,6BAA6B;AAC7B,OAAO;AACP,EAAE;AACF,+EAA+E;AAC/E,+EAA+E;AAC/E,+CAA+C;AAG/C,OAAO,EAAE,SAAS,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,SAAS,CAAA;AAEvD,gFAAgF;AAChF,IAAI,QAAQ,GAAiB,IAAI,CAAA;AAEjC,MAAM,UAAU,YAAY;IAC1B,IAAI,CAAC,QAAQ,EAAE,CAAC;QACd,MAAM,IAAI,KAAK,CACb,qDAAqD;YACrD,2CAA2C,CAC5C,CAAA;IACH,CAAC;IACD,OAAO,QAAQ,CAAA;AACjB,CAAC;AAED,gFAAgF;AAChF,8EAA8E;AAC9E,MAAM,UAAU,KAAK,CAAC,EAA2B;IAC/C,IAAI,QAAQ,EAAE,CAAC;QACb,MAAM,IAAI,KAAK,CAAC,6CAA6C,CAAC,CAAA;IAChE,CAAC;IACD,MAAM,CAAC,GAAG,SAAS,EAAE,CAAA;IACrB,QAAQ,GAAG,CAAC,CAAA;IACZ,IAAI,CAAC;QACH,MAAM,MAAM,GAAG,EAAE,EAAE,CAAA;QACnB,MAAM,OAAO,GAAG,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAA;QACzD,KAAK,MAAM,CAAC,IAAI,OAAO,EAAE,CAAC;YACxB,CAAC;YAAC,CAAC,CAAC,OAAoB,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,CAAA;QACrC,CAAC;IACH,CAAC;YAAS,CAAC;QACT,QAAQ,GAAG,IAAI,CAAA;IACjB,CAAC;IACD,OAAO,CAAC,CAAA;AACV,CAAC;AAED,yEAAyE;AACzE,yEAAyE;AACzE,4EAA4E;AAC5E,iCAAiC;AACjC,MAAM,UAAU,SAAS,CAAI,CAAQ,EAAE,EAAW;IAChD,IAAI,QAAQ,EAAE,CAAC;QACb,MAAM,IAAI,KAAK,CAAC,8DAA8D,CAAC,CAAA;IACjF,CAAC;IACD,QAAQ,GAAG,CAAC,CAAA;IACZ,IAAI,CAAC;QACH,OAAO,EAAE,EAAE,CAAA;IACb,CAAC;YAAS,CAAC;QACT,QAAQ,GAAG,IAAI,CAAA;IACjB,CAAC;AACH,CAAC;AAED,+EAA+E;AAC/E,6EAA6E;AAC7E,4EAA4E;AAC5E,sBAAsB;AAEtB,MAAM,UAAU,UAAU,CAAC,IAAY,EAAE,KAAY,EAAE,QAAe,KAAK;IACzE,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,IAAI,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,IAAI,KAAK,aAAa,IAAI,EAAE,CAAC,IAAI,KAAK,cAAc,CAAC,IAAI,EAAE,CAAC,IAAI,KAAK,IAAI,CAAC,EAAE,CAAC;QACpG,MAAM,IAAI,KAAK,CAAC,2BAA2B,IAAI,8BAA8B,CAAC,CAAA;IAChF,CAAC;IACD,MAAM,IAAI,GAAG,WAAW,CAAC,YAAY,CAAC,CAAA;IACtC,OAAO,KAAK,CAAC,CAAC,EAAE,aAAa,EAAE,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,IAAI,EAAS,CAAC,CAAA;AACrE,CAAC;AAED,MAAM,UAAU,WAAW,CAAC,IAAY,EAAE,KAAY,EAAE,QAAe,KAAK;IAC1E,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,IAAI,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,IAAI,KAAK,aAAa,IAAI,EAAE,CAAC,IAAI,KAAK,cAAc,CAAC,IAAI,EAAE,CAAC,IAAI,KAAK,IAAI,CAAC,EAAE,CAAC;QACpG,MAAM,IAAI,KAAK,CAAC,2BAA2B,IAAI,8BAA8B,CAAC,CAAA;IAChF,CAAC;IACD,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,OAAO,KAAK,CAAC,CAAC,EAAE,cAAc,EAAE,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,IAAI,EAAS,CAAC,CAAA;AACtE,CAAC;AAED,uFAAuF;AACvF,mFAAmF;AACnF,MAAM,UAAU,UAAU,CAAC,IAAY,EAAE,KAAY,EAAE,QAAe,KAAK,EAAE,SAAS,GAAG,CAAC;IACxF,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,IAAI,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,IAAI,KAAK,aAAa,IAAI,EAAE,CAAC,IAAI,KAAK,IAAI,CAAC,EAAE,CAAC;QACpE,MAAM,IAAI,KAAK,CAAC,2BAA2B,IAAI,8BAA8B,CAAC,CAAA;IAChF,CAAC;IACD,MAAM,IAAI,GAAG,WAAW,CAAC,YAAY,CAAC,CAAA;IACtC,OAAO,KAAK,CAAC,CAAC,EAAE,aAAa,EAAE,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,IAAI,EAAE,SAAS,EAAS,CAAC,CAAA;AAChF,CAAC"}

package/package.json

@@ -0,0 +1,62 @@
+{
+  "name": "tensorgrad",
+  "version": "0.0.1",
+  "description": "Tiny TypeScript-native tensor library with autograd, compiling to WebGPU. Train small models in the browser without hand-writing kernels.",
+  "license": "MIT",
+  "author": "Ben Albahari",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/typebulb/tensorgrad.git"
+  },
+  "homepage": "https://github.com/typebulb/tensorgrad#readme",
+  "bugs": {
+    "url": "https://github.com/typebulb/tensorgrad/issues"
+  },
+  "keywords": [
+    "webgpu",
+    "machine-learning",
+    "autograd",
+    "tensor",
+    "neural-network",
+    "transformer",
+    "browser",
+    "typescript"
+  ],
+  "type": "module",
+  "main": "./src/index.ts",
+  "types": "./src/index.ts",
+  "exports": {
+    ".": {
+      "types": "./src/index.ts",
+      "default": "./src/index.ts"
+    }
+  },
+  "//": "Source-only exports for monorepo dev. publishConfig flips to dist/*.js for npm publish.",
+  "publishConfig": {
+    "main": "./dist/index.js",
+    "types": "./dist/index.d.ts",
+    "exports": {
+      ".": {
+        "types": "./dist/index.d.ts",
+        "default": "./dist/index.js"
+      }
+    }
+  },
+  "files": [
+    "dist",
+    "src",
+    "SPEC.md",
+    "README.md",
+    "LICENSE"
+  ],
+  "scripts": {
+    "build": "tsc -p tsconfig.json",
+    "typecheck": "tsc -p tsconfig.json --noEmit",
+    "test": "tsx test/smoke.ts",
+    "prepublishOnly": "tsc -p tsconfig.json"
+  },
+  "devDependencies": {
+    "tsx": "*",
+    "typescript": "*"
+  }
+}

package/src/adam.ts

@@ -0,0 +1,95 @@
+// Adam optimizer, in-graph.
+//
+// `appendAdam` extends a graph that already has a forward pass + autograd-emitted
+// backward (i.e., has paramGrads from `appendGrad`) with the Adam update math.
+//
+// Per parameter P with gradient g:
+//   m_new = b1 * m + (1 - b1) * g
+//   v_new = b2 * v + (1 - b2) * g²
+//   p_new = p - lr * m_new / (sqrt(v_new) + eps)
+//
+// This is "Adam without bias correction" — the `1 / (1 - β^t)` factors are
+// dropped because computing them in-graph requires per-step uniforms or
+// awkward exp/log tricks. In practice the omission only affects the first
+// ~100 steps; convergence is unaffected.
+//
+// Returns writeback declarations the buffer planner uses to wire up the
+// "after step, copy the new value into the persistent home" path. m and v
+// are state_inputs (zero-initialized, persistent across steps); the param
+// updates are aliased back to the param buffers.
+
+import type { Tensor } from './ir.js'
+import type { Graph } from './ir.js'
+import type { WritebackDecl } from './buffers.js'
+import { traceInto, stateInput, tensorInput } from './trace.js'
+import { adamUpdateM, adamUpdateV, adamUpdateP } from './ops.js'
+
+export interface AdamConfig {
+  lr: number
+  b1?: number   // default 0.9
+  b2?: number   // default 0.999
+  eps?: number  // default 1e-8
+}
+
+export interface AdamResult {
+  /** Writebacks the buffer planner should wire into the runtime. */
+  writebacks: WritebackDecl[]
+  /** Name of the per-step scalar tensor_input. The runtime fills this each call
+   * with `lr * sqrt(1-b2^t)/(1-b1^t)` (Adam's bias-corrected effective LR). */
+  lrtInputName: string
+  /** Hyperparameters as captured (so the runtime can compute lrt). */
+  config: Required<AdamConfig>
+}
+
+/**
+ * Append Adam update ops to `graph`. Must be called inside an active trace
+ * context (or after a trace, since traceInto re-enters the graph).
+ *
+ * @param graph the graph (already containing forward + backward)
+ * @param paramGrads param name -> gradient tensor (output of `appendGrad`)
+ * @param paramTensors param name -> the param's leaf Tensor (the param_input).
+ *                     Needed because the param_input lives in the graph but we
+ *                     don't have a direct map by name in `Graph` — caller passes it.
+ * @param config Adam hyperparameters
+ */
+export function appendAdam(
+  graph: Graph,
+  paramGrads: Record<string, Tensor>,
+  paramTensors: Record<string, Tensor>,
+  config: AdamConfig,
+): AdamResult {
+  const fullConfig: Required<AdamConfig> = {
+    lr: config.lr,
+    b1: config.b1 ?? 0.9,
+    b2: config.b2 ?? 0.999,
+    eps: config.eps ?? 1e-8,
+  }
+  const writebacks: WritebackDecl[] = []
+  const lrtInputName = '_adam_lrt'
+
+  return traceInto(graph, () => {
+    // One scalar lrt input shared by every adam_update_p call. Runtime supplies
+    // it per step as `lr * sqrt(1-b2^t) / (1-b1^t)`.
+    const lrt = tensorInput(lrtInputName, [], 'f32')
+
+    for (const name of Object.keys(paramGrads)) {
+      const p = paramTensors[name]
+      const g = paramGrads[name]
+      if (!p) throw new Error(`appendAdam: missing param tensor for '${name}'`)
+      if (!g) throw new Error(`appendAdam: missing gradient for '${name}'`)
+
+      const mState = stateInput(`adam_m_${name}`, p.shape, 'f32', 0)
+      const vState = stateInput(`adam_v_${name}`, p.shape, 'f32', 0)
+
+      // Three fused kernels per parameter — one for each of m_new / v_new / p_new.
+      const newM = adamUpdateM(mState, g, fullConfig.b1)
+      const newV = adamUpdateV(vState, g, fullConfig.b2)
+      const newP = adamUpdateP(p, newM, newV, lrt, fullConfig.eps)
+
+      writebacks.push({ source: newM, destName: `adam_m_${name}`, destKind: 'state' })
+      writebacks.push({ source: newV, destName: `adam_v_${name}`, destKind: 'state' })
+      writebacks.push({ source: newP, destName: name,             destKind: 'param' })
+    }
+    return { writebacks, lrtInputName, config: fullConfig }
+  })
+}