tensorgrad 0.0.1

package/dist/ops.js

@@ -0,0 +1,270 @@
+// User-facing op surface.
+//
+// Each function here is a thin wrapper:
+//   1. capture the call site (for error attribution)
+//   2. validate input shapes via src/shape.ts (which throws on mismatch)
+//   3. compute the output shape and dtype
+//   4. append the op to the current Graph (held in module state by src/trace.ts)
+//   5. return the produced Tensor handle
+//
+// No actual numeric work happens here. These calls just build the IR.
+import { addOp, captureSite } from './ir.js';
+import { currentGraph } from './trace.js';
+import { inferElementwiseBinop, inferUnary, inferMeanLast, inferSumLast, inferReshape, inferTranspose, inferMatmul, inferMatmulBatched, inferOneHot, inferWhereCausal, inferSliceLastRange, inferBroadcastTo, inferSumToShape, inferReluGrad, inferWhere, ShapeError, } from './shape.js';
+// ----------------------------------------------------------------------------
+// Element-wise binops (add/sub/mul/div). Trailing-suffix broadcast.
+// ----------------------------------------------------------------------------
+/**
+ * Build an element-wise binop op (forward declaration only — appends to the
+ * graph). Used by both arithmetic ops (add/sub/mul/div, output dtype = input
+ * dtype) and comparisons (less/greater, output dtype = bool).
+ */
+function binopOp(name, kind, a, b, outDtype = a.dtype) {
+    const site = captureSite(name);
+    if (a.dtype !== b.dtype)
+        throw new ShapeError(`${name}: dtype mismatch (${a.dtype} vs ${b.dtype})`, site);
+    const outShape = inferElementwiseBinop(name, a.shape, b.shape, site);
+    return addOp(currentGraph(), kind, outShape, outDtype, site, { a: a.id, b: b.id });
+}
+// Element-wise binops. Second arg can be a Tensor or a JS number; the latter
+// dispatches to scalar-fused IR ops internally. `mul(x, 2)` and `mul(x, y)`
+// both work — matches every NumPy-shaped library.
+export function add(a, b) {
+    return typeof b === 'number' ? addScalar(a, b) : binopOp('add', 'add', a, b);
+}
+export function sub(a, b) {
+    return typeof b === 'number' ? addScalar(a, -b) : binopOp('sub', 'sub', a, b);
+}
+export function mul(a, b) {
+    return typeof b === 'number' ? mulScalar(a, b) : binopOp('mul', 'mul', a, b);
+}
+export function div(a, b) {
+    if (typeof b === 'number') {
+        if (b === 0)
+            throw new ShapeError(`div: scalar divisor cannot be zero`, captureSite('div'));
+        return mulScalar(a, 1 / b);
+    }
+    return binopOp('div', 'div', a, b);
+}
+// ----------------------------------------------------------------------------
+// Element-wise scalar binops (mul/add by JS number). Used for things like
+// `scores * (1/sqrt(d))` and `logits + 1e-5` where allocating a 0-d tensor
+// for the scalar is wasteful.
+// ----------------------------------------------------------------------------
+export function mulScalar(a, scalar) {
+    const site = captureSite('mulScalar');
+    return addOp(currentGraph(), 'mul_scalar', a.shape, a.dtype, site, { a: a.id, scalar });
+}
+export function addScalar(a, scalar) {
+    const site = captureSite('addScalar');
+    return addOp(currentGraph(), 'add_scalar', a.shape, a.dtype, site, { a: a.id, scalar });
+}
+// ----------------------------------------------------------------------------
+// Unary ops.
+// ----------------------------------------------------------------------------
+function unary(name, a) {
+    const site = captureSite(name);
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`${name}: requires f32, got ${a.dtype}`, site);
+    return addOp(currentGraph(), name, inferUnary(name, a.shape, site), 'f32', site, { a: a.id });
+}
+export const sqrt = (a) => unary('sqrt', a);
+export const rsqrt = (a) => unary('rsqrt', a);
+export const log = (a) => unary('log', a);
+export const exp = (a) => unary('exp', a);
+export const relu = (a) => unary('relu', a);
+// ----------------------------------------------------------------------------
+// Reductions over the last axis. To reduce along other axes, transpose first.
+// (This is intentional — keeps codegen and autograd small.)
+// ----------------------------------------------------------------------------
+export function meanLast(a) {
+    const site = captureSite('meanLast');
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`meanLast: requires f32, got ${a.dtype}`, site);
+    const outShape = inferMeanLast('meanLast', a.shape, site);
+    return addOp(currentGraph(), 'mean_last', outShape, a.dtype, site, { a: a.id });
+}
+export function sumLast(a) {
+    const site = captureSite('sumLast');
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`sumLast: requires f32, got ${a.dtype}`, site);
+    const outShape = inferSumLast('sumLast', a.shape, site);
+    return addOp(currentGraph(), 'sum_last', outShape, a.dtype, site, { a: a.id });
+}
+// ----------------------------------------------------------------------------
+// Shape ops.
+// ----------------------------------------------------------------------------
+export function reshape(a, newShape) {
+    const site = captureSite('reshape');
+    const outShape = inferReshape('reshape', a.shape, newShape, site);
+    return addOp(currentGraph(), 'reshape', outShape, a.dtype, site, { a: a.id, newShape: outShape });
+}
+export function transpose(a, perm) {
+    const site = captureSite('transpose');
+    const outShape = inferTranspose('transpose', a.shape, perm, site);
+    return addOp(currentGraph(), 'transpose', outShape, a.dtype, site, { a: a.id, perm });
+}
+// ----------------------------------------------------------------------------
+// Linear algebra.
+// ----------------------------------------------------------------------------
+export function matmul(a, b) {
+    const site = captureSite('matmul');
+    if (a.dtype !== 'f32' || b.dtype !== 'f32') {
+        throw new ShapeError(`matmul: requires f32, got ${a.dtype} and ${b.dtype}`, site);
+    }
+    const outShape = inferMatmul('matmul', a.shape, b.shape, site);
+    return addOp(currentGraph(), 'matmul', outShape, 'f32', site, { a: a.id, b: b.id });
+}
+export function matmulBatched(a, b) {
+    const site = captureSite('matmulBatched');
+    if (a.dtype !== 'f32' || b.dtype !== 'f32') {
+        throw new ShapeError(`matmulBatched: requires f32, got ${a.dtype} and ${b.dtype}`, site);
+    }
+    const outShape = inferMatmulBatched('matmulBatched', a.shape, b.shape, site);
+    return addOp(currentGraph(), 'matmul_batched', outShape, 'f32', site, { a: a.id, b: b.id });
+}
+// ----------------------------------------------------------------------------
+// Indexing / casting.
+// ----------------------------------------------------------------------------
+export function oneHot(indices, depth, dtype = 'f32') {
+    const site = captureSite('oneHot');
+    if (indices.dtype !== 'i32') {
+        throw new ShapeError(`oneHot: indices must be i32, got ${indices.dtype}`, site);
+    }
+    const outShape = inferOneHot('oneHot', indices.shape, depth, site);
+    return addOp(currentGraph(), 'one_hot', outShape, dtype, site, { indices: indices.id, depth, dtype });
+}
+// arange(n) → [n] of values [0, 1, ..., n-1]. Used for position embeddings.
+export function arange(n, dtype = 'i32') {
+    const site = captureSite('arange');
+    if (n <= 0 || !Number.isInteger(n)) {
+        throw new ShapeError(`arange: n must be a positive integer, got ${n}`, site);
+    }
+    return addOp(currentGraph(), 'arange', [n], dtype, site, { n, dtype });
+}
+// ----------------------------------------------------------------------------
+// ML primitives. Fused so autograd's transpose rule is straightforward and the
+// kernels can be hand-tuned for our specific shapes.
+// ----------------------------------------------------------------------------
+// Causal-masked softmax along the last axis. Shape preserved. Last two axes
+// must be square (TxT attention scores).
+export function softmaxCausalLast(a) {
+    const site = captureSite('softmaxCausalLast');
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`softmaxCausalLast: requires f32, got ${a.dtype}`, site);
+    inferWhereCausal('softmaxCausalLast', a.shape, site); // shape check (square last 2 axes)
+    return addOp(currentGraph(), 'softmax_causal_last', a.shape, 'f32', site, { a: a.id });
+}
+// Numerically-stable log-softmax along the last axis. Shape preserved.
+export function logSoftmaxLast(a) {
+    const site = captureSite('logSoftmaxLast');
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`logSoftmaxLast: requires f32, got ${a.dtype}`, site);
+    return addOp(currentGraph(), 'log_softmax_last', a.shape, 'f32', site, { a: a.id });
+}
+// Pre-softmax causal mask. Sets cells where (i < j) on the last two axes to
+// `fillValue` (typically -1e30). Lower-triangle entries pass through.
+// Use this when you want the masked scores explicitly (e.g. for capture);
+// for the common case, prefer softmaxCausalLast which fuses both.
+export function whereCausal(a, fillValue) {
+    const site = captureSite('whereCausal');
+    if (a.dtype !== 'f32')
+        throw new ShapeError(`whereCausal: requires f32, got ${a.dtype}`, site);
+    inferWhereCausal('whereCausal', a.shape, site);
+    return addOp(currentGraph(), 'where_causal', a.shape, 'f32', site, { a: a.id, fillValue });
+}
+// ----------------------------------------------------------------------------
+// Slicing.
+// ----------------------------------------------------------------------------
+// sliceLastRange(a, start, end): slice [start, end) along the last axis.
+// Used for splitting Q/K/V from a fused QKV matmul.
+export function sliceLastRange(a, start, end) {
+    const site = captureSite('sliceLastRange');
+    const outShape = inferSliceLastRange('sliceLastRange', a.shape, start, end, site);
+    return addOp(currentGraph(), 'slice_last_range', outShape, a.dtype, site, { a: a.id, start, end });
+}
+// ----------------------------------------------------------------------------
+// Broadcast / un-broadcast. Mostly used by autograd, but exposed in case user
+// code needs them (e.g. explicit broadcasting for clarity).
+// ----------------------------------------------------------------------------
+export function broadcastTo(a, targetShape) {
+    const site = captureSite('broadcastTo');
+    inferBroadcastTo('broadcastTo', a.shape, targetShape, site);
+    return addOp(currentGraph(), 'broadcast_to', targetShape, a.dtype, site, { a: a.id, targetShape });
+}
+export function sumToShape(a, targetShape) {
+    const site = captureSite('sumToShape');
+    inferSumToShape('sumToShape', a.shape, targetShape, site);
+    return addOp(currentGraph(), 'sum_to_shape', targetShape, a.dtype, site, { a: a.id, targetShape });
+}
+// ----------------------------------------------------------------------------
+// Constants.
+// ----------------------------------------------------------------------------
+// 0-d tensor with a constant value. Used by autograd to seed the loss cotangent.
+export function constScalar(value, dtype = 'f32') {
+    const site = captureSite('constScalar');
+    return addOp(currentGraph(), 'const_scalar', [], dtype, site, { value, dtype });
+}
+// ----------------------------------------------------------------------------
+// Autograd-internal helpers (exposed for users writing custom transpose rules).
+// ----------------------------------------------------------------------------
+// ----------------------------------------------------------------------------
+// Comparisons and selection.
+// ----------------------------------------------------------------------------
+// Comparisons reuse the binop helper but return bool.
+export const less = (a, b) => binopOp('less', 'less', a, b, 'bool');
+export const greater = (a, b) => binopOp('greater', 'greater', a, b, 'bool');
+// where(cond, a, b): elementwise select. cond is bool; a and b can be any matching dtype.
+export function where(cond, a, b) {
+    const site = captureSite('where');
+    if (cond.dtype !== 'bool')
+        throw new ShapeError(`where: cond must be bool, got ${cond.dtype}`, site);
+    if (a.dtype !== b.dtype)
+        throw new ShapeError(`where: a/b dtype mismatch (${a.dtype} vs ${b.dtype})`, site);
+    const outShape = inferWhere('where', cond.shape, a.shape, b.shape, site);
+    return addOp(currentGraph(), 'where', outShape, a.dtype, site, { cond: cond.id, a: a.id, b: b.id });
+}
+// reluGrad(x, dy) = dy where x > 0, else 0. Same shape as x. This is the
+// transpose rule for relu, exposed as an op so codegen can emit it.
+export function reluGrad(x, dy) {
+    const site = captureSite('reluGrad');
+    if (x.dtype !== 'f32' || dy.dtype !== 'f32') {
+        throw new ShapeError(`reluGrad: requires f32, got ${x.dtype} and ${dy.dtype}`, site);
+    }
+    const outShape = inferReluGrad('reluGrad', x.shape, dy.shape, site);
+    return addOp(currentGraph(), 'relu_grad', outShape, 'f32', site, { x: x.id, dy: dy.id });
+}
+// ----------------------------------------------------------------------------
+// Adam-fused ops. Each does its full per-element update in one kernel.
+// ----------------------------------------------------------------------------
+export function adamUpdateM(m, g, b1) {
+    const site = captureSite('adamUpdateM');
+    if (m.dtype !== 'f32' || g.dtype !== 'f32')
+        throw new ShapeError(`adamUpdateM: requires f32`, site);
+    if (m.shape.length !== g.shape.length || m.shape.some((d, i) => d !== g.shape[i])) {
+        throw new ShapeError(`adamUpdateM: shape mismatch`, site);
+    }
+    return addOp(currentGraph(), 'adam_update_m', m.shape, 'f32', site, { m: m.id, g: g.id, b1 });
+}
+export function adamUpdateV(v, g, b2) {
+    const site = captureSite('adamUpdateV');
+    if (v.dtype !== 'f32' || g.dtype !== 'f32')
+        throw new ShapeError(`adamUpdateV: requires f32`, site);
+    if (v.shape.length !== g.shape.length || v.shape.some((d, i) => d !== g.shape[i])) {
+        throw new ShapeError(`adamUpdateV: shape mismatch`, site);
+    }
+    return addOp(currentGraph(), 'adam_update_v', v.shape, 'f32', site, { v: v.id, g: g.id, b2 });
+}
+export function adamUpdateP(p, mNew, vNew, lrt, eps) {
+    const site = captureSite('adamUpdateP');
+    if (p.dtype !== 'f32')
+        throw new ShapeError(`adamUpdateP: requires f32`, site);
+    if (lrt.dtype !== 'f32' || lrt.shape.length !== 0) {
+        throw new ShapeError(`adamUpdateP: lrt must be a 0-d f32 scalar`, site);
+    }
+    if (p.shape.length !== mNew.shape.length || p.shape.some((d, i) => d !== mNew.shape[i])) {
+        throw new ShapeError(`adamUpdateP: p/mNew shape mismatch`, site);
+    }
+    return addOp(currentGraph(), 'adam_update_p', p.shape, 'f32', site, { p: p.id, mNew: mNew.id, vNew: vNew.id, lrt: lrt.id, eps });
+}
+//# sourceMappingURL=ops.js.map

package/dist/ops.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"ops.js","sourceRoot":"","sources":["../src/ops.ts"],"names":[],"mappings":"AAAA,0BAA0B;AAC1B,EAAE;AACF,wCAAwC;AACxC,qDAAqD;AACrD,yEAAyE;AACzE,0CAA0C;AAC1C,iFAAiF;AACjF,yCAAyC;AACzC,EAAE;AACF,sEAAsE;AAGtE,OAAO,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,SAAS,CAAA;AAC5C,OAAO,EAAE,YAAY,EAAE,MAAM,YAAY,CAAA;AACzC,OAAO,EACL,qBAAqB,EAAE,UAAU,EAAE,aAAa,EAAE,YAAY,EAC9D,YAAY,EAAE,cAAc,EAAE,WAAW,EAAE,kBAAkB,EAC7D,WAAW,EAAE,gBAAgB,EAAE,mBAAmB,EAClD,gBAAgB,EAAE,eAAe,EAAE,aAAa,EAAE,UAAU,EAC5D,UAAU,GACX,MAAM,YAAY,CAAA;AAEnB,+EAA+E;AAC/E,oEAAoE;AACpE,+EAA+E;AAE/E;;;;GAIG;AACH,SAAS,OAAO,CACd,IAAY,EACZ,IAAoB,EACpB,CAAS,EAAE,CAAS,EACpB,WAAkB,CAAC,CAAC,KAAK;IAEzB,MAAM,IAAI,GAAG,WAAW,CAAC,IAAI,CAAC,CAAA;IAC9B,IAAI,CAAC,CAAC,KAAK,KAAK,CAAC,CAAC,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,GAAG,IAAI,qBAAqB,CAAC,CAAC,KAAK,OAAO,CAAC,CAAC,KAAK,GAAG,EAAE,IAAI,CAAC,CAAA;IACzG,MAAM,QAAQ,GAAG,qBAAqB,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACpE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE,QAAQ,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACpF,CAAC;AAED,6EAA6E;AAC7E,4EAA4E;AAC5E,kDAAkD;AAClD,MAAM,UAAU,GAAG,CAAC,CAAS,EAAE,CAAkB;IAC/C,OAAO,OAAO,CAAC,KAAK,QAAQ,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AAC9E,CAAC;AACD,MAAM,UAAU,GAAG,CAAC,CAAS,EAAE,CAAkB;IAC/C,OAAO,OAAO,CAAC,KAAK,QAAQ,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AAC/E,CAAC;AACD,MAAM,UAAU,GAAG,CAAC,CAAS,EAAE,CAAkB;IAC/C,OAAO,OAAO,CAAC,KAAK,QAAQ,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AAC9E,CAAC;AACD,MAAM,UAAU,GAAG,CAAC,CAAS,EAAE,CAAkB;IAC/C,IAAI,OAAO,CAAC,KAAK,QAAQ,EAAE,CAAC;QAC1B,IAAI,CAAC,KAAK,CAAC;YAAE,MAAM,IAAI,UAAU,CAAC,oCAAoC,EAAE,WAAW,CAAC,KAAK,CAAC,CAAC,CAAA;QAC3F,OAAO,SAAS,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,CAAA;IAC5B,CAAC;IACD,OAAO,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,CAAC,CAAA;AACpC,CAAC;AAED,+EAA+E;AAC/E,0EAA0E;AAC1E,2EAA2E;AAC3E,8BAA8B;AAC9B,+EAA+E;AAE/E,MAAM,UAAU,SAAS,CAAC,CAAS,EAAE,MAAc;IACjD,MAAM,IAAI,GAAG,WAAW,CAAC,WAAW,CAAC,CAAA;IACrC,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,YAAY,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,MAAM,EAAE,CAAC,CAAA;AACzF,CAAC;AAED,MAAM,UAAU,SAAS,CAAC,CAAS,EAAE,MAAc;IACjD,MAAM,IAAI,GAAG,WAAW,CAAC,WAAW,CAAC,CAAA;IACrC,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,YAAY,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,MAAM,EAAE,CAAC,CAAA;AACzF,CAAC;AAED,+EAA+E;AAC/E,aAAa;AACb,+EAA+E;AAE/E,SAAS,KAAK,CAAC,IAA+C,EAAE,CAAS;IACvE,MAAM,IAAI,GAAG,WAAW,CAAC,IAAI,CAAC,CAAA;IAC9B,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,GAAG,IAAI,uBAAuB,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC1F,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,IAAI,EAAE,UAAU,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AAC/F,CAAC;AAED,MAAM,CAAC,MAAM,IAAI,GAAI,CAAC,CAAS,EAAU,EAAE,CAAC,KAAK,CAAC,MAAM,EAAG,CAAC,CAAC,CAAA;AAC7D,MAAM,CAAC,MAAM,KAAK,GAAG,CAAC,CAAS,EAAU,EAAE,CAAC,KAAK,CAAC,OAAO,EAAE,CAAC,CAAC,CAAA;AAC7D,MAAM,CAAC,MAAM,GAAG,GAAK,CAAC,CAAS,EAAU,EAAE,CAAC,KAAK,CAAC,KAAK,EAAI,CAAC,CAAC,CAAA;AAC7D,MAAM,CAAC,MAAM,GAAG,GAAK,CAAC,CAAS,EAAU,EAAE,CAAC,KAAK,CAAC,KAAK,EAAI,CAAC,CAAC,CAAA;AAC7D,MAAM,CAAC,MAAM,IAAI,GAAI,CAAC,CAAS,EAAU,EAAE,CAAC,KAAK,CAAC,MAAM,EAAG,CAAC,CAAC,CAAA;AAE7D,+EAA+E;AAC/E,8EAA8E;AAC9E,4DAA4D;AAC5D,+EAA+E;AAE/E,MAAM,UAAU,QAAQ,CAAC,CAAS;IAChC,MAAM,IAAI,GAAG,WAAW,CAAC,UAAU,CAAC,CAAA;IACpC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,+BAA+B,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC3F,MAAM,QAAQ,GAAG,aAAa,CAAC,UAAU,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACzD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,WAAW,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACjF,CAAC;AAED,MAAM,UAAU,OAAO,CAAC,CAAS;IAC/B,MAAM,IAAI,GAAG,WAAW,CAAC,SAAS,CAAC,CAAA;IACnC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,8BAA8B,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC1F,MAAM,QAAQ,GAAG,YAAY,CAAC,SAAS,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACvD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,UAAU,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AAChF,CAAC;AAED,+EAA+E;AAC/E,aAAa;AACb,+EAA+E;AAE/E,MAAM,UAAU,OAAO,CAAC,CAAS,EAAE,QAAe;IAChD,MAAM,IAAI,GAAG,WAAW,CAAC,SAAS,CAAC,CAAA;IACnC,MAAM,QAAQ,GAAG,YAAY,CAAC,SAAS,EAAE,CAAC,CAAC,KAAK,EAAE,QAAQ,EAAE,IAAI,CAAC,CAAA;IACjE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,SAAS,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,QAAQ,EAAE,QAAQ,EAAE,CAAC,CAAA;AACnG,CAAC;AAED,MAAM,UAAU,SAAS,CAAC,CAAS,EAAE,IAAuB;IAC1D,MAAM,IAAI,GAAG,WAAW,CAAC,WAAW,CAAC,CAAA;IACrC,MAAM,QAAQ,GAAG,cAAc,CAAC,WAAW,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,IAAI,CAAC,CAAA;IACjE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,WAAW,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,IAAI,EAAE,CAAC,CAAA;AACvF,CAAC;AAED,+EAA+E;AAC/E,kBAAkB;AAClB,+EAA+E;AAE/E,MAAM,UAAU,MAAM,CAAC,CAAS,EAAE,CAAS;IACzC,MAAM,IAAI,GAAG,WAAW,CAAC,QAAQ,CAAC,CAAA;IAClC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,EAAE,CAAC;QAC3C,MAAM,IAAI,UAAU,CAAC,6BAA6B,CAAC,CAAC,KAAK,QAAQ,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACnF,CAAC;IACD,MAAM,QAAQ,GAAG,WAAW,CAAC,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IAC9D,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,QAAQ,EAAE,QAAQ,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACrF,CAAC;AAED,MAAM,UAAU,aAAa,CAAC,CAAS,EAAE,CAAS;IAChD,MAAM,IAAI,GAAG,WAAW,CAAC,eAAe,CAAC,CAAA;IACzC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,EAAE,CAAC;QAC3C,MAAM,IAAI,UAAU,CAAC,oCAAoC,CAAC,CAAC,KAAK,QAAQ,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC1F,CAAC;IACD,MAAM,QAAQ,GAAG,kBAAkB,CAAC,eAAe,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IAC5E,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,gBAAgB,EAAE,QAAQ,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AAC7F,CAAC;AAED,+EAA+E;AAC/E,sBAAsB;AACtB,+EAA+E;AAE/E,MAAM,UAAU,MAAM,CAAC,OAAe,EAAE,KAAa,EAAE,QAAe,KAAK;IACzE,MAAM,IAAI,GAAG,WAAW,CAAC,QAAQ,CAAC,CAAA;IAClC,IAAI,OAAO,CAAC,KAAK,KAAK,KAAK,EAAE,CAAC;QAC5B,MAAM,IAAI,UAAU,CAAC,oCAAoC,OAAO,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACjF,CAAC;IACD,MAAM,QAAQ,GAAG,WAAW,CAAC,QAAQ,EAAE,OAAO,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,CAAC,CAAA;IAClE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,SAAS,EAAE,QAAQ,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,OAAO,EAAE,OAAO,CAAC,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAA;AACvG,CAAC;AAED,4EAA4E;AAC5E,MAAM,UAAU,MAAM,CAAC,CAAS,EAAE,QAAe,KAAK;IACpD,MAAM,IAAI,GAAG,WAAW,CAAC,QAAQ,CAAC,CAAA;IAClC,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,SAAS,CAAC,CAAC,CAAC,EAAE,CAAC;QACnC,MAAM,IAAI,UAAU,CAAC,6CAA6C,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IAC9E,CAAC;IACD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,QAAQ,EAAE,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,CAAA;AACxE,CAAC;AAED,+EAA+E;AAC/E,+EAA+E;AAC/E,qDAAqD;AACrD,+EAA+E;AAE/E,4EAA4E;AAC5E,yCAAyC;AACzC,MAAM,UAAU,iBAAiB,CAAC,CAAS;IACzC,MAAM,IAAI,GAAG,WAAW,CAAC,mBAAmB,CAAC,CAAA;IAC7C,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,wCAAwC,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACpG,gBAAgB,CAAC,mBAAmB,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA,CAAE,mCAAmC;IACzF,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,qBAAqB,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACxF,CAAC;AAED,uEAAuE;AACvE,MAAM,UAAU,cAAc,CAAC,CAAS;IACtC,MAAM,IAAI,GAAG,WAAW,CAAC,gBAAgB,CAAC,CAAA;IAC1C,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,qCAAqC,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACjG,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,kBAAkB,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACrF,CAAC;AAED,4EAA4E;AAC5E,sEAAsE;AACtE,0EAA0E;AAC1E,kEAAkE;AAClE,MAAM,UAAU,WAAW,CAAC,CAAS,EAAE,SAAiB;IACtD,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,kCAAkC,CAAC,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IAC9F,gBAAgB,CAAC,aAAa,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IAC9C,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,cAAc,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,SAAS,EAAE,CAAC,CAAA;AAC5F,CAAC;AAED,+EAA+E;AAC/E,WAAW;AACX,+EAA+E;AAE/E,yEAAyE;AACzE,oDAAoD;AACpD,MAAM,UAAU,cAAc,CAAC,CAAS,EAAE,KAAa,EAAE,GAAW;IAClE,MAAM,IAAI,GAAG,WAAW,CAAC,gBAAgB,CAAC,CAAA;IAC1C,MAAM,QAAQ,GAAG,mBAAmB,CAAC,gBAAgB,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,GAAG,EAAE,IAAI,CAAC,CAAA;IACjF,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,kBAAkB,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,KAAK,EAAE,GAAG,EAAE,CAAC,CAAA;AACpG,CAAC;AAED,+EAA+E;AAC/E,8EAA8E;AAC9E,4DAA4D;AAC5D,+EAA+E;AAE/E,MAAM,UAAU,WAAW,CAAC,CAAS,EAAE,WAAkB;IACvD,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,gBAAgB,CAAC,aAAa,EAAE,CAAC,CAAC,KAAK,EAAE,WAAW,EAAE,IAAI,CAAC,CAAA;IAC3D,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,cAAc,EAAE,WAAW,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,WAAW,EAAE,CAAC,CAAA;AACpG,CAAC;AAED,MAAM,UAAU,UAAU,CAAC,CAAS,EAAE,WAAkB;IACtD,MAAM,IAAI,GAAG,WAAW,CAAC,YAAY,CAAC,CAAA;IACtC,eAAe,CAAC,YAAY,EAAE,CAAC,CAAC,KAAK,EAAE,WAAW,EAAE,IAAI,CAAC,CAAA;IACzD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,cAAc,EAAE,WAAW,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,WAAW,EAAE,CAAC,CAAA;AACpG,CAAC;AAED,+EAA+E;AAC/E,aAAa;AACb,+EAA+E;AAE/E,iFAAiF;AACjF,MAAM,UAAU,WAAW,CAAC,KAAa,EAAE,QAAe,KAAK;IAC7D,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,cAAc,EAAE,EAAE,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAA;AACjF,CAAC;AAED,+EAA+E;AAC/E,gFAAgF;AAChF,+EAA+E;AAE/E,+EAA+E;AAC/E,6BAA6B;AAC7B,+EAA+E;AAE/E,sDAAsD;AACtD,MAAM,CAAC,MAAM,IAAI,GAAM,CAAC,CAAS,EAAE,CAAS,EAAU,EAAE,CAAC,OAAO,CAAC,MAAM,EAAK,MAAM,EAAK,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAA;AACpG,MAAM,CAAC,MAAM,OAAO,GAAG,CAAC,CAAS,EAAE,CAAS,EAAU,EAAE,CAAC,OAAO,CAAC,SAAS,EAAE,SAAS,EAAE,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAA;AAEpG,0FAA0F;AAC1F,MAAM,UAAU,KAAK,CAAC,IAAY,EAAE,CAAS,EAAE,CAAS;IACtD,MAAM,IAAI,GAAG,WAAW,CAAC,OAAO,CAAC,CAAA;IACjC,IAAI,IAAI,CAAC,KAAK,KAAK,MAAM;QAAE,MAAM,IAAI,UAAU,CAAC,iCAAiC,IAAI,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACpG,IAAI,CAAC,CAAC,KAAK,KAAK,CAAC,CAAC,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,8BAA8B,CAAC,CAAC,KAAK,OAAO,CAAC,CAAC,KAAK,GAAG,EAAE,IAAI,CAAC,CAAA;IAC3G,MAAM,QAAQ,GAAG,UAAU,CAAC,OAAO,EAAE,IAAI,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAC,CAAC,KAAK,EAAE,IAAI,EAAE,EAAE,IAAI,EAAE,IAAI,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,CAAA;AACrG,CAAC;AAED,yEAAyE;AACzE,oEAAoE;AACpE,MAAM,UAAU,QAAQ,CAAC,CAAS,EAAE,EAAU;IAC5C,MAAM,IAAI,GAAG,WAAW,CAAC,UAAU,CAAC,CAAA;IACpC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,IAAI,EAAE,CAAC,KAAK,KAAK,KAAK,EAAE,CAAC;QAC5C,MAAM,IAAI,UAAU,CAAC,+BAA+B,CAAC,CAAC,KAAK,QAAQ,EAAE,CAAC,KAAK,EAAE,EAAE,IAAI,CAAC,CAAA;IACtF,CAAC;IACD,MAAM,QAAQ,GAAG,aAAa,CAAC,UAAU,EAAE,CAAC,CAAC,KAAK,EAAE,EAAE,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACnE,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,WAAW,EAAE,QAAQ,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,EAAE,EAAE,CAAC,CAAA;AAC1F,CAAC;AAED,+EAA+E;AAC/E,uEAAuE;AACvE,+EAA+E;AAE/E,MAAM,UAAU,WAAW,CAAC,CAAS,EAAE,CAAS,EAAE,EAAU;IAC1D,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,2BAA2B,EAAE,IAAI,CAAC,CAAA;IACnG,IAAI,CAAC,CAAC,KAAK,CAAC,MAAM,KAAK,CAAC,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC;QAClF,MAAM,IAAI,UAAU,CAAC,6BAA6B,EAAE,IAAI,CAAC,CAAA;IAC3D,CAAC;IACD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,eAAe,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AAC/F,CAAC;AAED,MAAM,UAAU,WAAW,CAAC,CAAS,EAAE,CAAS,EAAE,EAAU;IAC1D,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,2BAA2B,EAAE,IAAI,CAAC,CAAA;IACnG,IAAI,CAAC,CAAC,KAAK,CAAC,MAAM,KAAK,CAAC,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC;QAClF,MAAM,IAAI,UAAU,CAAC,6BAA6B,EAAE,IAAI,CAAC,CAAA;IAC3D,CAAC;IACD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,eAAe,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AAC/F,CAAC;AAED,MAAM,UAAU,WAAW,CAAC,CAAS,EAAE,IAAY,EAAE,IAAY,EAAE,GAAW,EAAE,GAAW;IACzF,MAAM,IAAI,GAAG,WAAW,CAAC,aAAa,CAAC,CAAA;IACvC,IAAI,CAAC,CAAC,KAAK,KAAK,KAAK;QAAE,MAAM,IAAI,UAAU,CAAC,2BAA2B,EAAE,IAAI,CAAC,CAAA;IAC9E,IAAI,GAAG,CAAC,KAAK,KAAK,KAAK,IAAI,GAAG,CAAC,KAAK,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QAClD,MAAM,IAAI,UAAU,CAAC,2CAA2C,EAAE,IAAI,CAAC,CAAA;IACzE,CAAC;IACD,IAAI,CAAC,CAAC,KAAK,CAAC,MAAM,KAAK,IAAI,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,KAAK,IAAI,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC;QACxF,MAAM,IAAI,UAAU,CAAC,oCAAoC,EAAE,IAAI,CAAC,CAAA;IAClE,CAAC;IACD,OAAO,KAAK,CAAC,YAAY,EAAE,EAAE,eAAe,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAChE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,EAAE,IAAI,EAAE,IAAI,CAAC,EAAE,EAAE,IAAI,EAAE,IAAI,CAAC,EAAE,EAAE,GAAG,EAAE,GAAG,CAAC,EAAE,EAAE,GAAG,EAAE,CAAC,CAAA;AAChE,CAAC"}

package/dist/runtime.d.ts

@@ -0,0 +1,26 @@
+import type { BufferPlan } from './buffers.js';
+import type { KernelSpec } from './codegen.js';
+export interface CompiledRuntime {
+    /** Upload one or more parameter Float32Arrays to their GPU buffers. */
+    uploadParams(params: Record<string, Float32Array>): void;
+    /** Read all parameters back as Float32Arrays — used for UI panels. */
+    downloadParams(): Promise<Record<string, Float32Array>>;
+    /** Read all parameter gradients back. Mostly for verification / debugging. */
+    downloadParamGrads(): Promise<Record<string, Float32Array>>;
+    /**
+     * One full forward+backward step.
+     *   1. Uploads `inputs` (tokens, targets, masks) to input buffers.
+     *   2. Dispatches every kernel in order.
+     *   3. Reads back the loss scalar.
+     * Returns the loss as a JS number.
+     */
+    step(inputs: Record<string, Int32Array | Float32Array>): Promise<number>;
+    /** Free GPU resources. */
+    destroy(): void;
+}
+export interface RuntimeOpts {
+    /** Pre-acquired GPUDevice. If omitted, runtime requests its own. */
+    device?: GPUDevice;
+}
+export declare function createRuntime(plan: BufferPlan, kernels: KernelSpec[], lossBufferId: number, opts?: RuntimeOpts): Promise<CompiledRuntime>;
+//# sourceMappingURL=runtime.d.ts.map

package/dist/runtime.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"runtime.d.ts","sourceRoot":"","sources":["../src/runtime.ts"],"names":[],"mappings":"AAMA,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,cAAc,CAAA;AAC9C,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,cAAc,CAAA;AAM9C,MAAM,WAAW,eAAe;IAC9B,uEAAuE;IACvE,YAAY,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,GAAG,IAAI,CAAA;IACxD,sEAAsE;IACtE,cAAc,IAAI,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,CAAC,CAAA;IACvD,8EAA8E;IAC9E,kBAAkB,IAAI,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,CAAC,CAAA;IAC3D;;;;;;OAMG;IACH,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,GAAG,OAAO,CAAC,MAAM,CAAC,CAAA;IACxE,0BAA0B;IAC1B,OAAO,IAAI,IAAI,CAAA;CAChB;AAED,MAAM,WAAW,WAAW;IAC1B,oEAAoE;IACpE,MAAM,CAAC,EAAE,SAAS,CAAA;CACnB;AAQD,wBAAsB,aAAa,CACjC,IAAI,EAAE,UAAU,EAChB,OAAO,EAAE,UAAU,EAAE,EACrB,YAAY,EAAE,MAAM,EACpB,IAAI,GAAE,WAAgB,GACrB,OAAO,CAAC,eAAe,CAAC,CA8K1B"}

package/dist/runtime.js

@@ -0,0 +1,190 @@
+// WebGPU runtime. Reads a BufferPlan + KernelSpec[] (produced by codegen),
+// allocates real GPU buffers and pipelines, and provides a `step()` method
+// that uploads inputs, dispatches all kernels, and reads back outputs.
+//
+// Browser-only: this module needs `navigator.gpu` at runtime.
+// Inlined numeric values (per WebGPU spec) so this module is importable in Node
+// for codegen-only usage. The browser provides GPUBufferUsage as a global, but
+// referencing it at module scope would crash before any browser code runs.
+const STORAGE_RW = 0x80 /*STORAGE*/ | 0x8 /*COPY_DST*/ | 0x4; /*COPY_SRC*/
+const READBACK = 0x1 /*MAP_READ*/ | 0x8; /*COPY_DST*/
+export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
+    const device = opts.device ?? await acquireDevice();
+    const queue = device.queue;
+    // ---- Allocate one GPUBuffer per BufferSpec --------------------------------
+    // State buffers also get filled with their initValue at allocation time.
+    const buffers = new Map();
+    for (const spec of plan.buffers) {
+        const buf = device.createBuffer({
+            size: spec.byteSize,
+            usage: STORAGE_RW,
+            label: spec.name ?? `t${spec.id}-${spec.kind}`,
+        });
+        buffers.set(spec.id, buf);
+        if (spec.kind === 'state') {
+            // Fill with initValue (typically 0). Float and int both 4 bytes per element.
+            const elements = spec.byteSize / 4;
+            const init = spec.dtype === 'f32'
+                ? new Float32Array(elements).fill(spec.initValue ?? 0)
+                : new Int32Array(elements).fill(Math.trunc(spec.initValue ?? 0));
+            queue.writeBuffer(buf, 0, init);
+        }
+    }
+    // ---- Compile pipelines per kernel; cache by WGSL source -------------------
+    // Push an error scope around each shader+pipeline creation so we can surface
+    // the actual compile error rather than the cryptic "previous error" that
+    // comes from using an invalid pipeline at dispatch time.
+    const moduleCache = new Map();
+    const pipelines = [];
+    const probes = [];
+    for (const k of kernels) {
+        if (!k.wgsl) {
+            pipelines.push(null);
+            continue;
+        }
+        let module = moduleCache.get(k.wgsl);
+        if (!module) {
+            module = device.createShaderModule({ code: k.wgsl, label: k.opKind });
+            moduleCache.set(k.wgsl, module);
+        }
+        device.pushErrorScope('validation');
+        const pipeline = device.createComputePipeline({
+            layout: 'auto',
+            compute: { module, entryPoint: 'main' },
+            label: k.opKind,
+        });
+        pipelines.push(pipeline);
+        probes.push(device.popErrorScope().then(err => err ? { k, module: module, err } : null));
+    }
+    const probeResults = await Promise.all(probes);
+    const failures = probeResults.filter((p) => p != null);
+    if (failures.length > 0) {
+        const reports = [];
+        for (const { k, module, err } of failures) {
+            const info = await module.getCompilationInfo();
+            const messages = info.messages
+                .map(m => `  L${m.lineNum}:${m.linePos} [${m.type}] ${m.message}`)
+                .join('\n');
+            reports.push(`[shader compile error] ${k.opKind} (op #${k.opIndex}): ${err.message}\n` +
+                (messages || '  (no compilation messages)') +
+                `\n--- WGSL ---\n${k.wgsl}\n-----------`);
+        }
+        // eslint-disable-next-line no-console
+        console.error(reports.join('\n\n'));
+        throw new Error(`tensorgrad: ${failures.length} shader(s) failed to compile (see console).`);
+    }
+    // ---- Pre-build bind groups (static — buffer ids don't change per step) ---
+    const bindGroups = kernels.map((k, i) => {
+        const pipeline = pipelines[i];
+        if (!pipeline)
+            return null;
+        return device.createBindGroup({
+            layout: pipeline.getBindGroupLayout(0),
+            entries: k.bindings.map((bufId, idx) => ({
+                binding: idx,
+                resource: { buffer: buffers.get(bufId) },
+            })),
+        });
+    });
+    // ---- Loss readback staging buffer -----------------------------------------
+    const lossSpec = plan.buffers[lossBufferId];
+    const lossReadback = device.createBuffer({ size: lossSpec.byteSize, usage: READBACK });
+    // ---- step() ---------------------------------------------------------------
+    async function step(inputs) {
+        for (const [name, bufId] of plan.inputsByName) {
+            const data = inputs[name];
+            if (!data)
+                throw new Error(`tensorgrad: missing input '${name}'`);
+            const expectedBytes = plan.buffers[bufId].byteSize;
+            if (data.byteLength !== expectedBytes) {
+                throw new Error(`tensorgrad: input '${name}' has ${data.byteLength} bytes, expected ${expectedBytes}`);
+            }
+            // Cast to BufferSource: typed arrays are accepted by writeBuffer at runtime
+            // but TS may infer ArrayBufferLike (vs ArrayBuffer) under strict configs.
+            queue.writeBuffer(buffers.get(bufId), 0, data);
+        }
+        const encoder = device.createCommandEncoder({ label: 'tensorgrad-step' });
+        for (let i = 0; i < kernels.length; i++) {
+            const k = kernels[i];
+            if (!k.wgsl || k.threads === 0)
+                continue;
+            const pipeline = pipelines[i];
+            const bindGroup = bindGroups[i];
+            const pass = encoder.beginComputePass({ label: k.opKind });
+            pass.setPipeline(pipeline);
+            pass.setBindGroup(0, bindGroup);
+            // WebGPU caps each dispatch dimension at 65535 workgroups. Split into 2D
+            // when a kernel needs more than that on the X axis. Kernels compute their
+            // global index as `gid.x + gid.y * (65535 * workgroup_size)`, matching the
+            // stride we set here. For dispatches that fit in one row, gid.y is 0.
+            const wgCount = Math.max(1, Math.ceil(k.threads / k.workgroupSize));
+            const MAX_X = 65535;
+            const wgX = Math.min(wgCount, MAX_X);
+            const wgY = Math.ceil(wgCount / MAX_X);
+            pass.dispatchWorkgroups(wgX, wgY, 1);
+            pass.end();
+        }
+        // After all dispatches: writebacks (Adam state, updated params).
+        // copyBufferToBuffer is queued onto the same encoder so it's ordered after
+        // all kernel dispatches.
+        for (const wb of plan.writebacks) {
+            encoder.copyBufferToBuffer(buffers.get(wb.source), 0, buffers.get(wb.dest), 0, wb.bytes);
+        }
+        encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, lossReadback, 0, lossSpec.byteSize);
+        queue.submit([encoder.finish()]);
+        await lossReadback.mapAsync(GPUMapMode.READ);
+        const view = new Float32Array(lossReadback.getMappedRange().slice(0));
+        lossReadback.unmap();
+        return view[0];
+    }
+    // ---- uploadParams ---------------------------------------------------------
+    function uploadParams(params) {
+        for (const [name, bufId] of plan.paramsByName) {
+            const data = params[name];
+            if (!data)
+                continue;
+            queue.writeBuffer(buffers.get(bufId), 0, data);
+        }
+    }
+    // ---- download helpers -----------------------------------------------------
+    async function downloadFromMap(map) {
+        const stagings = [];
+        const encoder = device.createCommandEncoder({ label: 'tensorgrad-download' });
+        for (const [name, bufId] of map) {
+            const spec = plan.buffers[bufId];
+            const staging = device.createBuffer({ size: spec.byteSize, usage: READBACK });
+            encoder.copyBufferToBuffer(buffers.get(bufId), 0, staging, 0, spec.byteSize);
+            stagings.push({ name, buf: staging, bytes: spec.byteSize });
+        }
+        queue.submit([encoder.finish()]);
+        const out = {};
+        for (const s of stagings) {
+            await s.buf.mapAsync(GPUMapMode.READ);
+            out[s.name] = new Float32Array(s.buf.getMappedRange().slice(0));
+            s.buf.unmap();
+            s.buf.destroy();
+        }
+        return out;
+    }
+    return {
+        uploadParams,
+        downloadParams: () => downloadFromMap(plan.paramsByName),
+        downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),
+        step,
+        destroy: () => {
+            for (const b of buffers.values())
+                b.destroy();
+            lossReadback.destroy();
+        },
+    };
+}
+async function acquireDevice() {
+    if (typeof navigator === 'undefined' || !navigator.gpu) {
+        throw new Error('tensorgrad: WebGPU not available in this environment');
+    }
+    const adapter = await navigator.gpu.requestAdapter();
+    if (!adapter)
+        throw new Error('tensorgrad: no WebGPU adapter');
+    return await adapter.requestDevice();
+}
+//# sourceMappingURL=runtime.js.map

package/dist/runtime.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"runtime.js","sourceRoot":"","sources":["../src/runtime.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,2EAA2E;AAC3E,uEAAuE;AACvE,EAAE;AACF,8DAA8D;AAiC9D,gFAAgF;AAChF,+EAA+E;AAC/E,2EAA2E;AAC3E,MAAM,UAAU,GAAG,IAAI,CAAC,WAAW,GAAG,GAAG,CAAC,YAAY,GAAG,GAAG,CAAA,CAAC,YAAY;AACzE,MAAM,QAAQ,GAAG,GAAG,CAAC,YAAY,GAAG,GAAG,CAAA,CAAC,YAAY;AAEpD,MAAM,CAAC,KAAK,UAAU,aAAa,CACjC,IAAgB,EAChB,OAAqB,EACrB,YAAoB,EACpB,OAAoB,EAAE;IAEtB,MAAM,MAAM,GAAG,IAAI,CAAC,MAAM,IAAI,MAAM,aAAa,EAAE,CAAA;IACnD,MAAM,KAAK,GAAG,MAAM,CAAC,KAAK,CAAA;IAE1B,8EAA8E;IAC9E,yEAAyE;IACzE,MAAM,OAAO,GAAG,IAAI,GAAG,EAAqB,CAAA;IAC5C,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,OAAO,EAAE,CAAC;QAChC,MAAM,GAAG,GAAG,MAAM,CAAC,YAAY,CAAC;YAC9B,IAAI,EAAE,IAAI,CAAC,QAAQ;YACnB,KAAK,EAAE,UAAU;YACjB,KAAK,EAAE,IAAI,CAAC,IAAI,IAAI,IAAI,IAAI,CAAC,EAAE,IAAI,IAAI,CAAC,IAAI,EAAE;SAC/C,CAAC,CAAA;QACF,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,GAAG,CAAC,CAAA;QACzB,IAAI,IAAI,CAAC,IAAI,KAAK,OAAO,EAAE,CAAC;YAC1B,6EAA6E;YAC7E,MAAM,QAAQ,GAAG,IAAI,CAAC,QAAQ,GAAG,CAAC,CAAA;YAClC,MAAM,IAAI,GAAG,IAAI,CAAC,KAAK,KAAK,KAAK;gBAC/B,CAAC,CAAC,IAAI,YAAY,CAAC,QAAQ,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,SAAS,IAAI,CAAC,CAAC;gBACtD,CAAC,CAAC,IAAI,UAAU,CAAC,QAAQ,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,SAAS,IAAI,CAAC,CAAC,CAAC,CAAA;YAClE,KAAK,CAAC,WAAW,CAAC,GAAG,EAAE,CAAC,EAAE,IAA+B,CAAC,CAAA;QAC5D,CAAC;IACH,CAAC;IAED,8EAA8E;IAC9E,6EAA6E;IAC7E,yEAAyE;IACzE,yDAAyD;IACzD,MAAM,WAAW,GAAG,IAAI,GAAG,EAA2B,CAAA;IACtD,MAAM,SAAS,GAAkC,EAAE,CAAA;IAEnD,MAAM,MAAM,GAAiB,EAAE,CAAA;IAC/B,KAAK,MAAM,CAAC,IAAI,OAAO,EAAE,CAAC;QACxB,IAAI,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC;YAAC,SAAS,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;YAAC,SAAQ;QAAC,CAAC;QAC/C,IAAI,MAAM,GAAG,WAAW,CAAC,GAAG,CAAC,CAAC,CAAC,IAAI,CAAC,CAAA;QACpC,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,MAAM,GAAG,MAAM,CAAC,kBAAkB,CAAC,EAAE,IAAI,EAAE,CAAC,CAAC,IAAI,EAAE,KAAK,EAAE,CAAC,CAAC,MAAM,EAAE,CAAC,CAAA;YACrE,WAAW,CAAC,GAAG,CAAC,CAAC,CAAC,IAAI,EAAE,MAAM,CAAC,CAAA;QACjC,CAAC;QACD,MAAM,CAAC,cAAc,CAAC,YAAY,CAAC,CAAA;QACnC,MAAM,QAAQ,GAAG,MAAM,CAAC,qBAAqB,CAAC;YAC5C,MAAM,EAAE,MAAM;YACd,OAAO,EAAE,EAAE,MAAM,EAAE,UAAU,EAAE,MAAM,EAAE;YACvC,KAAK,EAAE,CAAC,CAAC,MAAM;SAChB,CAAC,CAAA;QACF,SAAS,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAA;QACxB,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,aAAa,EAAE,CAAC,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,EAAE,MAAO,EAAE,GAAG,EAAE,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAA;IAC3F,CAAC;IACD,MAAM,YAAY,GAAG,MAAM,OAAO,CAAC,GAAG,CAAC,MAAM,CAAC,CAAA;IAC9C,MAAM,QAAQ,GAAG,YAAY,CAAC,MAAM,CAAC,CAAC,CAAC,EAAkE,EAAE,CAAC,CAAC,IAAI,IAAI,CAAC,CAAA;IACtH,IAAI,QAAQ,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;QACxB,MAAM,OAAO,GAAa,EAAE,CAAA;QAC5B,KAAK,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,GAAG,EAAE,IAAI,QAAQ,EAAE,CAAC;YAC1C,MAAM,IAAI,GAAG,MAAM,MAAM,CAAC,kBAAkB,EAAE,CAAA;YAC9C,MAAM,QAAQ,GAAG,IAAI,CAAC,QAAQ;iBAC3B,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,MAAM,CAAC,CAAC,OAAO,IAAI,CAAC,CAAC,OAAO,KAAK,CAAC,CAAC,IAAI,KAAK,CAAC,CAAC,OAAO,EAAE,CAAC;iBACjE,IAAI,CAAC,IAAI,CAAC,CAAA;YACb,OAAO,CAAC,IAAI,CACV,0BAA0B,CAAC,CAAC,MAAM,SAAS,CAAC,CAAC,OAAO,MAAM,GAAG,CAAC,OAAO,IAAI;gBACzE,CAAC,QAAQ,IAAI,6BAA6B,CAAC;gBAC3C,mBAAmB,CAAC,CAAC,IAAI,eAAe,CACzC,CAAA;QACH,CAAC;QACD,sCAAsC;QACtC,OAAO,CAAC,KAAK,CAAC,OAAO,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,CAAA;QACnC,MAAM,IAAI,KAAK,CAAC,eAAe,QAAQ,CAAC,MAAM,6CAA6C,CAAC,CAAA;IAC9F,CAAC;IAED,6EAA6E;IAC7E,MAAM,UAAU,GAA4B,OAAO,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE;QAC/D,MAAM,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAC,CAAA;QAC7B,IAAI,CAAC,QAAQ;YAAE,OAAO,IAAI,CAAA;QAC1B,OAAO,MAAM,CAAC,eAAe,CAAC;YAC5B,MAAM,EAAE,QAAQ,CAAC,kBAAkB,CAAC,CAAC,CAAC;YACtC,OAAO,EAAE,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,EAAE,CAAC,CAAC;gBACvC,OAAO,EAAE,GAAG;gBACZ,QAAQ,EAAE,EAAE,MAAM,EAAE,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,EAAE;aAC1C,CAAC,CAAC;SACJ,CAAC,CAAA;IACJ,CAAC,CAAC,CAAA;IAEF,8EAA8E;IAC9E,MAAM,QAAQ,GAAG,IAAI,CAAC,OAAO,CAAC,YAAY,CAAE,CAAA;IAC5C,MAAM,YAAY,GAAG,MAAM,CAAC,YAAY,CAAC,EAAE,IAAI,EAAE,QAAQ,CAAC,QAAQ,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAA;IAEtF,8EAA8E;IAC9E,KAAK,UAAU,IAAI,CAAC,MAAiD;QACnE,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,MAAM,IAAI,GAAG,MAAM,CAAC,IAAI,CAAC,CAAA;YACzB,IAAI,CAAC,IAAI;gBAAE,MAAM,IAAI,KAAK,CAAC,8BAA8B,IAAI,GAAG,CAAC,CAAA;YACjE,MAAM,aAAa,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,QAAQ,CAAA;YACnD,IAAI,IAAI,CAAC,UAAU,KAAK,aAAa,EAAE,CAAC;gBACtC,MAAM,IAAI,KAAK,CAAC,sBAAsB,IAAI,SAAS,IAAI,CAAC,UAAU,oBAAoB,aAAa,EAAE,CAAC,CAAA;YACxG,CAAC;YACD,4EAA4E;YAC5E,0EAA0E;YAC1E,KAAK,CAAC,WAAW,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,EAAE,CAAC,EAAE,IAA+B,CAAC,CAAA;QAC5E,CAAC;QAED,MAAM,OAAO,GAAG,MAAM,CAAC,oBAAoB,CAAC,EAAE,KAAK,EAAE,iBAAiB,EAAE,CAAC,CAAA;QACzE,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACxC,MAAM,CAAC,GAAG,OAAO,CAAC,CAAC,CAAE,CAAA;YACrB,IAAI,CAAC,CAAC,CAAC,IAAI,IAAI,CAAC,CAAC,OAAO,KAAK,CAAC;gBAAE,SAAQ;YACxC,MAAM,QAAQ,GAAG,SAAS,CAAC,CAAC,CAAE,CAAA;YAC9B,MAAM,SAAS,GAAG,UAAU,CAAC,CAAC,CAAE,CAAA;YAChC,MAAM,IAAI,GAAG,OAAO,CAAC,gBAAgB,CAAC,EAAE,KAAK,EAAE,CAAC,CAAC,MAAM,EAAE,CAAC,CAAA;YAC1D,IAAI,CAAC,WAAW,CAAC,QAAQ,CAAC,CAAA;YAC1B,IAAI,CAAC,YAAY,CAAC,CAAC,EAAE,SAAS,CAAC,CAAA;YAC/B,yEAAyE;YACzE,0EAA0E;YAC1E,2EAA2E;YAC3E,sEAAsE;YACtE,MAAM,OAAO,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,OAAO,GAAG,CAAC,CAAC,aAAa,CAAC,CAAC,CAAA;YACnE,MAAM,KAAK,GAAG,KAAK,CAAA;YACnB,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,CAAC,OAAO,EAAE,KAAK,CAAC,CAAA;YACpC,MAAM,GAAG,GAAG,IAAI,CAAC,IAAI,CAAC,OAAO,GAAG,KAAK,CAAC,CAAA;YACtC,IAAI,CAAC,kBAAkB,CAAC,GAAG,EAAE,GAAG,EAAE,CAAC,CAAC,CAAA;YACpC,IAAI,CAAC,GAAG,EAAE,CAAA;QACZ,CAAC;QACD,iEAAiE;QACjE,2EAA2E;QAC3E,yBAAyB;QACzB,KAAK,MAAM,EAAE,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACjC,OAAO,CAAC,kBAAkB,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,MAAM,CAAE,EAAE,CAAC,EAAE,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,CAAE,EAAE,CAAC,EAAE,EAAE,CAAC,KAAK,CAAC,CAAA;QAC5F,CAAC;QACD,OAAO,CAAC,kBAAkB,CAAC,OAAO,CAAC,GAAG,CAAC,YAAY,CAAE,EAAE,CAAC,EAAE,YAAY,EAAE,CAAC,EAAE,QAAQ,CAAC,QAAQ,CAAC,CAAA;QAC7F,KAAK,CAAC,MAAM,CAAC,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAA;QAEhC,MAAM,YAAY,CAAC,QAAQ,CAAC,UAAU,CAAC,IAAI,CAAC,CAAA;QAC5C,MAAM,IAAI,GAAG,IAAI,YAAY,CAAC,YAAY,CAAC,cAAc,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAA;QACrE,YAAY,CAAC,KAAK,EAAE,CAAA;QACpB,OAAO,IAAI,CAAC,CAAC,CAAE,CAAA;IACjB,CAAC;IAED,8EAA8E;IAC9E,SAAS,YAAY,CAAC,MAAoC;QACxD,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,MAAM,IAAI,GAAG,MAAM,CAAC,IAAI,CAAC,CAAA;YACzB,IAAI,CAAC,IAAI;gBAAE,SAAQ;YACnB,KAAK,CAAC,WAAW,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,EAAE,CAAC,EAAE,IAA+B,CAAC,CAAA;QAC5E,CAAC;IACH,CAAC;IAED,8EAA8E;IAC9E,KAAK,UAAU,eAAe,CAAC,GAAwB;QACrD,MAAM,QAAQ,GAAsD,EAAE,CAAA;QACtE,MAAM,OAAO,GAAG,MAAM,CAAC,oBAAoB,CAAC,EAAE,KAAK,EAAE,qBAAqB,EAAE,CAAC,CAAA;QAC7E,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,GAAG,EAAE,CAAC;YAChC,MAAM,IAAI,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAA;YACjC,MAAM,OAAO,GAAG,MAAM,CAAC,YAAY,CAAC,EAAE,IAAI,EAAE,IAAI,CAAC,QAAQ,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAA;YAC7E,OAAO,CAAC,kBAAkB,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAA;YAC7E,QAAQ,CAAC,IAAI,CAAC,EAAE,IAAI,EAAE,GAAG,EAAE,OAAO,EAAE,KAAK,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAA;QAC7D,CAAC;QACD,KAAK,CAAC,MAAM,CAAC,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAA;QAChC,MAAM,GAAG,GAAiC,EAAE,CAAA;QAC5C,KAAK,MAAM,CAAC,IAAI,QAAQ,EAAE,CAAC;YACzB,MAAM,CAAC,CAAC,GAAG,CAAC,QAAQ,CAAC,UAAU,CAAC,IAAI,CAAC,CAAA;YACrC,GAAG,CAAC,CAAC,CAAC,IAAI,CAAC,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,GAAG,CAAC,cAAc,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAA;YAC/D,CAAC,CAAC,GAAG,CAAC,KAAK,EAAE,CAAA;YACb,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAA;QACjB,CAAC;QACD,OAAO,GAAG,CAAA;IACZ,CAAC;IAED,OAAO;QACL,YAAY;QACZ,cAAc,EAAE,GAAG,EAAE,CAAC,eAAe,CAAC,IAAI,CAAC,YAAY,CAAC;QACxD,kBAAkB,EAAE,GAAG,EAAE,CAAC,eAAe,CAAC,IAAI,CAAC,gBAAgB,CAAC;QAChE,IAAI;QACJ,OAAO,EAAE,GAAG,EAAE;YACZ,KAAK,MAAM,CAAC,IAAI,OAAO,CAAC,MAAM,EAAE;gBAAE,CAAC,CAAC,OAAO,EAAE,CAAA;YAC7C,YAAY,CAAC,OAAO,EAAE,CAAA;QACxB,CAAC;KACF,CAAA;AACH,CAAC;AAED,KAAK,UAAU,aAAa;IAC1B,IAAI,OAAO,SAAS,KAAK,WAAW,IAAI,CAAC,SAAS,CAAC,GAAG,EAAE,CAAC;QACvD,MAAM,IAAI,KAAK,CAAC,sDAAsD,CAAC,CAAA;IACzE,CAAC;IACD,MAAM,OAAO,GAAG,MAAM,SAAS,CAAC,GAAG,CAAC,cAAc,EAAE,CAAA;IACpD,IAAI,CAAC,OAAO;QAAE,MAAM,IAAI,KAAK,CAAC,+BAA+B,CAAC,CAAA;IAC9D,OAAO,MAAM,OAAO,CAAC,aAAa,EAAE,CAAA;AACtC,CAAC"}

package/dist/shape.d.ts

@@ -0,0 +1,24 @@
+import type { Shape, CallSite } from './ir.js';
+export declare class ShapeError extends Error {
+    constructor(message: string, site: CallSite | null);
+}
+export declare function shapesEqual(a: Shape, b: Shape): boolean;
+export declare function shapeSize(shape: Shape): number;
+export declare function showShape(shape: Shape): string;
+export declare function broadcastTrailing(a: Shape, b: Shape): Shape | null;
+export declare function inferElementwiseBinop(opName: string, aShape: Shape, bShape: Shape, site: CallSite | null): Shape;
+export declare function inferUnary(_opName: string, aShape: Shape, _site: CallSite | null): Shape;
+export declare function inferMeanLast(opName: string, aShape: Shape, site: CallSite | null): Shape;
+export declare function inferSumLast(opName: string, aShape: Shape, site: CallSite | null): Shape;
+export declare function inferReshape(opName: string, aShape: Shape, newShape: Shape, site: CallSite | null): Shape;
+export declare function inferTranspose(opName: string, aShape: Shape, perm: readonly number[], site: CallSite | null): Shape;
+export declare function inferMatmul(opName: string, aShape: Shape, bShape: Shape, site: CallSite | null): Shape;
+export declare function inferMatmulBatched(opName: string, aShape: Shape, bShape: Shape, site: CallSite | null): Shape;
+export declare function inferOneHot(opName: string, indicesShape: Shape, depth: number, site: CallSite | null): Shape;
+export declare function inferWhereCausal(opName: string, aShape: Shape, site: CallSite | null): Shape;
+export declare function inferSliceLastRange(opName: string, aShape: Shape, start: number, end: number, site: CallSite | null): Shape;
+export declare function inferBroadcastTo(opName: string, aShape: Shape, targetShape: Shape, site: CallSite | null): Shape;
+export declare function inferSumToShape(opName: string, aShape: Shape, targetShape: Shape, site: CallSite | null): Shape;
+export declare function inferWhere(opName: string, condShape: Shape, aShape: Shape, bShape: Shape, site: CallSite | null): Shape;
+export declare function inferReluGrad(opName: string, xShape: Shape, dyShape: Shape, site: CallSite | null): Shape;
+//# sourceMappingURL=shape.d.ts.map

package/dist/shape.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"shape.d.ts","sourceRoot":"","sources":["../src/shape.ts"],"names":[],"mappings":"AAkBA,OAAO,KAAK,EAAE,KAAK,EAAE,QAAQ,EAAE,MAAM,SAAS,CAAA;AAO9C,qBAAa,UAAW,SAAQ,KAAK;gBACvB,OAAO,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI;CAKnD;AAUD,wBAAgB,WAAW,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,KAAK,GAAG,OAAO,CAIvD;AAED,wBAAgB,SAAS,CAAC,KAAK,EAAE,KAAK,GAAG,MAAM,CAI9C;AAED,wBAAgB,SAAS,CAAC,KAAK,EAAE,KAAK,GAAG,MAAM,CAE9C;AAKD,wBAAgB,iBAAiB,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,KAAK,GAAG,KAAK,GAAG,IAAI,CAclE;AASD,wBAAgB,qBAAqB,CACnC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAClE,KAAK,CAWP;AAED,wBAAgB,UAAU,CAAC,OAAO,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAExF;AAED,wBAAgB,aAAa,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAIzF;AAED,wBAAgB,YAAY,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAIxF;AAED,wBAAgB,YAAY,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,QAAQ,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CA0BzG;AAED,wBAAgB,cAAc,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,SAAS,MAAM,EAAE,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAWnH;AAGD,wBAAgB,WAAW,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAStG;AAGD,wBAAgB,kBAAkB,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAoB7G;AAED,wBAAgB,WAAW,CAAC,MAAM,EAAE,MAAM,EAAE,YAAY,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAG5G;AAGD,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAM5F;AAED,wBAAgB,mBAAmB,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAO3H;AAID,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,WAAW,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAahH;AAID,wBAAgB,eAAe,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,WAAW,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAa/G;AAID,wBAAgB,UAAU,CAAC,MAAM,EAAE,MAAM,EAAE,SAAS,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAMvH;AAED,wBAAgB,aAAa,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,KAAK,CAKzG"}