tensorgrad 0.0.15 → 0.0.16

package/dist/shape.js

@@ -1,259 +0,0 @@
-// 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

@@ -1 +0,0 @@
-{"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.js

@@ -1,100 +0,0 @@
-// 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;
-// Module-local: whether `capture(name, t)` calls should register on the current
-// graph. True only during the user's forward trace; false during `traceInto`
-// (autograd / optimizer ops shouldn't accidentally publish gradient tensors).
-let _captureEnabled = false;
-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;
-}
-export function isCaptureEnabled() {
-    return _captureEnabled;
-}
-// 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;
-    _captureEnabled = true;
-    try {
-        const result = fn();
-        const outputs = Array.isArray(result) ? result : [result];
-        for (const t of outputs) {
-            ;
-            g.outputs.push(t.id);
-        }
-    }
-    finally {
-        _current = null;
-        _captureEnabled = false;
-    }
-    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.
-// Capture is intentionally disabled here — backward / optimizer rules
-// shouldn't publish their internal tensors via `capture()`.
-// Returns whatever `fn` returns.
-export function traceInto(g, fn) {
-    if (_current) {
-        throw new Error('tensorgrad: traceInto() called while another trace is active');
-    }
-    _current = g;
-    // _captureEnabled stays false (default) — explicit, but not toggled.
-    try {
-        return fn();
-    }
-    finally {
-        _current = null;
-    }
-}
-function assertNameUnused(g, name, kinds, label) {
-    if (g.ops.some(op => kinds.includes(op.kind) && op.name === name)) {
-        throw new Error(`tensorgrad: ${label} name '${name}' already used in this trace`);
-    }
-}
-export function paramInput(name, shape, dtype = 'f32') {
-    const g = currentGraph();
-    assertNameUnused(g, name, ['param_input', 'tensor_input'], 'input');
-    const site = captureSite('paramInput');
-    return addOp(g, 'param_input', shape, dtype, site, { name });
-}
-export function tensorInput(name, shape, dtype = 'f32') {
-    const g = currentGraph();
-    assertNameUnused(g, name, ['param_input', 'tensor_input'], 'input');
-    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();
-    assertNameUnused(g, name, ['state_input'], 'state');
-    const site = captureSite('stateInput');
-    return addOp(g, 'state_input', shape, dtype, site, { name, initValue });
-}
-//# sourceMappingURL=trace.js.map

package/dist/trace.js.map

@@ -1 +0,0 @@
-{"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;AACjC,gFAAgF;AAChF,6EAA6E;AAC7E,8EAA8E;AAC9E,IAAI,eAAe,GAAG,KAAK,CAAA;AAE3B,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,MAAM,UAAU,gBAAgB;IAC9B,OAAO,eAAe,CAAA;AACxB,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,eAAe,GAAG,IAAI,CAAA;IACtB,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;QACf,eAAe,GAAG,KAAK,CAAA;IACzB,CAAC;IACD,OAAO,CAAC,CAAA;AACV,CAAC;AAED,yEAAyE;AACzE,yEAAyE;AACzE,4EAA4E;AAC5E,sEAAsE;AACtE,4DAA4D;AAC5D,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,qEAAqE;IACrE,IAAI,CAAC;QACH,OAAO,EAAE,EAAE,CAAA;IACb,CAAC;YAAS,CAAC;QACT,QAAQ,GAAG,IAAI,CAAA;IACjB,CAAC;AACH,CAAC;AAUD,SAAS,gBAAgB,CAAC,CAAQ,EAAE,IAAY,EAAE,KAAuB,EAAE,KAAa;IACtF,IAAI,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,CAAC,KAAK,CAAC,QAAQ,CAAC,EAAE,CAAC,IAAsB,CAAC,IAAK,EAAwB,CAAC,IAAI,KAAK,IAAI,CAAC,EAAE,CAAC;QAC3G,MAAM,IAAI,KAAK,CAAC,eAAe,KAAK,UAAU,IAAI,8BAA8B,CAAC,CAAA;IACnF,CAAC;AACH,CAAC;AAED,MAAM,UAAU,UAAU,CAAC,IAAY,EAAE,KAAY,EAAE,QAAe,KAAK;IACzE,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,gBAAgB,CAAC,CAAC,EAAE,IAAI,EAAE,CAAC,aAAa,EAAE,cAAc,CAAC,EAAE,OAAO,CAAC,CAAA;IACnE,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,gBAAgB,CAAC,CAAC,EAAE,IAAI,EAAE,CAAC,aAAa,EAAE,cAAc,CAAC,EAAE,OAAO,CAAC,CAAA;IACnE,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,gBAAgB,CAAC,CAAC,EAAE,IAAI,EAAE,CAAC,aAAa,CAAC,EAAE,OAAO,CAAC,CAAA;IACnD,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"}