tensorgrad 0.0.1 → 0.0.4

package/dist/compile.js

@@ -12,7 +12,7 @@ import { appendGrad } from './grad.js';
 import { appendAdam } from './adam.js';
 import { planBuffers } from './buffers.js';
 import { emitKernels } from './codegen.js';
-import { createRuntime } from './runtime.js';
+import { createRuntime, createForwardRuntime } from './runtime.js';
 import { Module, materializeParams } from './module.js';
 /** Trace + autograd + buffer-plan + codegen, without touching WebGPU. */
 export function compileToIR(traceFn) {
@@ -30,9 +30,13 @@ export async function compile(traceFn, opts = {}) {
     return Object.assign(runtime, { ir });
 }
 /**
- * Compile a Module-based model. The forward function takes the materialized
- * model and returns the loss tensor (typically by also calling tensorInput
- * for tokens/targets/masks inside).
+ * Compile a Module-based model. Pass a *factory* `() => new Model()`, not the
+ * model instance itself: compilation mutates the tree (every `ParamSentinel`
+ * field becomes a real `Tensor`), so the instance is consumed and shouldn't be
+ * referenced afterwards. Re-call the factory if you need a fresh tree.
+ *
+ * The forward function takes the materialized model and returns the loss
+ * tensor.
  *
  * Walks the module tree to materialize params with auto-derived names, then
  * runs trace → grad → adam → buffer plan → codegen → runtime.
@@ -41,36 +45,117 @@ export async function compile(traceFn, opts = {}) {
  * internal step count and injects the bias-corrected `lrt` scalar each call;
  * users don't need to provide it themselves.
  */
-export async function compileModule(model, forward, opts = {}) {
+export async function compileModule(modelFactory, forward, opts = {}) {
     const inputDecls = opts.inputs ?? [];
-    let paramTensors = {};
+    const model = modelFactory();
+    let materialized = { tensors: {}, initFns: {} };
     const graph = trace(() => {
-        paramTensors = materializeParams(model);
+        materialized = materializeParams(model);
         const inputTensors = inputDecls.map(d => tensorInput(d.name, d.shape, d.dtype ?? 'f32'));
         return forward(model, ...inputTensors);
     });
     const { paramGrads, loss } = appendGrad(graph);
     let adamResult;
     if (opts.adam) {
-        adamResult = appendAdam(graph, paramGrads, paramTensors, opts.adam);
+        adamResult = appendAdam(graph, paramGrads, materialized.tensors, opts.adam);
     }
     const plan = planBuffers(graph, paramGrads, adamResult?.writebacks ?? []);
     const kernels = emitKernels(graph, plan);
     const lossBufferId = plan.tensorToBuffer.get(loss.id);
     const runtime = await createRuntime(plan, kernels, lossBufferId, opts);
     // If Adam is enabled, wrap step() to track the step count and supply lrt.
+    // Wrap resetOptimizerState() too, so a reset zeros m/v *and* the bias-correction
+    // counter — otherwise the next step would skip Adam's warmup phase.
     if (adamResult) {
         const { lrtInputName, config } = adamResult;
         let t = 0;
         const lrtBuf = new Float32Array(1);
         const innerStep = runtime.step.bind(runtime);
-        runtime.step = async (inputs) => {
+        const innerReset = runtime.resetOptimizerState.bind(runtime);
+        const wrappedStep = (inputs, opts) => {
             t++;
             lrtBuf[0] = config.lr * Math.sqrt(1 - Math.pow(config.b2, t)) / (1 - Math.pow(config.b1, t));
-            return innerStep({ ...inputs, [lrtInputName]: lrtBuf });
+            const merged = { ...inputs, [lrtInputName]: lrtBuf };
+            return opts?.withCaptures ? innerStep(merged, { withCaptures: true }) : innerStep(merged);
+        };
+        runtime.step = wrappedStep;
+        runtime.resetOptimizerState = () => {
+            t = 0;
+            innerReset();
         };
     }
+    const { initFns } = materialized;
+    const uploadInitialParams = () => {
+        const out = {};
+        for (const [name, bufId] of plan.paramsByName) {
+            const shape = plan.buffers[bufId].shape;
+            const size = shape.reduce((a, b) => a * b, 1);
+            const initFn = initFns[name];
+            if (!initFn)
+                throw new Error(`uploadInitialParams: no init for param '${name}'`);
+            out[name] = initFn(size, shape);
+        }
+        runtime.uploadParams(out);
+    };
     const ir = { graph, paramGrads, loss, plan, kernels };
-    return Object.assign(runtime, { ir });
+    return Object.assign(runtime, { ir, uploadInitialParams });
+}
+// ============================================================================
+// Forward-only compile
+// ============================================================================
+/**
+ * Compile a Module-based model in forward-only mode (no autograd, no Adam).
+ * The forward function returns the output tensor (e.g., logits) instead of a
+ * scalar loss; runtime exposes `run(inputs)` returning the full output as a
+ * `Float32Array`.
+ *
+ * **Sharing params with a training compile.** Pass `opts.sharedParams =
+ * trainCompiled.params` to bind this graph's param buffers to an existing
+ * training runtime's GPU buffers — every train step is then immediately
+ * visible to `run()` calls here, no copies. The forward graph's
+ * `uploadInitialParams()` skips any param covered by `sharedParams`.
+ *
+ * Typical use: a B=1 inference graph alongside a B=512 training graph,
+ * built from the same `Module` factory.
+ */
+export async function compileForward(modelFactory, forward, opts = {}) {
+    const inputDecls = opts.inputs ?? [];
+    const model = modelFactory();
+    let materialized = { tensors: {}, initFns: {} };
+    const graph = trace(() => {
+        materialized = materializeParams(model);
+        const inputTensors = inputDecls.map(d => tensorInput(d.name, d.shape, d.dtype ?? 'f32'));
+        return forward(model, ...inputTensors);
+    });
+    const plan = planBuffers(graph, /* paramGrads */ {});
+    const kernels = emitKernels(graph, plan);
+    const outputTensor = graph.tensors[graph.outputs[0]];
+    const outputBufferId = plan.tensorToBuffer.get(outputTensor.id);
+    const runtime = await createForwardRuntime(plan, kernels, outputBufferId, opts);
+    const sharedParams = opts.sharedParams;
+    const { initFns } = materialized;
+    const uploadInitialParams = () => {
+        const out = {};
+        let needsUpload = false;
+        for (const [name, bufId] of plan.paramsByName) {
+            // Skip params covered by sharedParams — those are owned by the providing
+            // compile and already initialized there.
+            if (sharedParams?.has(name))
+                continue;
+            const shape = plan.buffers[bufId].shape;
+            const size = shape.reduce((a, b) => a * b, 1);
+            const initFn = initFns[name];
+            if (!initFn)
+                throw new Error(`uploadInitialParams: no init for param '${name}'`);
+            out[name] = initFn(size, shape);
+            needsUpload = true;
+        }
+        if (needsUpload)
+            runtime.uploadParams(out, { partial: !!sharedParams });
+    };
+    // CompiledIR.loss is the field name; for forward-only, it carries the user's
+    // returned tensor (e.g., logits). Same shape conceptually; just no autograd.
+    const ir = { graph, paramGrads: {}, loss: outputTensor, plan, kernels };
+    return Object.assign(runtime, { ir, uploadInitialParams });
 }
 //# sourceMappingURL=compile.js.map

package/dist/compile.js.map

@@ -1 +1 @@
-{"version":3,"file":"compile.js","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,EAAE;AACF,oBAAoB;AACpB,sEAAsE;AACtE,iEAAiE;AACjE,0EAA0E;AAC1E,0EAA0E;AAC1E,2EAA2E;AAC3E,mEAAmE;AAGnE,OAAO,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,YAAY,CAAA;AAC/C,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAA0C,MAAM,cAAc,CAAA;AACpF,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,MAAM,aAAa,CAAA;AAmBvD,yEAAyE;AACzE,MAAM,UAAU,WAAW,CAAC,OAAqB;IAC/C,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,CAAA;IAC5B,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAC9C,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,CAAC,CAAA;IAC3C,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;AACnD,CAAC;AAED,0EAA0E;AAC1E,MAAM,CAAC,KAAK,UAAU,OAAO,CAAC,OAAqB,EAAE,OAAoB,EAAE;IACzE,MAAM,EAAE,GAAG,WAAW,CAAC,OAAO,CAAC,CAAA;IAC/B,MAAM,YAAY,GAAG,EAAE,CAAC,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IAC5D,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,EAAE,CAAC,IAAI,EAAE,EAAE,CAAC,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAC5E,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AACvC,CAAC;AAgBD;;;;;;;;;;;GAWG;AACH,MAAM,CAAC,KAAK,UAAU,aAAa,CACjC,KAAQ,EACR,OAA8C,EAC9C,OAA6B,EAAE;IAE/B,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,IAAI,EAAE,CAAA;IACpC,IAAI,YAAY,GAA2B,EAAE,CAAA;IAC7C,MAAM,KAAK,GAAG,KAAK,CAAC,GAAG,EAAE;QACvB,YAAY,GAAG,iBAAiB,CAAC,KAAK,CAAC,CAAA;QACvC,MAAM,YAAY,GAAG,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,KAAK,CAAC,CAAC,CAAA;QACxF,OAAO,OAAO,CAAC,KAAK,EAAE,GAAG,YAAY,CAAC,CAAA;IACxC,CAAC,CAAC,CAAA;IAEF,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAE9C,IAAI,UAAqD,CAAA;IACzD,IAAI,IAAI,CAAC,IAAI,EAAE,CAAC;QACd,UAAU,GAAG,UAAU,CAAC,KAAK,EAAE,UAAU,EAAE,YAAY,EAAE,IAAI,CAAC,IAAI,CAAC,CAAA;IACrE,CAAC;IAED,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,EAAE,UAAU,EAAE,UAAU,IAAI,EAAE,CAAC,CAAA;IACzE,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,MAAM,YAAY,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IACtD,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,IAAI,EAAE,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAEtE,0EAA0E;IAC1E,IAAI,UAAU,EAAE,CAAC;QACf,MAAM,EAAE,YAAY,EAAE,MAAM,EAAE,GAAG,UAAU,CAAA;QAC3C,IAAI,CAAC,GAAG,CAAC,CAAA;QACT,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAA;QAClC,MAAM,SAAS,GAAG,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,CAAA;QAC5C,OAAO,CAAC,IAAI,GAAG,KAAK,EAAE,MAAM,EAAE,EAAE;YAC9B,CAAC,EAAE,CAAA;YACH,MAAM,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,EAAE,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,CAAA;YAC5F,OAAO,SAAS,CAAC,EAAE,GAAG,MAAM,EAAE,CAAC,YAAY,CAAC,EAAE,MAAM,EAAE,CAAC,CAAA;QACzD,CAAC,CAAA;IACH,CAAC;IAED,MAAM,EAAE,GAAe,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;IACjE,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AACvC,CAAC"}
+{"version":3,"file":"compile.js","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,EAAE;AACF,oBAAoB;AACpB,sEAAsE;AACtE,iEAAiE;AACjE,0EAA0E;AAC1E,0EAA0E;AAC1E,2EAA2E;AAC3E,mEAAmE;AAGnE,OAAO,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,YAAY,CAAA;AAC/C,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAgE,MAAM,cAAc,CAAA;AAChI,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,MAAM,aAAa,CAAA;AAmBvD,yEAAyE;AACzE,MAAM,UAAU,WAAW,CAAC,OAAqB;IAC/C,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,CAAA;IAC5B,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAC9C,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,CAAC,CAAA;IAC3C,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;AACnD,CAAC;AAED,0EAA0E;AAC1E,MAAM,CAAC,KAAK,UAAU,OAAO,CAAC,OAAqB,EAAE,OAAoB,EAAE;IACzE,MAAM,EAAE,GAAG,WAAW,CAAC,OAAO,CAAC,CAAA;IAC/B,MAAM,YAAY,GAAG,EAAE,CAAC,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IAC5D,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,EAAE,CAAC,IAAI,EAAE,EAAE,CAAC,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAC5E,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AACvC,CAAC;AAqBD;;;;;;;;;;;;;;;GAeG;AACH,MAAM,CAAC,KAAK,UAAU,aAAa,CACjC,YAAqB,EACrB,OAA8C,EAC9C,OAA6B,EAAE;IAE/B,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,IAAI,EAAE,CAAA;IACpC,MAAM,KAAK,GAAG,YAAY,EAAE,CAAA;IAC5B,IAAI,YAAY,GAAyC,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,CAAA;IACrF,MAAM,KAAK,GAAG,KAAK,CAAC,GAAG,EAAE;QACvB,YAAY,GAAG,iBAAiB,CAAC,KAAK,CAAC,CAAA;QACvC,MAAM,YAAY,GAAG,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,KAAK,CAAC,CAAC,CAAA;QACxF,OAAO,OAAO,CAAC,KAAK,EAAE,GAAG,YAAY,CAAC,CAAA;IACxC,CAAC,CAAC,CAAA;IAEF,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAE9C,IAAI,UAAqD,CAAA;IACzD,IAAI,IAAI,CAAC,IAAI,EAAE,CAAC;QACd,UAAU,GAAG,UAAU,CAAC,KAAK,EAAE,UAAU,EAAE,YAAY,CAAC,OAAO,EAAE,IAAI,CAAC,IAAI,CAAC,CAAA;IAC7E,CAAC;IAED,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,EAAE,UAAU,EAAE,UAAU,IAAI,EAAE,CAAC,CAAA;IACzE,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,MAAM,YAAY,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IACtD,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,IAAI,EAAE,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAEtE,0EAA0E;IAC1E,iFAAiF;IACjF,oEAAoE;IACpE,IAAI,UAAU,EAAE,CAAC;QACf,MAAM,EAAE,YAAY,EAAE,MAAM,EAAE,GAAG,UAAU,CAAA;QAC3C,IAAI,CAAC,GAAG,CAAC,CAAA;QACT,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAA;QAClC,MAAM,SAAS,GAAG,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAA4B,CAAA;QACvE,MAAM,UAAU,GAAG,OAAO,CAAC,mBAAmB,CAAC,IAAI,CAAC,OAAO,CAAC,CAAA;QAC5D,MAAM,WAAW,GAAG,CAClB,MAAiD,EACjD,IAAiC,EAC2C,EAAE;YAC9E,CAAC,EAAE,CAAA;YACH,MAAM,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,EAAE,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,CAAA;YAC5F,MAAM,MAAM,GAAG,EAAE,GAAG,MAAM,EAAE,CAAC,YAAY,CAAC,EAAE,MAAM,EAAE,CAAA;YACpD,OAAO,IAAI,EAAE,YAAY,CAAC,CAAC,CAAC,SAAS,CAAC,MAAM,EAAE,EAAE,YAAY,EAAE,IAAI,EAAE,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,MAAM,CAAC,CAAA;QAC3F,CAAC,CAAA;QACD,OAAO,CAAC,IAAI,GAAG,WAAsC,CAAA;QACrD,OAAO,CAAC,mBAAmB,GAAG,GAAG,EAAE;YACjC,CAAC,GAAG,CAAC,CAAA;YACL,UAAU,EAAE,CAAA;QACd,CAAC,CAAA;IACH,CAAC;IAED,MAAM,EAAE,OAAO,EAAE,GAAG,YAAY,CAAA;IAChC,MAAM,mBAAmB,GAAG,GAAG,EAAE;QAC/B,MAAM,GAAG,GAAiC,EAAE,CAAA;QAC5C,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,KAAK,CAAA;YACxC,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAA;YAC7C,MAAM,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAA;YAC5B,IAAI,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,2CAA2C,IAAI,GAAG,CAAC,CAAA;YAChF,GAAG,CAAC,IAAI,CAAC,GAAG,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;QACjC,CAAC;QACD,OAAO,CAAC,YAAY,CAAC,GAAG,CAAC,CAAA;IAC3B,CAAC,CAAA;IAED,MAAM,EAAE,GAAe,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;IACjE,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,mBAAmB,EAAE,CAAC,CAAA;AAC5D,CAAC;AAED,+EAA+E;AAC/E,uBAAuB;AACvB,+EAA+E;AAE/E;;;;;;;;;;;;;;GAcG;AACH,MAAM,CAAC,KAAK,UAAU,cAAc,CAClC,YAAqB,EACrB,OAA8C,EAC9C,OAA8B,EAAE;IAEhC,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,IAAI,EAAE,CAAA;IACpC,MAAM,KAAK,GAAG,YAAY,EAAE,CAAA;IAC5B,IAAI,YAAY,GAAyC,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,CAAA;IACrF,MAAM,KAAK,GAAG,KAAK,CAAC,GAAG,EAAE;QACvB,YAAY,GAAG,iBAAiB,CAAC,KAAK,CAAC,CAAA;QACvC,MAAM,YAAY,GAAG,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,KAAK,CAAC,CAAC,CAAA;QACxF,OAAO,OAAO,CAAC,KAAK,EAAE,GAAG,YAAY,CAAC,CAAA;IACxC,CAAC,CAAC,CAAA;IAEF,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,gBAAgB,CAAC,EAAE,CAAC,CAAA;IACpD,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,MAAM,YAAY,GAAG,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC,CAAE,CAAE,CAAA;IACtD,MAAM,cAAc,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,YAAY,CAAC,EAAE,CAAE,CAAA;IAChE,MAAM,OAAO,GAAG,MAAM,oBAAoB,CAAC,IAAI,EAAE,OAAO,EAAE,cAAc,EAAE,IAAI,CAAC,CAAA;IAE/E,MAAM,YAAY,GAAG,IAAI,CAAC,YAAY,CAAA;IACtC,MAAM,EAAE,OAAO,EAAE,GAAG,YAAY,CAAA;IAChC,MAAM,mBAAmB,GAAG,GAAG,EAAE;QAC/B,MAAM,GAAG,GAAiC,EAAE,CAAA;QAC5C,IAAI,WAAW,GAAG,KAAK,CAAA;QACvB,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,yEAAyE;YACzE,yCAAyC;YACzC,IAAI,YAAY,EAAE,GAAG,CAAC,IAAI,CAAC;gBAAE,SAAQ;YACrC,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,KAAK,CAAA;YACxC,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAA;YAC7C,MAAM,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAA;YAC5B,IAAI,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,2CAA2C,IAAI,GAAG,CAAC,CAAA;YAChF,GAAG,CAAC,IAAI,CAAC,GAAG,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;YAC/B,WAAW,GAAG,IAAI,CAAA;QACpB,CAAC;QACD,IAAI,WAAW;YAAE,OAAO,CAAC,YAAY,CAAC,GAAG,EAAE,EAAE,OAAO,EAAE,CAAC,CAAC,YAAY,EAAE,CAAC,CAAA;IACzE,CAAC,CAAA;IAED,6EAA6E;IAC7E,6EAA6E;IAC7E,MAAM,EAAE,GAAe,EAAE,KAAK,EAAE,UAAU,EAAE,EAAE,EAAE,IAAI,EAAE,YAAY,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;IACnF,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,mBAAmB,EAAE,CAAC,CAAA;AAC5D,CAAC"}

package/dist/index.d.ts

@@ -1,12 +1,14 @@
 export type { Tensor, Shape, Dtype, OpNode, Graph, CallSite } from './ir.js';
 export { ShapeError } from './shape.js';
 export { trace, traceInto, paramInput, tensorInput, stateInput } from './trace.js';
+export { capture } from './capture.js';
 export { add, sub, mul, div, sqrt, rsqrt, log, exp, relu, less, greater, where, meanLast, sumLast, reshape, transpose, matmul, matmulBatched, oneHot, arange, softmaxCausalLast, logSoftmaxLast, whereCausal, sliceLastRange, } from './ops.js';
 export { appendGrad, type GradResult } from './grad.js';
 export { appendAdam, type AdamConfig, type AdamResult } from './adam.js';
 export { planBuffers, type BufferPlan, type BufferSpec, type Writeback, type WritebackDecl } from './buffers.js';
 export { emitKernels, type KernelSpec } from './codegen.js';
-export { createRuntime, type CompiledRuntime, type RuntimeOpts } from './runtime.js';
-export { compile, compileToIR, compileModule, type CompiledIR, type CompileModuleOptions, type InputDecl } from './compile.js';
-export { Module, materializeParams } from './module.js';
+export { createRuntime, createForwardRuntime, type CompiledRuntime, type CompiledForward, type RuntimeOpts, type StepOptions, type StepWithCaptures, type RunOptions, type RunWithCaptures } from './runtime.js';
+export { compile, compileToIR, compileModule, compileForward, type CompiledIR, type CompileModuleOptions, type CompileForwardOptions, type InputDecl } from './compile.js';
+export { Module, materializeParams, type InitSpec, type ParamOptions, type MaterializedParams } from './module.js';
+export * as nn from './nn.js';
 //# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAKA,YAAY,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,QAAQ,EAAE,MAAM,SAAS,CAAA;AAC5E,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAEL,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG,EAElB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI,EAE3B,IAAI,EAAE,OAAO,EAAE,KAAK,EAEpB,QAAQ,EAAE,OAAO,EAEjB,OAAO,EAAE,SAAS,EAElB,MAAM,EAAE,aAAa,EAErB,MAAM,EAAE,MAAM,EAEd,iBAAiB,EAAE,cAAc,EAAE,WAAW,EAE9C,cAAc,GACf,MAAM,UAAU,CAAA;AAMjB,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,SAAS,EAAE,KAAK,aAAa,EAAE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,MAAM,cAAc,CAAA;AACpF,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,KAAK,UAAU,EAAE,KAAK,oBAAoB,EAAE,KAAK,SAAS,EAAE,MAAM,cAAc,CAAA;AAC9H,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,MAAM,aAAa,CAAA"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAKA,YAAY,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,QAAQ,EAAE,MAAM,SAAS,CAAA;AAC5E,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAAE,OAAO,EAAE,MAAM,cAAc,CAAA;AACtC,OAAO,EAEL,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG,EAElB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI,EAE3B,IAAI,EAAE,OAAO,EAAE,KAAK,EAEpB,QAAQ,EAAE,OAAO,EAEjB,OAAO,EAAE,SAAS,EAElB,MAAM,EAAE,aAAa,EAErB,MAAM,EAAE,MAAM,EAEd,iBAAiB,EAAE,cAAc,EAAE,WAAW,EAE9C,cAAc,GACf,MAAM,UAAU,CAAA;AAMjB,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,SAAS,EAAE,KAAK,aAAa,EAAE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAE,KAAK,eAAe,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,KAAK,WAAW,EAAE,KAAK,gBAAgB,EAAE,KAAK,UAAU,EAAE,KAAK,eAAe,EAAE,MAAM,cAAc,CAAA;AAChN,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,cAAc,EAAE,KAAK,UAAU,EAAE,KAAK,oBAAoB,EAAE,KAAK,qBAAqB,EAAE,KAAK,SAAS,EAAE,MAAM,cAAc,CAAA;AAC1K,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,KAAK,QAAQ,EAAE,KAAK,YAAY,EAAE,KAAK,kBAAkB,EAAE,MAAM,aAAa,CAAA;AAClH,OAAO,KAAK,EAAE,MAAM,SAAS,CAAA"}

package/dist/index.js

@@ -4,6 +4,7 @@
 // codegen / compile() (Phase 3+) come later.
 export { ShapeError } from './shape.js';
 export { trace, traceInto, paramInput, tensorInput, stateInput } from './trace.js';
+export { capture } from './capture.js';
 export { 
 // Element-wise arithmetic. The binops accept Tensor or JS-number for the second arg.
 add, sub, mul, div, 
@@ -31,7 +32,8 @@ export { appendGrad } from './grad.js';
 export { appendAdam } from './adam.js';
 export { planBuffers } from './buffers.js';
 export { emitKernels } from './codegen.js';
-export { createRuntime } from './runtime.js';
-export { compile, compileToIR, compileModule } from './compile.js';
+export { createRuntime, createForwardRuntime } from './runtime.js';
+export { compile, compileToIR, compileModule, compileForward } from './compile.js';
 export { Module, materializeParams } from './module.js';
+export * as nn from './nn.js';
 //# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -1 +1 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,+CAA+C;AAC/C,EAAE;AACF,8EAA8E;AAC9E,6CAA6C;AAG7C,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO;AACL,qFAAqF;AACrF,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG;AAClB,qBAAqB;AACrB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI;AAC3B,uBAAuB;AACvB,IAAI,EAAE,OAAO,EAAE,KAAK;AACpB,yEAAyE;AACzE,QAAQ,EAAE,OAAO;AACjB,YAAY;AACZ,OAAO,EAAE,SAAS;AAClB,iBAAiB;AACjB,MAAM,EAAE,aAAa;AACrB,qBAAqB;AACrB,MAAM,EAAE,MAAM;AACd,4CAA4C;AAC5C,iBAAiB,EAAE,cAAc,EAAE,WAAW;AAC9C,UAAU;AACV,cAAc,GACf,MAAM,UAAU,CAAA;AAEjB,sFAAsF;AACtF,8EAA8E;AAC9E,2EAA2E;AAC3E,qDAAqD;AACrD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAoC,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAwE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAA0C,MAAM,cAAc,CAAA;AACpF,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAA8D,MAAM,cAAc,CAAA;AAC9H,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,MAAM,aAAa,CAAA"}
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,+CAA+C;AAC/C,EAAE;AACF,8EAA8E;AAC9E,6CAA6C;AAG7C,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAAE,OAAO,EAAE,MAAM,cAAc,CAAA;AACtC,OAAO;AACL,qFAAqF;AACrF,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG;AAClB,qBAAqB;AACrB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI;AAC3B,uBAAuB;AACvB,IAAI,EAAE,OAAO,EAAE,KAAK;AACpB,yEAAyE;AACzE,QAAQ,EAAE,OAAO;AACjB,YAAY;AACZ,OAAO,EAAE,SAAS;AAClB,iBAAiB;AACjB,MAAM,EAAE,aAAa;AACrB,qBAAqB;AACrB,MAAM,EAAE,MAAM;AACd,4CAA4C;AAC5C,iBAAiB,EAAE,cAAc,EAAE,WAAW;AAC9C,UAAU;AACV,cAAc,GACf,MAAM,UAAU,CAAA;AAEjB,sFAAsF;AACtF,8EAA8E;AAC9E,2EAA2E;AAC3E,qDAAqD;AACrD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAoC,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAwE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAgJ,MAAM,cAAc,CAAA;AAChN,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,cAAc,EAA0F,MAAM,cAAc,CAAA;AAC1K,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAA6D,MAAM,aAAa,CAAA;AAClH,OAAO,KAAK,EAAE,MAAM,SAAS,CAAA"}

package/dist/ir.d.ts

@@ -162,6 +162,7 @@ export type OpNode = {
     vNew: number;
     lrt: number;
     eps: number;
+    decayShrink: number;
 } | {
     kind: 'slice_last_range';
     out: number;
@@ -193,6 +194,7 @@ export interface Graph {
     readonly ops: OpNode[];
     readonly tensors: Tensor[];
     readonly outputs: number[];
+    readonly captures: Map<string, number>;
 }
 export declare function makeGraph(): Graph;
 export declare function addTensor(g: Graph, shape: Shape, dtype: Dtype, source: number | null, site: CallSite | null): Tensor;

package/dist/ir.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"ir.d.ts","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAcA,MAAM,MAAM,KAAK,GAAG,KAAK,GAAG,KAAK,GAAG,MAAM,CAAA;AAC1C,MAAM,MAAM,KAAK,GAAG,SAAS,MAAM,EAAE,CAAA;AAIrC,MAAM,WAAW,MAAM;IACrB,QAAQ,CAAC,EAAE,EAAE,MAAM,CAAA;IACnB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IACrB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IAErB,QAAQ,CAAC,MAAM,EAAE,MAAM,GAAG,IAAI,CAAA;IAG9B,QAAQ,CAAC,IAAI,EAAE,QAAQ,GAAG,IAAI,CAAA;CAC/B;AAED,MAAM,WAAW,QAAQ;IACvB,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAA;IAEvB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAA;CACvB;AAQD,MAAM,MAAM,MAAM,GAGd;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAElD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAInD;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAGrE;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAC9D;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAG9D;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACxC;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACzC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGxC;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAC7C;IAAE,IAAI,EAAE,UAAU,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG5C;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,QAAQ,EAAE,KAAK,CAAA;CAAE,GAC5D;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,SAAS,MAAM,EAAE,CAAA;CAAE,GAMtE;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAErD;IAAE,IAAI,EAAE,gBAAgB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG7D;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,OAAO,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAC9E;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAGxD;IAAE,IAAI,EAAE,qBAAqB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvD;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAIpD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAKnE;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACnD;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGtD;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAMlE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GACxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GAKxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAA;CAAE,GAMvG;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAA;CAAE,GAGhF;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAGpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAEpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAElE;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,CAAA;AAI7D,MAAM,WAAW,KAAK;IACpB,QAAQ,CAAC,GAAG,EAAE,MAAM,EAAE,CAAA;IACtB,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAG1B,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;CAC3B;AAED,wBAAgB,SAAS,IAAI,KAAK,CAEjC;AAGD,wBAAgB,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,MAAM,CAKpH;AAMD,wBAAgB,KAAK,CAAC,CAAC,SAAS,MAAM,CAAC,MAAM,CAAC,EAC5C,CAAC,EAAE,KAAK,EACR,IAAI,EAAE,CAAC,EACP,KAAK,EAAE,KAAK,EACZ,KAAK,EAAE,KAAK,EACZ,IAAI,EAAE,QAAQ,GAAG,IAAI,EACrB,MAAM,EAAE,IAAI,CAAC,OAAO,CAAC,MAAM,EAAE;IAAE,IAAI,EAAE,CAAC,CAAA;CAAE,CAAC,EAAE,MAAM,GAAG,KAAK,CAAC,GACzD,MAAM,CAMR;AAID,wBAAgB,WAAW,CAAC,MAAM,EAAE,MAAM,GAAG,QAAQ,CAIpD;AAID,wBAAgB,UAAU,CAAC,IAAI,EAAE,QAAQ,GAAG,MAAM,CAYjD"}
+{"version":3,"file":"ir.d.ts","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAcA,MAAM,MAAM,KAAK,GAAG,KAAK,GAAG,KAAK,GAAG,MAAM,CAAA;AAC1C,MAAM,MAAM,KAAK,GAAG,SAAS,MAAM,EAAE,CAAA;AAIrC,MAAM,WAAW,MAAM;IACrB,QAAQ,CAAC,EAAE,EAAE,MAAM,CAAA;IACnB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IACrB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IAErB,QAAQ,CAAC,MAAM,EAAE,MAAM,GAAG,IAAI,CAAA;IAG9B,QAAQ,CAAC,IAAI,EAAE,QAAQ,GAAG,IAAI,CAAA;CAC/B;AAED,MAAM,WAAW,QAAQ;IACvB,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAA;IAEvB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAA;CACvB;AAQD,MAAM,MAAM,MAAM,GAGd;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAElD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAInD;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAGrE;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAC9D;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAG9D;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACxC;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACzC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGxC;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAC7C;IAAE,IAAI,EAAE,UAAU,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG5C;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,QAAQ,EAAE,KAAK,CAAA;CAAE,GAC5D;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,SAAS,MAAM,EAAE,CAAA;CAAE,GAMtE;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAErD;IAAE,IAAI,EAAE,gBAAgB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG7D;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,OAAO,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAC9E;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAGxD;IAAE,IAAI,EAAE,qBAAqB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvD;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAIpD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAKnE;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACnD;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGtD;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAMlE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GACxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GAOxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,MAAM,CAAA;CAAE,GAM5H;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAA;CAAE,GAGhF;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAGpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAEpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAElE;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,CAAA;AAI7D,MAAM,WAAW,KAAK;IACpB,QAAQ,CAAC,GAAG,EAAE,MAAM,EAAE,CAAA;IACtB,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAG1B,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAI1B,QAAQ,CAAC,QAAQ,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CACvC;AAED,wBAAgB,SAAS,IAAI,KAAK,CAEjC;AAGD,wBAAgB,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,MAAM,CAKpH;AAMD,wBAAgB,KAAK,CAAC,CAAC,SAAS,MAAM,CAAC,MAAM,CAAC,EAC5C,CAAC,EAAE,KAAK,EACR,IAAI,EAAE,CAAC,EACP,KAAK,EAAE,KAAK,EACZ,KAAK,EAAE,KAAK,EACZ,IAAI,EAAE,QAAQ,GAAG,IAAI,EACrB,MAAM,EAAE,IAAI,CAAC,OAAO,CAAC,MAAM,EAAE;IAAE,IAAI,EAAE,CAAC,CAAA;CAAE,CAAC,EAAE,MAAM,GAAG,KAAK,CAAC,GACzD,MAAM,CAMR;AAID,wBAAgB,WAAW,CAAC,MAAM,EAAE,MAAM,GAAG,QAAQ,CAIpD;AAID,wBAAgB,UAAU,CAAC,IAAI,EAAE,QAAQ,GAAG,MAAM,CAYjD"}

package/dist/ir.js

@@ -12,7 +12,7 @@
 // Design intent: keep this file boring. No tracing logic, no shape inference,
 // no codegen — those live in their own modules and consume `Graph` / `OpNode`.
 export function makeGraph() {
-    return { ops: [], tensors: [], outputs: [] };
+    return { ops: [], tensors: [], outputs: [], captures: new Map() };
 }
 // Internal: register a fresh tensor in the graph and return its id.
 export function addTensor(g, shape, dtype, source, site) {

package/dist/ir.js.map

@@ -1 +1 @@
-{"version":3,"file":"ir.js","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAAA,uDAAuD;AACvD,EAAE;AACF,gFAAgF;AAChF,+EAA+E;AAC/E,8EAA8E;AAC9E,EAAE;AACF,0DAA0D;AAC1D,uCAAuC;AACvC,8EAA8E;AAC9E,wFAAwF;AACxF,EAAE;AACF,8EAA8E;AAC9E,+EAA+E;AAmI/E,MAAM,UAAU,SAAS;IACvB,OAAO,EAAE,GAAG,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,CAAA;AAC9C,CAAC;AAED,oEAAoE;AACpE,MAAM,UAAU,SAAS,CAAC,CAAQ,EAAE,KAAY,EAAE,KAAY,EAAE,MAAqB,EAAE,IAAqB;IAC1G,MAAM,EAAE,GAAG,CAAC,CAAC,OAAO,CAAC,MAAM,CAAA;IAC3B,MAAM,CAAC,GAAW,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,CAAA;IACpD,CAAC,CAAC,OAAO,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;IACjB,OAAO,CAAC,CAAA;AACV,CAAC;AAED,kFAAkF;AAClF,0EAA0E;AAC1E,+EAA+E;AAC/E,kFAAkF;AAClF,MAAM,UAAU,KAAK,CACnB,CAAQ,EACR,IAAO,EACP,KAAY,EACZ,KAAY,EACZ,IAAqB,EACrB,MAA0D;IAE1D,MAAM,OAAO,GAAG,CAAC,CAAC,GAAG,CAAC,MAAM,CAAA;IAC5B,MAAM,GAAG,GAAG,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,OAAO,EAAE,IAAI,CAAC,CAAA;IACrD,MAAM,IAAI,GAAG,EAAE,IAAI,EAAE,GAAG,EAAE,GAAG,CAAC,EAAE,EAAE,GAAG,MAAM,EAAkC,CAAA;IAC7E,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;IAChB,OAAO,GAAG,CAAA;AACZ,CAAC;AAED,0EAA0E;AAC1E,iFAAiF;AACjF,MAAM,UAAU,WAAW,CAAC,MAAc;IACxC,+EAA+E;IAC/E,MAAM,KAAK,GAAG,CAAC,IAAI,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,EAAE,CAAA;IACvC,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,CAAA;AAC1B,CAAC;AAED,8EAA8E;AAC9E,2DAA2D;AAC3D,MAAM,UAAU,UAAU,CAAC,IAAc;IACvC,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,CAAC,CAAA;IACpC,2EAA2E;IAC3E,iEAAiE;IACjE,MAAM,UAAU,GAAa,EAAE,CAAA;IAC/B,KAAK,MAAM,IAAI,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC;QAClC,IAAI,IAAI,CAAC,QAAQ,CAAC,kBAAkB,CAAC,IAAI,IAAI,CAAC,QAAQ,CAAC,qBAAqB,CAAC;YAAE,SAAQ;QACvF,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,EAAE,CAAC,CAAA;QAC5B,IAAI,UAAU,CAAC,MAAM,IAAI,CAAC;YAAE,MAAK;IACnC,CAAC;IACD,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO,IAAI,IAAI,CAAC,MAAM,yBAAyB,CAAA;IAC5E,OAAO,IAAI,IAAI,CAAC,MAAM,QAAQ,UAAU,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAA;AACzD,CAAC"}
+{"version":3,"file":"ir.js","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAAA,uDAAuD;AACvD,EAAE;AACF,gFAAgF;AAChF,+EAA+E;AAC/E,8EAA8E;AAC9E,EAAE;AACF,0DAA0D;AAC1D,uCAAuC;AACvC,8EAA8E;AAC9E,wFAAwF;AACxF,EAAE;AACF,8EAA8E;AAC9E,+EAA+E;AAyI/E,MAAM,UAAU,SAAS;IACvB,OAAO,EAAE,GAAG,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,QAAQ,EAAE,IAAI,GAAG,EAAE,EAAE,CAAA;AACnE,CAAC;AAED,oEAAoE;AACpE,MAAM,UAAU,SAAS,CAAC,CAAQ,EAAE,KAAY,EAAE,KAAY,EAAE,MAAqB,EAAE,IAAqB;IAC1G,MAAM,EAAE,GAAG,CAAC,CAAC,OAAO,CAAC,MAAM,CAAA;IAC3B,MAAM,CAAC,GAAW,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,CAAA;IACpD,CAAC,CAAC,OAAO,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;IACjB,OAAO,CAAC,CAAA;AACV,CAAC;AAED,kFAAkF;AAClF,0EAA0E;AAC1E,+EAA+E;AAC/E,kFAAkF;AAClF,MAAM,UAAU,KAAK,CACnB,CAAQ,EACR,IAAO,EACP,KAAY,EACZ,KAAY,EACZ,IAAqB,EACrB,MAA0D;IAE1D,MAAM,OAAO,GAAG,CAAC,CAAC,GAAG,CAAC,MAAM,CAAA;IAC5B,MAAM,GAAG,GAAG,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,OAAO,EAAE,IAAI,CAAC,CAAA;IACrD,MAAM,IAAI,GAAG,EAAE,IAAI,EAAE,GAAG,EAAE,GAAG,CAAC,EAAE,EAAE,GAAG,MAAM,EAAkC,CAAA;IAC7E,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;IAChB,OAAO,GAAG,CAAA;AACZ,CAAC;AAED,0EAA0E;AAC1E,iFAAiF;AACjF,MAAM,UAAU,WAAW,CAAC,MAAc;IACxC,+EAA+E;IAC/E,MAAM,KAAK,GAAG,CAAC,IAAI,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,EAAE,CAAA;IACvC,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,CAAA;AAC1B,CAAC;AAED,8EAA8E;AAC9E,2DAA2D;AAC3D,MAAM,UAAU,UAAU,CAAC,IAAc;IACvC,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,CAAC,CAAA;IACpC,2EAA2E;IAC3E,iEAAiE;IACjE,MAAM,UAAU,GAAa,EAAE,CAAA;IAC/B,KAAK,MAAM,IAAI,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC;QAClC,IAAI,IAAI,CAAC,QAAQ,CAAC,kBAAkB,CAAC,IAAI,IAAI,CAAC,QAAQ,CAAC,qBAAqB,CAAC;YAAE,SAAQ;QACvF,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,EAAE,CAAC,CAAA;QAC5B,IAAI,UAAU,CAAC,MAAM,IAAI,CAAC;YAAE,MAAK;IACnC,CAAC;IACD,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO,IAAI,IAAI,CAAC,MAAM,yBAAyB,CAAA;IAC5E,OAAO,IAAI,IAAI,CAAC,MAAM,QAAQ,UAAU,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAA;AACzD,CAAC"}

package/dist/module.d.ts

@@ -1,4 +1,21 @@
 import type { Tensor, Shape, Dtype } from './ir.js';
+/** How a parameter's initial values are produced.
+ *  - `'randn'` — Gaussian, with `scale` (default 0.02). The common case for
+ *    weight matrices and embeddings.
+ *  - `'zeros'` — fill with 0. Common for biases and LayerNorm beta.
+ *  - `'ones'`  — fill with 1. Common for LayerNorm gain.
+ *  - Custom function — receives total element count and shape, returns the
+ *    Float32Array. Use for fan-in scaling or any non-standard scheme.
+ */
+export type InitSpec = 'randn' | 'zeros' | 'ones' | ((size: number, shape: readonly number[]) => Float32Array);
+export interface ParamOptions {
+    dtype?: Dtype;
+    /** Init kind. Default: `'randn'`. */
+    init?: InitSpec;
+    /** Std dev for `'randn'`. Default 0.02. Ignored for non-randn init. */
+    scale?: number;
+}
+type InitFn = (size: number, shape: readonly number[]) => Float32Array;
 export declare abstract class Module {
     /**
      * Declare a learnable parameter at this module. Must be called from inside
@@ -6,16 +23,25 @@ export declare abstract 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.
      */
-    protected param(shape: Shape, dtype?: Dtype): Tensor;
+    protected param(shape: Shape, opts?: ParamOptions): Tensor;
+}
+export interface MaterializedParams {
+    /** Map from auto-derived path (e.g. `layers.0.attn.W_q`) to its Tensor. */
+    tensors: Record<string, Tensor>;
+    /** Init function per param path. Used by `uploadInitialParams`. */
+    initFns: Record<string, InitFn>;
 }
 /**
  * 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 declare function materializeParams(root: Module): Record<string, Tensor>;
+export declare function materializeParams(root: Module): MaterializedParams;
+export {};
 //# sourceMappingURL=module.d.ts.map

package/dist/module.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"module.d.ts","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AA2BA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAoBnD,8BAAsB,MAAM;IAC1B;;;;;;;OAOG;IACH,SAAS,CAAC,KAAK,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM;CAI5D;AAMD;;;;;;GAMG;AACH,wBAAgB,iBAAiB,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAUtE"}
+{"version":3,"file":"module.d.ts","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AA2BA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAOnD;;;;;;;GAOG;AACH,MAAM,MAAM,QAAQ,GAChB,OAAO,GACP,OAAO,GACP,MAAM,GACN,CAAC,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,MAAM,EAAE,KAAK,YAAY,CAAC,CAAA;AAE9D,MAAM,WAAW,YAAY;IAC3B,KAAK,CAAC,EAAE,KAAK,CAAA;IACb,qCAAqC;IACrC,IAAI,CAAC,EAAE,QAAQ,CAAA;IACf,uEAAuE;IACvE,KAAK,CAAC,EAAE,MAAM,CAAA;CACf;AAED,KAAK,MAAM,GAAG,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,MAAM,EAAE,KAAK,YAAY,CAAA;AA2CtE,8BAAsB,MAAM;IAC1B;;;;;;;;OAQG;IACH,SAAS,CAAC,KAAK,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,CAAC,EAAE,YAAY,GAAG,MAAM;CAK3D;AAMD,MAAM,WAAW,kBAAkB;IACjC,2EAA2E;IAC3E,OAAO,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IAC/B,mEAAmE;IACnE,OAAO,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CAChC;AAED;;;;;;;GAOG;AACH,wBAAgB,iBAAiB,CAAC,IAAI,EAAE,MAAM,GAAG,kBAAkB,CAYlE"}

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/ops.d.ts

@@ -31,5 +31,5 @@ export declare function where(cond: Tensor, a: Tensor, b: Tensor): Tensor;
 export declare function reluGrad(x: Tensor, dy: Tensor): Tensor;
 export declare function adamUpdateM(m: Tensor, g: Tensor, b1: number): Tensor;
 export declare function adamUpdateV(v: Tensor, g: Tensor, b2: number): Tensor;
-export declare function adamUpdateP(p: Tensor, mNew: Tensor, vNew: Tensor, lrt: Tensor, eps: number): Tensor;
+export declare function adamUpdateP(p: Tensor, mNew: Tensor, vNew: Tensor, lrt: Tensor, eps: number, decayShrink?: number): Tensor;
 //# sourceMappingURL=ops.d.ts.map

package/dist/ops.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"ops.d.ts","sourceRoot":"","sources":["../src/ops.ts"],"names":[],"mappings":"AAWA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAU,MAAM,SAAS,CAAA;AAmC3D,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAMzD;AAQD,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAG3D;AAED,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAG3D;AAYD,eAAO,MAAM,IAAI,GAAK,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,KAAK,GAAI,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,GAAG,GAAM,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,GAAG,GAAM,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,IAAI,GAAK,GAAG,MAAM,KAAG,MAA2B,CAAA;AAO7D,wBAAgB,QAAQ,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAK1C;AAED,wBAAgB,OAAO,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAKzC;AAMD,wBAAgB,OAAO,CAAC,CAAC,EAAE,MAAM,EAAE,QAAQ,EAAE,KAAK,GAAG,MAAM,CAI1D;AAED,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,SAAS,MAAM,EAAE,GAAG,MAAM,CAIpE;AAMD,wBAAgB,MAAM,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAOnD;AAED,wBAAgB,aAAa,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAO1D;AAMD,wBAAgB,MAAM,CAAC,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAOnF;AAGD,wBAAgB,MAAM,CAAC,CAAC,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAM9D;AASD,wBAAgB,iBAAiB,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAKnD;AAGD,wBAAgB,cAAc,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAIhD;AAMD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,MAAM,GAAG,MAAM,CAKhE;AAQD,wBAAgB,cAAc,CAAC,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,MAAM,CAI5E;AAOD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,WAAW,EAAE,KAAK,GAAG,MAAM,CAIjE;AAED,wBAAgB,UAAU,CAAC,CAAC,EAAE,MAAM,EAAE,WAAW,EAAE,KAAK,GAAG,MAAM,CAIhE;AAOD,wBAAgB,WAAW,CAAC,KAAK,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAGvE;AAWD,eAAO,MAAM,IAAI,GAAO,GAAG,MAAM,EAAE,GAAG,MAAM,KAAG,MAAqD,CAAA;AACpG,eAAO,MAAM,OAAO,GAAI,GAAG,MAAM,EAAE,GAAG,MAAM,KAAG,MAAqD,CAAA;AAGpG,wBAAgB,KAAK,CAAC,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAMhE;AAID,wBAAgB,QAAQ,CAAC,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOtD;AAMD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOpE;AAED,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOpE;AAED,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,MAAM,CAWnG"}
+{"version":3,"file":"ops.d.ts","sourceRoot":"","sources":["../src/ops.ts"],"names":[],"mappings":"AAWA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAU,MAAM,SAAS,CAAA;AAmC3D,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAEzD;AACD,wBAAgB,GAAG,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,GAAG,MAAM,CAMzD;AAQD,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAG3D;AAED,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAG3D;AAYD,eAAO,MAAM,IAAI,GAAK,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,KAAK,GAAI,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,GAAG,GAAM,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,GAAG,GAAM,GAAG,MAAM,KAAG,MAA2B,CAAA;AAC7D,eAAO,MAAM,IAAI,GAAK,GAAG,MAAM,KAAG,MAA2B,CAAA;AAO7D,wBAAgB,QAAQ,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAK1C;AAED,wBAAgB,OAAO,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAKzC;AAMD,wBAAgB,OAAO,CAAC,CAAC,EAAE,MAAM,EAAE,QAAQ,EAAE,KAAK,GAAG,MAAM,CAI1D;AAED,wBAAgB,SAAS,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,SAAS,MAAM,EAAE,GAAG,MAAM,CAIpE;AAMD,wBAAgB,MAAM,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAOnD;AAED,wBAAgB,aAAa,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAO1D;AAMD,wBAAgB,MAAM,CAAC,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAOnF;AAGD,wBAAgB,MAAM,CAAC,CAAC,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAM9D;AASD,wBAAgB,iBAAiB,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAKnD;AAGD,wBAAgB,cAAc,CAAC,CAAC,EAAE,MAAM,GAAG,MAAM,CAIhD;AAMD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,MAAM,GAAG,MAAM,CAKhE;AAQD,wBAAgB,cAAc,CAAC,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,MAAM,CAI5E;AAOD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,WAAW,EAAE,KAAK,GAAG,MAAM,CAIjE;AAED,wBAAgB,UAAU,CAAC,CAAC,EAAE,MAAM,EAAE,WAAW,EAAE,KAAK,GAAG,MAAM,CAIhE;AAOD,wBAAgB,WAAW,CAAC,KAAK,EAAE,MAAM,EAAE,KAAK,GAAE,KAAa,GAAG,MAAM,CAGvE;AAWD,eAAO,MAAM,IAAI,GAAO,GAAG,MAAM,EAAE,GAAG,MAAM,KAAG,MAAqD,CAAA;AACpG,eAAO,MAAM,OAAO,GAAI,GAAG,MAAM,EAAE,GAAG,MAAM,KAAG,MAAqD,CAAA;AAGpG,wBAAgB,KAAK,CAAC,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,GAAG,MAAM,CAMhE;AAID,wBAAgB,QAAQ,CAAC,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOtD;AAMD,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOpE;AAED,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,EAAE,EAAE,MAAM,GAAG,MAAM,CAOpE;AAED,wBAAgB,WAAW,CAAC,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,WAAW,GAAE,MAAU,GAAG,MAAM,CAW5H"}

package/dist/ops.js

@@ -255,7 +255,7 @@ export function adamUpdateV(v, g, b2) {
     }
     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) {
+export function adamUpdateP(p, mNew, vNew, lrt, eps, decayShrink = 1) {
     const site = captureSite('adamUpdateP');
     if (p.dtype !== 'f32')
         throw new ShapeError(`adamUpdateP: requires f32`, site);
@@ -265,6 +265,6 @@ export function adamUpdateP(p, mNew, vNew, lrt, eps) {
     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 });
+    return addOp(currentGraph(), 'adam_update_p', p.shape, 'f32', site, { p: p.id, mNew: mNew.id, vNew: vNew.id, lrt: lrt.id, eps, decayShrink });
 }
 //# sourceMappingURL=ops.js.map