tensorgrad 0.0.2 → 0.0.4

package/dist/trace.js

@@ -16,6 +16,10 @@
 import { makeGraph, addOp, captureSite } from './ir.js';
 // Module-local: the graph being built right now, or null if no trace is active.
 let _current = null;
+// Module-local: whether `capture(name, t)` calls should register on the current
+// graph. True only during the user's forward trace; false during `traceInto`
+// (autograd / optimizer ops shouldn't accidentally publish gradient tensors).
+let _captureEnabled = false;
 export function currentGraph() {
     if (!_current) {
         throw new Error('tensorgrad: ops can only be called inside trace(). ' +
@@ -23,6 +27,9 @@ export function currentGraph() {
     }
     return _current;
 }
+export function isCaptureEnabled() {
+    return _captureEnabled;
+}
 // Run `fn` with a fresh graph as the current one; capture and return the graph.
 // `fn` must return the tensor (or array of tensors) to mark as graph outputs.
 export function trace(fn) {
@@ -31,6 +38,7 @@ export function trace(fn) {
     }
     const g = makeGraph();
     _current = g;
+    _captureEnabled = true;
     try {
         const result = fn();
         const outputs = Array.isArray(result) ? result : [result];
@@ -41,18 +49,22 @@ export function trace(fn) {
     }
     finally {
         _current = null;
+        _captureEnabled = false;
     }
     return g;
 }
 // Re-enter an existing graph to append more ops. Used by autograd to add
 // backward ops to a graph that's already been traced. `fn` runs with the
 // supplied graph as the current one; any ops it calls append to that graph.
+// Capture is intentionally disabled here — backward / optimizer rules
+// shouldn't publish their internal tensors via `capture()`.
 // Returns whatever `fn` returns.
 export function traceInto(g, fn) {
     if (_current) {
         throw new Error('tensorgrad: traceInto() called while another trace is active');
     }
     _current = g;
+    // _captureEnabled stays false (default) — explicit, but not toggled.
     try {
         return fn();
     }

package/dist/trace.js.map

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

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tensorgrad",
-  "version": "0.0.2",
+  "version": "0.0.4",
   "description": "Tiny TypeScript-native tensor library with autograd, compiling to WebGPU. Train small models in the browser without hand-writing kernels.",
   "license": "MIT",
   "author": "Ben Albahari",
@@ -45,7 +45,6 @@
   "files": [
     "dist",
     "src",
-    "SPEC.md",
     "README.md",
     "LICENSE"
   ],

package/src/buffers.ts

@@ -47,6 +47,7 @@ export interface BufferPlan {
   inputsByName: Map<string, number>           // name -> buffer id
   paramGradsByName: Map<string, number>       // name -> buffer id
   statesByName: Map<string, number>           // name -> buffer id (persistent state homes)
+  capturesByName: Map<string, number>         // name -> buffer id (activation captures)
   outputBufferIds: number[]                   // graph.outputs mapped through
   /** End-of-step writebacks (Adam updates for params, m, v, etc.) */
   writebacks: Writeback[]
@@ -169,5 +170,17 @@ export function planBuffers(
     return { source: sourceBufId, dest: destBufId, bytes: sourceSpec.byteSize }
   })
 
-  return { buffers, tensorToBuffer, paramsByName, inputsByName, paramGradsByName, statesByName, outputBufferIds, writebacks }
+  // Resolve graph.captures (name -> tensor id) to (name -> buffer id).
+  // No pinning needed at the planner level: each tensor already has its own
+  // buffer (see "v1 strategy" comment at top — no pooling yet).
+  const capturesByName = new Map<string, number>()
+  for (const [name, tensorId] of graph.captures) {
+    const bufId = tensorToBuffer.get(tensorId)
+    if (bufId === undefined) {
+      throw new Error(`planBuffers: capture '${name}' references unknown tensor #${tensorId}`)
+    }
+    capturesByName.set(name, bufId)
+  }
+
+  return { buffers, tensorToBuffer, paramsByName, inputsByName, paramGradsByName, statesByName, capturesByName, outputBufferIds, writebacks }
 }

package/src/capture.ts

@@ -0,0 +1,36 @@
+// Activation capture — opt-in readback of intermediate tensors at training step.
+//
+// Usage (inside the user's forward pass):
+//
+//   import { capture } from 'tensorgrad'
+//
+//   function attentionFwd(p, x) {
+//     const scores = mul(matmulBatched(q, kT), SCALE_QK)
+//     const attn = capture(`attn.${layerIdx}`, softmaxCausalLast(scores))
+//     return matmulBatched(attn, v)
+//   }
+//
+// Pass-through return type: `capture(name, t)` returns `t` unchanged so it
+// inlines at the point of computation. Behind the scenes it registers `t.id`
+// against `name` on the current graph; runtime exposes the registered tensors
+// via `step(inputs, { withCaptures: true })`.
+//
+// Outside the user's forward trace (during `appendGrad` / `appendAdam`'s
+// `traceInto` re-entry), `capture()` is a no-op — gradient and optimizer
+// internals shouldn't accidentally publish themselves to the UI.
+
+import type { Tensor } from './ir.js'
+import { currentGraph, isCaptureEnabled } from './trace.js'
+
+export function capture<T extends Tensor>(name: string, t: T): T {
+  if (!isCaptureEnabled()) return t
+  const g = currentGraph()
+  if (g.captures.has(name)) {
+    throw new Error(
+      `capture: name '${name}' already registered. Use unique names ` +
+      `(e.g. \`attn.\${layerIdx}\`) when capturing across a loop.`,
+    )
+  }
+  g.captures.set(name, t.id)
+  return t
+}

package/src/compile.ts

@@ -14,7 +14,7 @@ import { appendGrad, type GradResult } from './grad.js'
 import { appendAdam, type AdamConfig } from './adam.js'
 import { planBuffers, type BufferPlan } from './buffers.js'
 import { emitKernels, type KernelSpec } from './codegen.js'
-import { createRuntime, type CompiledRuntime, type RuntimeOpts } from './runtime.js'
+import { createRuntime, createForwardRuntime, type CompiledRuntime, type CompiledForward, type RuntimeOpts } from './runtime.js'
 import { Module, materializeParams } from './module.js'
 
 /** Declares one input tensor of the model's forward function. Order matches
@@ -65,10 +65,19 @@ export interface CompileModuleOptions extends RuntimeOpts {
   adam?: AdamConfig
 }
 
+export interface CompileForwardOptions extends RuntimeOpts {
+  /** Per-step data inputs to the forward function. */
+  inputs?: InputDecl[]
+}
+
 /**
- * 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.
@@ -78,14 +87,15 @@ export interface CompileModuleOptions extends RuntimeOpts {
  * users don't need to provide it themselves.
  */
 export async function compileModule<M extends Module>(
-  model: M,
+  modelFactory: () => M,
   forward: (m: M, ...inputs: Tensor[]) => Tensor,
   opts: CompileModuleOptions = {},
-): Promise<CompiledRuntime & { ir: CompiledIR }> {
+): Promise<CompiledRuntime & { ir: CompiledIR; uploadInitialParams: () => void }> {
   const inputDecls = opts.inputs ?? []
-  let paramTensors: Record<string, Tensor> = {}
+  const model = modelFactory()
+  let materialized: ReturnType<typeof materializeParams> = { 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)
   })
@@ -94,7 +104,7 @@ export async function compileModule<M extends Module>(
 
   let adamResult: ReturnType<typeof appendAdam> | undefined
   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 ?? [])
@@ -103,18 +113,107 @@ export async function compileModule<M extends Module>(
   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 innerStep = runtime.step.bind(runtime) as CompiledRuntime['step']
+    const innerReset = runtime.resetOptimizerState.bind(runtime)
+    const wrappedStep = (
+      inputs: Record<string, Int32Array | Float32Array>,
+      opts?: { withCaptures?: boolean },
+    ): Promise<number | { loss: number; captures: Record<string, Float32Array> }> => {
       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 as CompiledRuntime['step']
+    runtime.resetOptimizerState = () => {
+      t = 0
+      innerReset()
+    }
+  }
+
+  const { initFns } = materialized
+  const uploadInitialParams = () => {
+    const out: Record<string, Float32Array> = {}
+    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: CompiledIR = { 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<M extends Module>(
+  modelFactory: () => M,
+  forward: (m: M, ...inputs: Tensor[]) => Tensor,
+  opts: CompileForwardOptions = {},
+): Promise<CompiledForward & { ir: CompiledIR; uploadInitialParams: () => void }> {
+  const inputDecls = opts.inputs ?? []
+  const model = modelFactory()
+  let materialized: ReturnType<typeof materializeParams> = { 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: Record<string, Float32Array> = {}
+    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: CompiledIR = { graph, paramGrads: {}, loss: outputTensor, plan, kernels }
+  return Object.assign(runtime, { ir, uploadInitialParams })
 }

package/src/index.ts

@@ -6,6 +6,7 @@
 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 {
   // Element-wise arithmetic. The binops accept Tensor or JS-number for the second arg.
   add, sub, mul, div,
@@ -35,6 +36,7 @@ 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'

package/src/ir.ts

@@ -141,10 +141,14 @@ export interface Graph {
   // Names of tensors that should be exposed as outputs of the compiled function.
   // Set by the trace driver; for a loss function, this is `[lossTensor]`.
   readonly outputs: number[]
+  // Tensors registered for activation readback via `capture(name, t)`.
+  // Keyed by user-supplied name; insertion order preserved. Empty when no
+  // captures registered (the common training case — zero overhead).
+  readonly captures: Map<string, number>
 }
 
 export function makeGraph(): Graph {
-  return { ops: [], tensors: [], outputs: [] }
+  return { ops: [], tensors: [], outputs: [], captures: new Map() }
 }
 
 // Internal: register a fresh tensor in the graph and return its id.

package/src/module.ts

@@ -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 {
@@ -28,6 +28,54 @@
 import type { Tensor, Shape, Dtype } from './ir.js'
 import { paramInput } from './trace.js'
 
+// ============================================================================
+// Init metadata
+// ============================================================================
+
+/** 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
+
+function boxMuller(): number {
+  return Math.sqrt(-2 * Math.log(Math.max(1e-10, Math.random()))) * Math.cos(2 * Math.PI * Math.random())
+}
+
+function resolveInit(opts: ParamOptions | undefined): InitFn {
+  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
 // ============================================================================
@@ -38,7 +86,11 @@ import { paramInput } from './trace.js'
 // only valid post-materialization (which is always before forward runs).
 
 class ParamSentinel {
-  constructor(public readonly shape: Shape, public readonly dtype: Dtype) {}
+  constructor(
+    public readonly shape: Shape,
+    public readonly dtype: Dtype,
+    public readonly initFn: InitFn,
+  ) {}
 }
 
 // ============================================================================
@@ -52,11 +104,13 @@ export 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 = 'f32'): Tensor {
+  protected param(shape: Shape, opts?: ParamOptions): Tensor {
+    const dtype = opts?.dtype ?? 'f32'
     // Lie to TypeScript: the sentinel becomes a Tensor at materialize time.
-    return new ParamSentinel(shape, dtype) as unknown as Tensor
+    return new ParamSentinel(shape, dtype, resolveInit(opts)) as unknown as Tensor
   }
 }
 
@@ -64,23 +118,33 @@ export abstract class Module {
 // Tree walking
 // ============================================================================
 
+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 function materializeParams(root: Module): Record<string, Tensor> {
-  const out: Record<string, Tensor> = {}
+export function materializeParams(root: Module): MaterializedParams {
+  const tensors: Record<string, Tensor> = {}
+  const initFns: Record<string, InitFn> = {}
   visit(root, '', (path, val, owner, key) => {
     if (val instanceof ParamSentinel) {
       const t = paramInput(path, val.shape, val.dtype)
       ;(owner as any)[key] = t
-      out[path] = t
+      tensors[path] = t
+      initFns[path] = val.initFn
     }
   })
-  return out
+  return { tensors, initFns }
 }
 
 // ----------------------------------------------------------------------------

package/src/nn.ts

@@ -0,0 +1,59 @@
+// 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 type { Tensor } from './ir.js'
+import { add, matmul, sub, mul, div, sqrt, meanLast } from './ops.js'
+
+// ----------------------------------------------------------------------------
+// Linear: y = x @ W (+ b)
+// ----------------------------------------------------------------------------
+
+export class Linear extends Module {
+  W: Tensor
+  b: Tensor | null
+  constructor(public readonly inDim: number, public readonly outDim: number, withBias = true) {
+    super()
+    this.W = this.param([inDim, outDim])                      // randn, scale 0.02
+    this.b = withBias ? this.param([outDim], { init: 'zeros' }) : null
+  }
+}
+
+export function linearFwd(p: Linear, x: Tensor): Tensor {
+  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 {
+  g: Tensor
+  b: Tensor
+  constructor(public readonly d: number, public readonly eps: number = 1e-5) {
+    super()
+    this.g = this.param([d], { init: 'ones' })
+    this.b = this.param([d], { init: 'zeros' })
+  }
+}
+
+export function layerNormFwd(p: LayerNorm, x: Tensor): Tensor {
+  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)
+}