tensorgrad 0.0.2 → 0.0.4

package/dist/module.js

@@ -6,8 +6,8 @@
 //     W: Tensor; b: Tensor
 //     constructor(inDim: number, outDim: number) {
 //       super()
-//       this.W = this.param([inDim, outDim])
-//       this.b = this.param([outDim])
+//       this.W = this.param([inDim, outDim])               // randn, scale 0.02
+//       this.b = this.param([outDim], { init: 'zeros' })
 //     }
 //   }
 //   class Block extends Module {
@@ -25,6 +25,28 @@
 // and writeback wiring. Forward functions are pure and stateless — they
 // take the materialized model and inputs, return a Tensor.
 import { paramInput } from './trace.js';
+function boxMuller() {
+    return Math.sqrt(-2 * Math.log(Math.max(1e-10, Math.random()))) * Math.cos(2 * Math.PI * Math.random());
+}
+function resolveInit(opts) {
+    const init = opts?.init ?? 'randn';
+    if (init === 'randn') {
+        const scale = opts?.scale ?? 0.02;
+        return (size) => {
+            const arr = new Float32Array(size);
+            for (let i = 0; i < size; i++)
+                arr[i] = boxMuller() * scale;
+            return arr;
+        };
+    }
+    if (init === 'zeros')
+        return (size) => new Float32Array(size);
+    if (init === 'ones')
+        return (size) => { const a = new Float32Array(size); a.fill(1); return a; };
+    if (typeof init === 'function')
+        return init;
+    throw new Error(`Unknown init: ${String(init)}`);
+}
 // ============================================================================
 // Internals: param sentinel
 // ============================================================================
@@ -36,9 +58,11 @@ import { paramInput } from './trace.js';
 class ParamSentinel {
     shape;
     dtype;
-    constructor(shape, dtype) {
+    initFn;
+    constructor(shape, dtype, initFn) {
         this.shape = shape;
         this.dtype = dtype;
+        this.initFn = initFn;
     }
 }
 // ============================================================================
@@ -51,33 +75,35 @@ export class Module {
      * that gets replaced with a real Tensor at compile time.
      *
      * The parameter's name is auto-derived from its property path in the model
-     * tree (e.g. `layers.0.attn.W_q`).
+     * tree (e.g. `layers.0.attn.W_q`). Init metadata travels with the param;
+     * call `compiled.uploadInitialParams()` to apply it after compile.
      */
-    param(shape, dtype = 'f32') {
+    param(shape, opts) {
+        const dtype = opts?.dtype ?? 'f32';
         // Lie to TypeScript: the sentinel becomes a Tensor at materialize time.
-        return new ParamSentinel(shape, dtype);
+        return new ParamSentinel(shape, dtype, resolveInit(opts));
     }
 }
-// ============================================================================
-// Tree walking
-// ============================================================================
 /**
  * Walk the module tree and replace every ParamSentinel with a real Tensor
  * created via `paramInput(autoName, ...)`. Must be called inside an active
  * trace context (paramInput appends to the current graph).
  *
- * Returns a flat record of `{ path: tensor }` for every materialized param.
+ * Returns the param tensors keyed by path, plus init functions for use by
+ * `uploadInitialParams`.
  */
 export function materializeParams(root) {
-    const out = {};
+    const tensors = {};
+    const initFns = {};
     visit(root, '', (path, val, owner, key) => {
         if (val instanceof ParamSentinel) {
             const t = paramInput(path, val.shape, val.dtype);
             owner[key] = t;
-            out[path] = t;
+            tensors[path] = t;
+            initFns[path] = val.initFn;
         }
     });
-    return out;
+    return { tensors, initFns };
 }
 function visit(node, path, visitor) {
     if (node === null || node === undefined)

package/dist/module.js.map

@@ -1 +1 @@
-{"version":3,"file":"module.js","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AAAA,6EAA6E;AAC7E,EAAE;AACF,+CAA+C;AAC/C,EAAE;AACF,kCAAkC;AAClC,2BAA2B;AAC3B,mDAAmD;AACnD,gBAAgB;AAChB,6CAA6C;AAC7C,sCAAsC;AACtC,QAAQ;AACR,MAAM;AACN,iCAAiC;AACjC,8BAA8B;AAC9B,+BAA+B;AAC/B,MAAM;AACN,iCAAiC;AACjC,mCAAmC;AACnC,+CAA+C;AAC/C,MAAM;AACN,EAAE;AACF,wEAAwE;AACxE,0EAA0E;AAC1E,0EAA0E;AAC1E,wEAAwE;AACxE,2DAA2D;AAG3D,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAEvC,+EAA+E;AAC/E,4BAA4B;AAC5B,+EAA+E;AAC/E,EAAE;AACF,6EAA6E;AAC7E,4EAA4E;AAC5E,2EAA2E;AAC3E,yEAAyE;AAEzE,MAAM,aAAa;IACW;IAA8B;IAA1D,YAA4B,KAAY,EAAkB,KAAY;QAA1C,UAAK,GAAL,KAAK,CAAO;QAAkB,UAAK,GAAL,KAAK,CAAO;IAAG,CAAC;CAC3E;AAED,+EAA+E;AAC/E,oBAAoB;AACpB,+EAA+E;AAE/E,MAAM,OAAgB,MAAM;IAC1B;;;;;;;OAOG;IACO,KAAK,CAAC,KAAY,EAAE,QAAe,KAAK;QAChD,wEAAwE;QACxE,OAAO,IAAI,aAAa,CAAC,KAAK,EAAE,KAAK,CAAsB,CAAA;IAC7D,CAAC;CACF;AAED,+EAA+E;AAC/E,eAAe;AACf,+EAA+E;AAE/E;;;;;;GAMG;AACH,MAAM,UAAU,iBAAiB,CAAC,IAAY;IAC5C,MAAM,GAAG,GAA2B,EAAE,CAAA;IACtC,KAAK,CAAC,IAAI,EAAE,EAAE,EAAE,CAAC,IAAI,EAAE,GAAG,EAAE,KAAK,EAAE,GAAG,EAAE,EAAE;QACxC,IAAI,GAAG,YAAY,aAAa,EAAE,CAAC;YACjC,MAAM,CAAC,GAAG,UAAU,CAAC,IAAI,EAAE,GAAG,CAAC,KAAK,EAAE,GAAG,CAAC,KAAK,CAAC,CAC/C;YAAC,KAAa,CAAC,GAAG,CAAC,GAAG,CAAC,CAAA;YACxB,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,CAAA;QACf,CAAC;IACH,CAAC,CAAC,CAAA;IACF,OAAO,GAAG,CAAA;AACZ,CAAC;AAaD,SAAS,KAAK,CAAC,IAAa,EAAE,IAAY,EAAE,OAAgB;IAC1D,IAAI,IAAI,KAAK,IAAI,IAAI,IAAI,KAAK,SAAS;QAAE,OAAM;IAC/C,IAAI,OAAO,IAAI,KAAK,QAAQ;QAAE,OAAM;IAEpC,IAAI,IAAI,YAAY,MAAM,EAAE,CAAC;QAC3B,KAAK,MAAM,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,IAAc,CAAC,EAAE,CAAC;YAC9C,MAAM,KAAK,GAAI,IAAY,CAAC,GAAG,CAAC,CAAA;YAChC,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,IAAI,GAAG,EAAE,CAAC,CAAC,CAAC,GAAG,CAAA;YAC/C,UAAU,CAAC,KAAK,EAAE,SAAS,EAAE,IAAI,EAAE,GAAG,EAAE,OAAO,CAAC,CAAA;QAClD,CAAC;QACD,OAAM;IACR,CAAC;IACD,IAAI,KAAK,CAAC,OAAO,CAAC,IAAI,CAAC,EAAE,CAAC;QACxB,IAAI,CAAC,OAAO,CAAC,CAAC,IAAI,EAAE,CAAC,EAAE,EAAE;YACvB,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAA;YACnD,UAAU,CAAC,IAAI,EAAE,SAAS,EAAE,IAAyB,EAAE,CAAC,EAAE,OAAO,CAAC,CAAA;QACpE,CAAC,CAAC,CAAA;QACF,OAAM;IACR,CAAC;IACD,2EAA2E;IAC3E,uBAAuB;AACzB,CAAC;AAED,SAAS,UAAU,CAAC,KAAc,EAAE,IAAY,EAAE,KAAa,EAAE,GAAoB,EAAE,OAAgB;IACrG,IAAI,KAAK,YAAY,MAAM,IAAI,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE,CAAC;QACpD,KAAK,CAAC,KAAK,EAAE,IAAI,EAAE,OAAO,CAAC,CAAA;IAC7B,CAAC;SAAM,CAAC;QACN,OAAO,CAAC,IAAI,EAAE,KAAK,EAAE,KAAK,EAAE,GAAG,CAAC,CAAA;IAClC,CAAC;AACH,CAAC"}
+{"version":3,"file":"module.js","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AAAA,6EAA6E;AAC7E,EAAE;AACF,+CAA+C;AAC/C,EAAE;AACF,kCAAkC;AAClC,2BAA2B;AAC3B,mDAAmD;AACnD,gBAAgB;AAChB,gFAAgF;AAChF,yDAAyD;AACzD,QAAQ;AACR,MAAM;AACN,iCAAiC;AACjC,8BAA8B;AAC9B,+BAA+B;AAC/B,MAAM;AACN,iCAAiC;AACjC,mCAAmC;AACnC,+CAA+C;AAC/C,MAAM;AACN,EAAE;AACF,wEAAwE;AACxE,0EAA0E;AAC1E,0EAA0E;AAC1E,wEAAwE;AACxE,2DAA2D;AAG3D,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AA8BvC,SAAS,SAAS;IAChB,OAAO,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,IAAI,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,EAAE,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC,CAAA;AACzG,CAAC;AAED,SAAS,WAAW,CAAC,IAA8B;IACjD,MAAM,IAAI,GAAG,IAAI,EAAE,IAAI,IAAI,OAAO,CAAA;IAClC,IAAI,IAAI,KAAK,OAAO,EAAE,CAAC;QACrB,MAAM,KAAK,GAAG,IAAI,EAAE,KAAK,IAAI,IAAI,CAAA;QACjC,OAAO,CAAC,IAAI,EAAE,EAAE;YACd,MAAM,GAAG,GAAG,IAAI,YAAY,CAAC,IAAI,CAAC,CAAA;YAClC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,EAAE,CAAC,EAAE;gBAAE,GAAG,CAAC,CAAC,CAAC,GAAG,SAAS,EAAE,GAAG,KAAK,CAAA;YAC3D,OAAO,GAAG,CAAA;QACZ,CAAC,CAAA;IACH,CAAC;IACD,IAAI,IAAI,KAAK,OAAO;QAAE,OAAO,CAAC,IAAI,EAAE,EAAE,CAAC,IAAI,YAAY,CAAC,IAAI,CAAC,CAAA;IAC7D,IAAI,IAAI,KAAK,MAAM;QAAE,OAAO,CAAC,IAAI,EAAE,EAAE,GAAG,MAAM,CAAC,GAAG,IAAI,YAAY,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAA,CAAC,CAAC,CAAA;IAC/F,IAAI,OAAO,IAAI,KAAK,UAAU;QAAE,OAAO,IAAI,CAAA;IAC3C,MAAM,IAAI,KAAK,CAAC,iBAAiB,MAAM,CAAC,IAAI,CAAC,EAAE,CAAC,CAAA;AAClD,CAAC;AAED,+EAA+E;AAC/E,4BAA4B;AAC5B,+EAA+E;AAC/E,EAAE;AACF,6EAA6E;AAC7E,4EAA4E;AAC5E,2EAA2E;AAC3E,yEAAyE;AAEzE,MAAM,aAAa;IAEC;IACA;IACA;IAHlB,YACkB,KAAY,EACZ,KAAY,EACZ,MAAc;QAFd,UAAK,GAAL,KAAK,CAAO;QACZ,UAAK,GAAL,KAAK,CAAO;QACZ,WAAM,GAAN,MAAM,CAAQ;IAC7B,CAAC;CACL;AAED,+EAA+E;AAC/E,oBAAoB;AACpB,+EAA+E;AAE/E,MAAM,OAAgB,MAAM;IAC1B;;;;;;;;OAQG;IACO,KAAK,CAAC,KAAY,EAAE,IAAmB;QAC/C,MAAM,KAAK,GAAG,IAAI,EAAE,KAAK,IAAI,KAAK,CAAA;QAClC,wEAAwE;QACxE,OAAO,IAAI,aAAa,CAAC,KAAK,EAAE,KAAK,EAAE,WAAW,CAAC,IAAI,CAAC,CAAsB,CAAA;IAChF,CAAC;CACF;AAaD;;;;;;;GAOG;AACH,MAAM,UAAU,iBAAiB,CAAC,IAAY;IAC5C,MAAM,OAAO,GAA2B,EAAE,CAAA;IAC1C,MAAM,OAAO,GAA2B,EAAE,CAAA;IAC1C,KAAK,CAAC,IAAI,EAAE,EAAE,EAAE,CAAC,IAAI,EAAE,GAAG,EAAE,KAAK,EAAE,GAAG,EAAE,EAAE;QACxC,IAAI,GAAG,YAAY,aAAa,EAAE,CAAC;YACjC,MAAM,CAAC,GAAG,UAAU,CAAC,IAAI,EAAE,GAAG,CAAC,KAAK,EAAE,GAAG,CAAC,KAAK,CAAC,CAC/C;YAAC,KAAa,CAAC,GAAG,CAAC,GAAG,CAAC,CAAA;YACxB,OAAO,CAAC,IAAI,CAAC,GAAG,CAAC,CAAA;YACjB,OAAO,CAAC,IAAI,CAAC,GAAG,GAAG,CAAC,MAAM,CAAA;QAC5B,CAAC;IACH,CAAC,CAAC,CAAA;IACF,OAAO,EAAE,OAAO,EAAE,OAAO,EAAE,CAAA;AAC7B,CAAC;AAaD,SAAS,KAAK,CAAC,IAAa,EAAE,IAAY,EAAE,OAAgB;IAC1D,IAAI,IAAI,KAAK,IAAI,IAAI,IAAI,KAAK,SAAS;QAAE,OAAM;IAC/C,IAAI,OAAO,IAAI,KAAK,QAAQ;QAAE,OAAM;IAEpC,IAAI,IAAI,YAAY,MAAM,EAAE,CAAC;QAC3B,KAAK,MAAM,GAAG,IAAI,MAAM,CAAC,IAAI,CAAC,IAAc,CAAC,EAAE,CAAC;YAC9C,MAAM,KAAK,GAAI,IAAY,CAAC,GAAG,CAAC,CAAA;YAChC,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,IAAI,GAAG,EAAE,CAAC,CAAC,CAAC,GAAG,CAAA;YAC/C,UAAU,CAAC,KAAK,EAAE,SAAS,EAAE,IAAI,EAAE,GAAG,EAAE,OAAO,CAAC,CAAA;QAClD,CAAC;QACD,OAAM;IACR,CAAC;IACD,IAAI,KAAK,CAAC,OAAO,CAAC,IAAI,CAAC,EAAE,CAAC;QACxB,IAAI,CAAC,OAAO,CAAC,CAAC,IAAI,EAAE,CAAC,EAAE,EAAE;YACvB,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,GAAG,IAAI,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAA;YACnD,UAAU,CAAC,IAAI,EAAE,SAAS,EAAE,IAAyB,EAAE,CAAC,EAAE,OAAO,CAAC,CAAA;QACpE,CAAC,CAAC,CAAA;QACF,OAAM;IACR,CAAC;IACD,2EAA2E;IAC3E,uBAAuB;AACzB,CAAC;AAED,SAAS,UAAU,CAAC,KAAc,EAAE,IAAY,EAAE,KAAa,EAAE,GAAoB,EAAE,OAAgB;IACrG,IAAI,KAAK,YAAY,MAAM,IAAI,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE,CAAC;QACpD,KAAK,CAAC,KAAK,EAAE,IAAI,EAAE,OAAO,CAAC,CAAA;IAC7B,CAAC;SAAM,CAAC;QACN,OAAO,CAAC,IAAI,EAAE,KAAK,EAAE,KAAK,EAAE,GAAG,CAAC,CAAA;IAClC,CAAC;AACH,CAAC"}

package/dist/nn.d.ts

@@ -0,0 +1,19 @@
+import { Module } from './module.js';
+import type { Tensor } from './ir.js';
+export declare class Linear extends Module {
+    readonly inDim: number;
+    readonly outDim: number;
+    W: Tensor;
+    b: Tensor | null;
+    constructor(inDim: number, outDim: number, withBias?: boolean);
+}
+export declare function linearFwd(p: Linear, x: Tensor): Tensor;
+export declare class LayerNorm extends Module {
+    readonly d: number;
+    readonly eps: number;
+    g: Tensor;
+    b: Tensor;
+    constructor(d: number, eps?: number);
+}
+export declare function layerNormFwd(p: LayerNorm, x: Tensor): Tensor;
+//# sourceMappingURL=nn.d.ts.map

package/dist/nn.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"nn.d.ts","sourceRoot":"","sources":["../src/nn.ts"],"names":[],"mappings":"AAeA,OAAO,EAAE,MAAM,EAAE,MAAM,aAAa,CAAA;AACpC,OAAO,KAAK,EAAE,MAAM,EAAE,MAAM,SAAS,CAAA;AAOrC,qBAAa,MAAO,SAAQ,MAAM;aAGJ,KAAK,EAAE,MAAM;aAAkB,MAAM,EAAE,MAAM;IAFzE,CAAC,EAAE,MAAM,CAAA;IACT,CAAC,EAAE,MAAM,GAAG,IAAI,CAAA;gBACY,KAAK,EAAE,MAAM,EAAkB,MAAM,EAAE,MAAM,EAAE,QAAQ,UAAO;CAK3F;AAED,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAGtD;AAMD,qBAAa,SAAU,SAAQ,MAAM;aAGP,CAAC,EAAE,MAAM;aAAkB,GAAG,EAAE,MAAM;IAFlE,CAAC,EAAE,MAAM,CAAA;IACT,CAAC,EAAE,MAAM,CAAA;gBACmB,CAAC,EAAE,MAAM,EAAkB,GAAG,GAAE,MAAa;CAK1E;AAED,wBAAgB,YAAY,CAAC,CAAC,EAAE,SAAS,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAM5D"}

package/dist/nn.js

@@ -0,0 +1,60 @@
+// Standard "batteries-included" Module subclasses for the most common layers.
+//
+// JAX-style: each class declares its params (and their init); the forward is a
+// plain function the user calls with `(module, x)`. No subclassing, no method
+// dispatch — keeps the autograd-traced computation visible at the call site.
+//
+// Import as a namespace:
+//
+//   import { nn } from 'tensorgrad'
+//   class Block extends Module {
+//     ln  = new nn.LayerNorm(D)
+//     ffn = new nn.Linear(D, 4 * D)
+//   }
+//   const y = nn.linearFwd(p.ffn, nn.layerNormFwd(p.ln, x))
+import { Module } from './module.js';
+import { add, matmul, sub, mul, div, sqrt, meanLast } from './ops.js';
+// ----------------------------------------------------------------------------
+// Linear: y = x @ W (+ b)
+// ----------------------------------------------------------------------------
+export class Linear extends Module {
+    inDim;
+    outDim;
+    W;
+    b;
+    constructor(inDim, outDim, withBias = true) {
+        super();
+        this.inDim = inDim;
+        this.outDim = outDim;
+        this.W = this.param([inDim, outDim]); // randn, scale 0.02
+        this.b = withBias ? this.param([outDim], { init: 'zeros' }) : null;
+    }
+}
+export function linearFwd(p, x) {
+    const out = matmul(x, p.W);
+    return p.b ? add(out, p.b) : out;
+}
+// ----------------------------------------------------------------------------
+// LayerNorm — normalizes over the last axis. eps defaults to 1e-5.
+// ----------------------------------------------------------------------------
+export class LayerNorm extends Module {
+    d;
+    eps;
+    g;
+    b;
+    constructor(d, eps = 1e-5) {
+        super();
+        this.d = d;
+        this.eps = eps;
+        this.g = this.param([d], { init: 'ones' });
+        this.b = this.param([d], { init: 'zeros' });
+    }
+}
+export function layerNormFwd(p, x) {
+    const m = meanLast(x);
+    const c = sub(x, m);
+    const v = meanLast(mul(c, c));
+    const stdev = sqrt(add(v, p.eps));
+    return add(mul(div(c, stdev), p.g), p.b);
+}
+//# sourceMappingURL=nn.js.map

package/dist/nn.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"nn.js","sourceRoot":"","sources":["../src/nn.ts"],"names":[],"mappings":"AAAA,8EAA8E;AAC9E,EAAE;AACF,+EAA+E;AAC/E,8EAA8E;AAC9E,6EAA6E;AAC7E,EAAE;AACF,yBAAyB;AACzB,EAAE;AACF,oCAAoC;AACpC,iCAAiC;AACjC,gCAAgC;AAChC,oCAAoC;AACpC,MAAM;AACN,4DAA4D;AAE5D,OAAO,EAAE,MAAM,EAAE,MAAM,aAAa,CAAA;AAEpC,OAAO,EAAE,GAAG,EAAE,MAAM,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI,EAAE,QAAQ,EAAE,MAAM,UAAU,CAAA;AAErE,+EAA+E;AAC/E,0BAA0B;AAC1B,+EAA+E;AAE/E,MAAM,OAAO,MAAO,SAAQ,MAAM;IAGJ;IAA+B;IAF3D,CAAC,CAAQ;IACT,CAAC,CAAe;IAChB,YAA4B,KAAa,EAAkB,MAAc,EAAE,QAAQ,GAAG,IAAI;QACxF,KAAK,EAAE,CAAA;QADmB,UAAK,GAAL,KAAK,CAAQ;QAAkB,WAAM,GAAN,MAAM,CAAQ;QAEvE,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC,CAAA,CAAsB,oBAAoB;QAC9E,IAAI,CAAC,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAC,CAAC,CAAC,IAAI,CAAA;IACpE,CAAC;CACF;AAED,MAAM,UAAU,SAAS,CAAC,CAAS,EAAE,CAAS;IAC5C,MAAM,GAAG,GAAG,MAAM,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAA;IAC1B,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,CAAA;AAClC,CAAC;AAED,+EAA+E;AAC/E,mEAAmE;AACnE,+EAA+E;AAE/E,MAAM,OAAO,SAAU,SAAQ,MAAM;IAGP;IAA2B;IAFvD,CAAC,CAAQ;IACT,CAAC,CAAQ;IACT,YAA4B,CAAS,EAAkB,MAAc,IAAI;QACvE,KAAK,EAAE,CAAA;QADmB,MAAC,GAAD,CAAC,CAAQ;QAAkB,QAAG,GAAH,GAAG,CAAe;QAEvE,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,EAAE,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,CAAA;QAC1C,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,EAAE,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAA;IAC7C,CAAC;CACF;AAED,MAAM,UAAU,YAAY,CAAC,CAAY,EAAE,CAAS;IAClD,MAAM,CAAC,GAAG,QAAQ,CAAC,CAAC,CAAC,CAAA;IACrB,MAAM,CAAC,GAAG,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAA;IACnB,MAAM,CAAC,GAAG,QAAQ,CAAC,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAA;IAC7B,MAAM,KAAK,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,CAAA;IACjC,OAAO,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,EAAE,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAA;AAC1C,CAAC"}

package/dist/runtime.d.ts

@@ -1,8 +1,41 @@
 import type { BufferPlan } from './buffers.js';
 import type { KernelSpec } from './codegen.js';
+export interface UploadParamsOptions {
+    /** Skip the "missing param" check, allowing the caller to update only some
+     *  params and leave the rest at their current GPU values. Extra (unknown)
+     *  keys are still rejected — that's always a typo. Default: false. */
+    partial?: boolean;
+}
+export interface StepOptions {
+    /** Read back tensors registered via `capture(name, t)` during the trace.
+     *  When false/unset, `step()` returns just the loss number; staging buffers
+     *  for captures are not allocated. When true, returns `{ loss, captures }`
+     *  and lazily allocates one staging buffer per capture on first use. */
+    withCaptures?: boolean;
+}
+export interface StepWithCaptures {
+    loss: number;
+    captures: Record<string, Float32Array>;
+}
+export interface RunOptions {
+    /** Read back tensors registered via `capture(name, t)` during the trace.
+     *  When false/unset, `run()` returns just the output Float32Array.
+     *  When true, returns `{ output, captures }`. */
+    withCaptures?: boolean;
+}
+export interface RunWithCaptures {
+    output: Float32Array;
+    captures: Record<string, Float32Array>;
+}
 export interface CompiledRuntime {
-    /** Upload one or more parameter Float32Arrays to their GPU buffers. */
-    uploadParams(params: Record<string, Float32Array>): void;
+    /** Map of param name -> the underlying GPUBuffer. Pass to a sibling compile
+     *  via `sharedParams` to share without copies — every step on this runtime
+     *  is immediately visible to anyone reading these buffers. */
+    params: Map<string, GPUBuffer>;
+    /** Upload parameter Float32Arrays to their GPU buffers. By default, requires
+     *  *all* params to be present; throws on any unknown or missing key. Pass
+     *  `{ partial: true }` to skip the missing-key check. */
+    uploadParams(params: Record<string, Float32Array>, opts?: UploadParamsOptions): 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. */
@@ -11,16 +44,58 @@ export interface CompiledRuntime {
      * 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.
+     *   3. Reads back the loss scalar (and any registered captures, if requested).
+     * Default returns the loss as a JS number; with `{ withCaptures: true }`
+     * returns `{ loss, captures }` where `captures` is keyed by the names passed
+     * to `capture(...)` during the trace.
      */
     step(inputs: Record<string, Int32Array | Float32Array>): Promise<number>;
+    step(inputs: Record<string, Int32Array | Float32Array>, opts: {
+        withCaptures: true;
+    }): Promise<StepWithCaptures>;
+    step(inputs: Record<string, Int32Array | Float32Array>, opts: StepOptions): Promise<number | StepWithCaptures>;
+    /** Like `step()` but returns the full output Float32Array instead of just
+     *  its first element. For training graphs this is rarely useful (the output
+     *  *is* a scalar loss); it's the primary API for forward-only compiles. */
+    run(inputs: Record<string, Int32Array | Float32Array>): Promise<Float32Array>;
+    run(inputs: Record<string, Int32Array | Float32Array>, opts: {
+        withCaptures: true;
+    }): Promise<RunWithCaptures>;
+    run(inputs: Record<string, Int32Array | Float32Array>, opts: RunOptions): Promise<Float32Array | RunWithCaptures>;
+    /** Re-zero all optimizer state buffers (Adam's m/v) in place. Pair with
+     *  `uploadInitialParams()` for a full training reset without recompile. */
+    resetOptimizerState(): void;
     /** Free GPU resources. */
     destroy(): void;
 }
+/** Forward-only compiled runtime — produced by `compileForward`. No optimizer,
+ *  no backward. Returns the output tensor (not just a scalar) per `run()` call. */
+export interface CompiledForward {
+    params: Map<string, GPUBuffer>;
+    uploadParams(params: Record<string, Float32Array>, opts?: UploadParamsOptions): void;
+    downloadParams(): Promise<Record<string, Float32Array>>;
+    /** Forward-only dispatch. Returns the graph's output tensor as a Float32Array
+     *  (the user's returned tensor from the forward function, in row-major order).
+     *  With `{ withCaptures: true }`, returns `{ output, captures }`. */
+    run(inputs: Record<string, Int32Array | Float32Array>): Promise<Float32Array>;
+    run(inputs: Record<string, Int32Array | Float32Array>, opts: {
+        withCaptures: true;
+    }): Promise<RunWithCaptures>;
+    run(inputs: Record<string, Int32Array | Float32Array>, opts: RunOptions): Promise<Float32Array | RunWithCaptures>;
+    destroy(): void;
+}
 export interface RuntimeOpts {
     /** Pre-acquired GPUDevice. If omitted, runtime requests its own. */
     device?: GPUDevice;
+    /** External param buffers to bind in place of allocating fresh ones, keyed
+     *  by param name. Used to share params between a training compile and a
+     *  sibling forward-only compile (e.g., a B=1 inference graph). When a name
+     *  is in this map, the runtime reuses the provided GPUBuffer; otherwise it
+     *  allocates as usual. */
+    sharedParams?: Map<string, GPUBuffer>;
 }
 export declare function createRuntime(plan: BufferPlan, kernels: KernelSpec[], lossBufferId: number, opts?: RuntimeOpts): Promise<CompiledRuntime>;
+/** Same machinery as `createRuntime`, narrower public API: no step,
+ *  no resetOptimizerState, no downloadParamGrads. Used by `compileForward`. */
+export declare function createForwardRuntime(plan: BufferPlan, kernels: KernelSpec[], outputBufferId: number, opts?: RuntimeOpts): Promise<CompiledForward>;
 //# sourceMappingURL=runtime.d.ts.map

package/dist/runtime.d.ts.map

@@ -1 +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"}
+{"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,mBAAmB;IAClC;;0EAEsE;IACtE,OAAO,CAAC,EAAE,OAAO,CAAA;CAClB;AAED,MAAM,WAAW,WAAW;IAC1B;;;4EAGwE;IACxE,YAAY,CAAC,EAAE,OAAO,CAAA;CACvB;AAED,MAAM,WAAW,gBAAgB;IAC/B,IAAI,EAAE,MAAM,CAAA;IACZ,QAAQ,EAAE,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,CAAA;CACvC;AAED,MAAM,WAAW,UAAU;IACzB;;qDAEiD;IACjD,YAAY,CAAC,EAAE,OAAO,CAAA;CACvB;AAED,MAAM,WAAW,eAAe;IAC9B,MAAM,EAAE,YAAY,CAAA;IACpB,QAAQ,EAAE,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,CAAA;CACvC;AAED,MAAM,WAAW,eAAe;IAC9B;;kEAE8D;IAC9D,MAAM,EAAE,GAAG,CAAC,MAAM,EAAE,SAAS,CAAC,CAAA;IAC9B;;6DAEyD;IACzD,YAAY,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,EAAE,IAAI,CAAC,EAAE,mBAAmB,GAAG,IAAI,CAAA;IACpF,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;;;;;;;;OAQG;IACH,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,GAAG,OAAO,CAAC,MAAM,CAAC,CAAA;IACxE,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE;QAAE,YAAY,EAAE,IAAI,CAAA;KAAE,GAAG,OAAO,CAAC,gBAAgB,CAAC,CAAA;IAChH,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE,WAAW,GAAG,OAAO,CAAC,MAAM,GAAG,gBAAgB,CAAC,CAAA;IAC9G;;+EAE2E;IAC3E,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,GAAG,OAAO,CAAC,YAAY,CAAC,CAAA;IAC7E,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE;QAAE,YAAY,EAAE,IAAI,CAAA;KAAE,GAAG,OAAO,CAAC,eAAe,CAAC,CAAA;IAC9G,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE,UAAU,GAAG,OAAO,CAAC,YAAY,GAAG,eAAe,CAAC,CAAA;IACjH;+EAC2E;IAC3E,mBAAmB,IAAI,IAAI,CAAA;IAC3B,0BAA0B;IAC1B,OAAO,IAAI,IAAI,CAAA;CAChB;AAED;mFACmF;AACnF,MAAM,WAAW,eAAe;IAC9B,MAAM,EAAE,GAAG,CAAC,MAAM,EAAE,SAAS,CAAC,CAAA;IAC9B,YAAY,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,EAAE,IAAI,CAAC,EAAE,mBAAmB,GAAG,IAAI,CAAA;IACpF,cAAc,IAAI,OAAO,CAAC,MAAM,CAAC,MAAM,EAAE,YAAY,CAAC,CAAC,CAAA;IACvD;;yEAEqE;IACrE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,GAAG,OAAO,CAAC,YAAY,CAAC,CAAA;IAC7E,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE;QAAE,YAAY,EAAE,IAAI,CAAA;KAAE,GAAG,OAAO,CAAC,eAAe,CAAC,CAAA;IAC9G,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,GAAG,YAAY,CAAC,EAAE,IAAI,EAAE,UAAU,GAAG,OAAO,CAAC,YAAY,GAAG,eAAe,CAAC,CAAA;IACjH,OAAO,IAAI,IAAI,CAAA;CAChB;AAED,MAAM,WAAW,WAAW;IAC1B,oEAAoE;IACpE,MAAM,CAAC,EAAE,SAAS,CAAA;IAClB;;;;8BAI0B;IAC1B,YAAY,CAAC,EAAE,GAAG,CAAC,MAAM,EAAE,SAAS,CAAC,CAAA;CACtC;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,CAkU1B;AAED;+EAC+E;AAC/E,wBAAsB,oBAAoB,CACxC,IAAI,EAAE,UAAU,EAChB,OAAO,EAAE,UAAU,EAAE,EACrB,cAAc,EAAE,MAAM,EACtB,IAAI,GAAE,WAAgB,GACrB,OAAO,CAAC,eAAe,CAAC,CAS1B"}

package/dist/runtime.js

@@ -13,8 +13,21 @@ export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
     const queue = device.queue;
     // ---- Allocate one GPUBuffer per BufferSpec --------------------------------
     // State buffers also get filled with their initValue at allocation time.
+    // Param buffers may be supplied externally via opts.sharedParams; in that
+    // case we reuse the provided GPUBuffer instead of allocating, and the
+    // sibling compile that owns it is responsible for upload + lifetime.
     const buffers = new Map();
+    const sharedParams = opts.sharedParams;
     for (const spec of plan.buffers) {
+        if (spec.kind === 'param' && sharedParams?.has(spec.name)) {
+            const shared = sharedParams.get(spec.name);
+            if (shared.size !== spec.byteSize) {
+                throw new Error(`sharedParams: size mismatch for '${spec.name}' — supplied ${shared.size} bytes, ` +
+                    `compiled graph expects ${spec.byteSize}.`);
+            }
+            buffers.set(spec.id, shared);
+            continue;
+        }
         const buf = device.createBuffer({
             size: spec.byteSize,
             usage: STORAGE_RW,
@@ -30,6 +43,13 @@ export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
             queue.writeBuffer(buf, 0, init);
         }
     }
+    // Track which params are externally owned — those are skipped on destroy().
+    const ownedBufferIds = new Set();
+    for (const spec of plan.buffers) {
+        if (!(spec.kind === 'param' && sharedParams?.has(spec.name))) {
+            ownedBufferIds.add(spec.id);
+        }
+    }
     // ---- 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
@@ -86,11 +106,39 @@ export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
             })),
         });
     });
-    // ---- Loss readback staging buffer -----------------------------------------
-    const lossSpec = plan.buffers[lossBufferId];
-    const lossReadback = device.createBuffer({ size: lossSpec.byteSize, usage: READBACK });
-    // ---- step() ---------------------------------------------------------------
-    async function step(inputs) {
+    // ---- Output readback staging buffer ---------------------------------------
+    // `outputBufferId` is the graph's main output (loss for training, the user's
+    // returned tensor for forward-only). step() reads back its first element;
+    // run() reads back the full Float32Array.
+    const outputSpec = plan.buffers[lossBufferId];
+    const outputReadback = device.createBuffer({ size: outputSpec.byteSize, usage: READBACK });
+    // ---- Capture readback staging buffers (lazy) ------------------------------
+    // Allocated on first `step({ withCaptures: true })` call and reused across
+    // subsequent calls. When the graph has no captures registered or when the
+    // caller never opts in, no extra GPU memory is allocated.
+    let captureStagings = null;
+    function ensureCaptureStagings() {
+        if (captureStagings)
+            return captureStagings;
+        captureStagings = new Map();
+        for (const [name, bufId] of plan.capturesByName) {
+            const spec = plan.buffers[bufId];
+            const staging = device.createBuffer({ size: spec.byteSize, usage: READBACK, label: `cap-${name}` });
+            captureStagings.set(name, staging);
+        }
+        return captureStagings;
+    }
+    // ---- dispatch() — shared core for step() and run() -----------------------
+    // Uploads inputs, dispatches all kernels (in order), queues writebacks, copies
+    // the output buffer into its staging, optionally copies captures into theirs,
+    // submits, and reads back. Returns the full output Float32Array; step() takes
+    // [0] for scalar loss, run() returns it whole.
+    async function dispatch(inputs, wantCaptures) {
+        if (wantCaptures && plan.capturesByName.size === 0) {
+            throw new Error(`withCaptures=true but no capture(...) calls were registered during ` +
+                `the trace. Add capture('name', tensor) inside your forward pass for ` +
+                `the intermediates you want read back.`);
+        }
         for (const [name, bufId] of plan.inputsByName) {
             const data = inputs[name];
             if (!data)
@@ -124,25 +172,73 @@ export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
             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.
+        // After all dispatches: writebacks (Adam state, updated params). Empty for
+        // forward-only compiles.
         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);
+        encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, outputReadback, 0, outputSpec.byteSize);
+        // Capture readbacks (only when opted in). Queued before submit so they
+        // observe the same kernel outputs as the main output.
+        let stagings = null;
+        if (wantCaptures) {
+            stagings = ensureCaptureStagings();
+            for (const [name, bufId] of plan.capturesByName) {
+                const spec = plan.buffers[bufId];
+                encoder.copyBufferToBuffer(buffers.get(bufId), 0, stagings.get(name), 0, spec.byteSize);
+            }
+        }
         queue.submit([encoder.finish()]);
-        await lossReadback.mapAsync(GPUMapMode.READ);
-        const view = new Float32Array(lossReadback.getMappedRange().slice(0));
-        lossReadback.unmap();
-        return view[0];
+        await outputReadback.mapAsync(GPUMapMode.READ);
+        const output = new Float32Array(outputReadback.getMappedRange().slice(0));
+        outputReadback.unmap();
+        if (!wantCaptures)
+            return { output, captures: null };
+        const captures = {};
+        for (const [name, staging] of stagings) {
+            await staging.mapAsync(GPUMapMode.READ);
+            captures[name] = new Float32Array(staging.getMappedRange().slice(0));
+            staging.unmap();
+        }
+        return { output, captures };
+    }
+    async function step(inputs, opts) {
+        const r = await dispatch(inputs, opts?.withCaptures === true);
+        if (opts?.withCaptures)
+            return { loss: r.output[0], captures: r.captures };
+        return r.output[0];
+    }
+    async function run(inputs, opts) {
+        const r = await dispatch(inputs, opts?.withCaptures === true);
+        if (opts?.withCaptures)
+            return { output: r.output, captures: r.captures };
+        return r.output;
     }
     // ---- uploadParams ---------------------------------------------------------
-    function uploadParams(params) {
+    function uploadParams(params, opts) {
+        const partial = opts?.partial ?? false;
+        for (const name of Object.keys(params)) {
+            if (!plan.paramsByName.has(name)) {
+                throw new Error(`uploadParams: unknown param '${name}'. ` +
+                    `Known: ${[...plan.paramsByName.keys()].sort().join(', ')}`);
+            }
+        }
+        if (!partial) {
+            for (const name of plan.paramsByName.keys()) {
+                if (!(name in params)) {
+                    throw new Error(`uploadParams: missing param '${name}'. ` +
+                        `Pass { partial: true } if you mean to update only some params.`);
+                }
+            }
+        }
         for (const [name, bufId] of plan.paramsByName) {
             const data = params[name];
             if (!data)
                 continue;
+            const expected = plan.buffers[bufId].byteSize / 4;
+            if (data.length !== expected) {
+                throw new Error(`uploadParams: '${name}' has ${data.length} elements, expected ${expected}`);
+            }
             queue.writeBuffer(buffers.get(bufId), 0, data);
         }
     }
@@ -166,16 +262,54 @@ export async function createRuntime(plan, kernels, lossBufferId, opts = {}) {
         }
         return out;
     }
+    function resetOptimizerState() {
+        for (const spec of plan.buffers) {
+            if (spec.kind !== 'state')
+                continue;
+            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(buffers.get(spec.id), 0, init);
+        }
+    }
+    // Build the params map AFTER buffer allocation so it points at the actual
+    // GPUBuffers (shared or freshly allocated).
+    const params = new Map();
+    for (const [name, bufId] of plan.paramsByName) {
+        params.set(name, buffers.get(bufId));
+    }
+    const destroy = () => {
+        for (const [id, b] of buffers) {
+            if (ownedBufferIds.has(id))
+                b.destroy();
+        }
+        outputReadback.destroy();
+        if (captureStagings)
+            for (const b of captureStagings.values())
+                b.destroy();
+    };
     return {
+        params,
         uploadParams,
         downloadParams: () => downloadFromMap(plan.paramsByName),
         downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),
         step,
-        destroy: () => {
-            for (const b of buffers.values())
-                b.destroy();
-            lossReadback.destroy();
-        },
+        run,
+        resetOptimizerState,
+        destroy,
+    };
+}
+/** Same machinery as `createRuntime`, narrower public API: no step,
+ *  no resetOptimizerState, no downloadParamGrads. Used by `compileForward`. */
+export async function createForwardRuntime(plan, kernels, outputBufferId, opts = {}) {
+    const full = await createRuntime(plan, kernels, outputBufferId, opts);
+    return {
+        params: full.params,
+        uploadParams: full.uploadParams,
+        downloadParams: full.downloadParams,
+        run: full.run,
+        destroy: full.destroy,
     };
 }
 async function acquireDevice() {

package/dist/runtime.js.map

@@ -1 +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"}
+{"version":3,"file":"runtime.js","sourceRoot":"","sources":["../src/runtime.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,2EAA2E;AAC3E,uEAAuE;AACvE,EAAE;AACF,8DAA8D;AAyG9D,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,0EAA0E;IAC1E,sEAAsE;IACtE,qEAAqE;IACrE,MAAM,OAAO,GAAG,IAAI,GAAG,EAAqB,CAAA;IAC5C,MAAM,YAAY,GAAG,IAAI,CAAC,YAAY,CAAA;IACtC,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,OAAO,EAAE,CAAC;QAChC,IAAI,IAAI,CAAC,IAAI,KAAK,OAAO,IAAI,YAAY,EAAE,GAAG,CAAC,IAAI,CAAC,IAAK,CAAC,EAAE,CAAC;YAC3D,MAAM,MAAM,GAAG,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,IAAK,CAAE,CAAA;YAC5C,IAAI,MAAM,CAAC,IAAI,KAAK,IAAI,CAAC,QAAQ,EAAE,CAAC;gBAClC,MAAM,IAAI,KAAK,CACb,oCAAoC,IAAI,CAAC,IAAI,gBAAgB,MAAM,CAAC,IAAI,UAAU;oBAClF,0BAA0B,IAAI,CAAC,QAAQ,GAAG,CAC3C,CAAA;YACH,CAAC;YACD,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,EAAE,MAAM,CAAC,CAAA;YAC5B,SAAQ;QACV,CAAC;QACD,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;IACD,4EAA4E;IAC5E,MAAM,cAAc,GAAG,IAAI,GAAG,EAAU,CAAA;IACxC,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,OAAO,EAAE,CAAC;QAChC,IAAI,CAAC,CAAC,IAAI,CAAC,IAAI,KAAK,OAAO,IAAI,YAAY,EAAE,GAAG,CAAC,IAAI,CAAC,IAAK,CAAC,CAAC,EAAE,CAAC;YAC9D,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC,CAAA;QAC7B,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,6EAA6E;IAC7E,0EAA0E;IAC1E,0CAA0C;IAC1C,MAAM,UAAU,GAAG,IAAI,CAAC,OAAO,CAAC,YAAY,CAAE,CAAA;IAC9C,MAAM,cAAc,GAAG,MAAM,CAAC,YAAY,CAAC,EAAE,IAAI,EAAE,UAAU,CAAC,QAAQ,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAA;IAE1F,8EAA8E;IAC9E,2EAA2E;IAC3E,0EAA0E;IAC1E,0DAA0D;IAC1D,IAAI,eAAe,GAAkC,IAAI,CAAA;IACzD,SAAS,qBAAqB;QAC5B,IAAI,eAAe;YAAE,OAAO,eAAe,CAAA;QAC3C,eAAe,GAAG,IAAI,GAAG,EAAE,CAAA;QAC3B,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,cAAc,EAAE,CAAC;YAChD,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,KAAK,EAAE,OAAO,IAAI,EAAE,EAAE,CAAC,CAAA;YACnG,eAAe,CAAC,GAAG,CAAC,IAAI,EAAE,OAAO,CAAC,CAAA;QACpC,CAAC;QACD,OAAO,eAAe,CAAA;IACxB,CAAC;IAED,6EAA6E;IAC7E,+EAA+E;IAC/E,8EAA8E;IAC9E,8EAA8E;IAC9E,+CAA+C;IAC/C,KAAK,UAAU,QAAQ,CACrB,MAAiD,EACjD,YAAqB;QAErB,IAAI,YAAY,IAAI,IAAI,CAAC,cAAc,CAAC,IAAI,KAAK,CAAC,EAAE,CAAC;YACnD,MAAM,IAAI,KAAK,CACb,qEAAqE;gBACrE,sEAAsE;gBACtE,uCAAuC,CACxC,CAAA;QACH,CAAC;QACD,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,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,cAAc,EAAE,CAAC,EAAE,UAAU,CAAC,QAAQ,CAAC,CAAA;QACjG,uEAAuE;QACvE,sDAAsD;QACtD,IAAI,QAAQ,GAAkC,IAAI,CAAA;QAClD,IAAI,YAAY,EAAE,CAAC;YACjB,QAAQ,GAAG,qBAAqB,EAAE,CAAA;YAClC,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,cAAc,EAAE,CAAC;gBAChD,MAAM,IAAI,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAA;gBACjC,OAAO,CAAC,kBAAkB,CAAC,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,EAAE,CAAC,EAAE,QAAQ,CAAC,GAAG,CAAC,IAAI,CAAE,EAAE,CAAC,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAA;YAC3F,CAAC;QACH,CAAC;QACD,KAAK,CAAC,MAAM,CAAC,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAA;QAEhC,MAAM,cAAc,CAAC,QAAQ,CAAC,UAAU,CAAC,IAAI,CAAC,CAAA;QAC9C,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,cAAc,CAAC,cAAc,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAA;QACzE,cAAc,CAAC,KAAK,EAAE,CAAA;QAEtB,IAAI,CAAC,YAAY;YAAE,OAAO,EAAE,MAAM,EAAE,QAAQ,EAAE,IAAI,EAAE,CAAA;QAEpD,MAAM,QAAQ,GAAiC,EAAE,CAAA;QACjD,KAAK,MAAM,CAAC,IAAI,EAAE,OAAO,CAAC,IAAI,QAAS,EAAE,CAAC;YACxC,MAAM,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAC,IAAI,CAAC,CAAA;YACvC,QAAQ,CAAC,IAAI,CAAC,GAAG,IAAI,YAAY,CAAC,OAAO,CAAC,cAAc,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAA;YACpE,OAAO,CAAC,KAAK,EAAE,CAAA;QACjB,CAAC;QACD,OAAO,EAAE,MAAM,EAAE,QAAQ,EAAE,CAAA;IAC7B,CAAC;IAMD,KAAK,UAAU,IAAI,CACjB,MAAiD,EACjD,IAAkB;QAElB,MAAM,CAAC,GAAG,MAAM,QAAQ,CAAC,MAAM,EAAE,IAAI,EAAE,YAAY,KAAK,IAAI,CAAC,CAAA;QAC7D,IAAI,IAAI,EAAE,YAAY;YAAE,OAAO,EAAE,IAAI,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,EAAE,QAAQ,EAAE,CAAC,CAAC,QAAS,EAAE,CAAA;QAC5E,OAAO,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,CAAA;IACrB,CAAC;IAMD,KAAK,UAAU,GAAG,CAChB,MAAiD,EACjD,IAAiB;QAEjB,MAAM,CAAC,GAAG,MAAM,QAAQ,CAAC,MAAM,EAAE,IAAI,EAAE,YAAY,KAAK,IAAI,CAAC,CAAA;QAC7D,IAAI,IAAI,EAAE,YAAY;YAAE,OAAO,EAAE,MAAM,EAAE,CAAC,CAAC,MAAM,EAAE,QAAQ,EAAE,CAAC,CAAC,QAAS,EAAE,CAAA;QAC1E,OAAO,CAAC,CAAC,MAAM,CAAA;IACjB,CAAC;IAED,8EAA8E;IAC9E,SAAS,YAAY,CAAC,MAAoC,EAAE,IAA0B;QACpF,MAAM,OAAO,GAAG,IAAI,EAAE,OAAO,IAAI,KAAK,CAAA;QACtC,KAAK,MAAM,IAAI,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAC;YACvC,IAAI,CAAC,IAAI,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC;gBACjC,MAAM,IAAI,KAAK,CACb,gCAAgC,IAAI,KAAK;oBACzC,UAAU,CAAC,GAAG,IAAI,CAAC,YAAY,CAAC,IAAI,EAAE,CAAC,CAAC,IAAI,EAAE,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,CAC5D,CAAA;YACH,CAAC;QACH,CAAC;QACD,IAAI,CAAC,OAAO,EAAE,CAAC;YACb,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,YAAY,CAAC,IAAI,EAAE,EAAE,CAAC;gBAC5C,IAAI,CAAC,CAAC,IAAI,IAAI,MAAM,CAAC,EAAE,CAAC;oBACtB,MAAM,IAAI,KAAK,CACb,gCAAgC,IAAI,KAAK;wBACzC,gEAAgE,CACjE,CAAA;gBACH,CAAC;YACH,CAAC;QACH,CAAC;QACD,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,MAAM,QAAQ,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,QAAQ,GAAG,CAAC,CAAA;YAClD,IAAI,IAAI,CAAC,MAAM,KAAK,QAAQ,EAAE,CAAC;gBAC7B,MAAM,IAAI,KAAK,CAAC,kBAAkB,IAAI,SAAS,IAAI,CAAC,MAAM,uBAAuB,QAAQ,EAAE,CAAC,CAAA;YAC9F,CAAC;YACD,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,SAAS,mBAAmB;QAC1B,KAAK,MAAM,IAAI,IAAI,IAAI,CAAC,OAAO,EAAE,CAAC;YAChC,IAAI,IAAI,CAAC,IAAI,KAAK,OAAO;gBAAE,SAAQ;YACnC,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,OAAO,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAE,EAAE,CAAC,EAAE,IAA+B,CAAC,CAAA;QAC9E,CAAC;IACH,CAAC;IAED,0EAA0E;IAC1E,4CAA4C;IAC5C,MAAM,MAAM,GAAG,IAAI,GAAG,EAAqB,CAAA;IAC3C,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;QAC9C,MAAM,CAAC,GAAG,CAAC,IAAI,EAAE,OAAO,CAAC,GAAG,CAAC,KAAK,CAAE,CAAC,CAAA;IACvC,CAAC;IAED,MAAM,OAAO,GAAG,GAAG,EAAE;QACnB,KAAK,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,IAAI,OAAO,EAAE,CAAC;YAC9B,IAAI,cAAc,CAAC,GAAG,CAAC,EAAE,CAAC;gBAAE,CAAC,CAAC,OAAO,EAAE,CAAA;QACzC,CAAC;QACD,cAAc,CAAC,OAAO,EAAE,CAAA;QACxB,IAAI,eAAe;YAAE,KAAK,MAAM,CAAC,IAAI,eAAe,CAAC,MAAM,EAAE;gBAAE,CAAC,CAAC,OAAO,EAAE,CAAA;IAC5E,CAAC,CAAA;IAED,OAAO;QACL,MAAM;QACN,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,GAAG;QACH,mBAAmB;QACnB,OAAO;KACR,CAAA;AACH,CAAC;AAED;+EAC+E;AAC/E,MAAM,CAAC,KAAK,UAAU,oBAAoB,CACxC,IAAgB,EAChB,OAAqB,EACrB,cAAsB,EACtB,OAAoB,EAAE;IAEtB,MAAM,IAAI,GAAG,MAAM,aAAa,CAAC,IAAI,EAAE,OAAO,EAAE,cAAc,EAAE,IAAI,CAAC,CAAA;IACrE,OAAO;QACL,MAAM,EAAE,IAAI,CAAC,MAAM;QACnB,YAAY,EAAE,IAAI,CAAC,YAAY;QAC/B,cAAc,EAAE,IAAI,CAAC,cAAc;QACnC,GAAG,EAAE,IAAI,CAAC,GAAG;QACb,OAAO,EAAE,IAAI,CAAC,OAAO;KACtB,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/trace.d.ts

@@ -1,5 +1,6 @@
 import type { Graph, Tensor, Shape, Dtype } from './ir.js';
 export declare function currentGraph(): Graph;
+export declare function isCaptureEnabled(): boolean;
 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;

package/dist/trace.d.ts.map

@@ -1 +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"}
+{"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;AAU1D,wBAAgB,YAAY,IAAI,KAAK,CAQpC;AAED,wBAAgB,gBAAgB,IAAI,OAAO,CAE1C;AAID,wBAAgB,KAAK,CAAC,EAAE,EAAE,MAAM,MAAM,GAAG,MAAM,EAAE,GAAG,KAAK,CAkBxD;AAQD,wBAAgB,SAAS,CAAC,CAAC,EAAE,CAAC,EAAE,KAAK,EAAE,EAAE,EAAE,MAAM,CAAC,GAAG,CAAC,CAWrD;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"}