tensorgrad 0.0.2 → 0.0.5

package/src/adam.ts

@@ -15,6 +15,12 @@
 // `sqrt(1-b2^t)/(1-b1^t)`; that's why convergence isn't affected by the
 // first-step warmup that bias-correction-free Adam suffers.
 //
+// **Static vs scheduled lr.** When `config.lr` is a number, decayShrink is
+// baked into the kernel as a literal. When it's a function `(step) => lr`,
+// decayShrink for decayed params becomes a per-step scalar input that the
+// runtime updates each call (computed from the current step's lr). lrt is
+// always per-step; the bias-correction factor changes every step regardless.
+//
 // Returns writeback declarations the buffer planner uses to wire up the
 // "after step, copy the new value into the persistent home" path. m and v
 // are state_inputs (zero-initialized, persistent across steps); the param
@@ -27,7 +33,11 @@ import { traceInto, stateInput, tensorInput } from './trace.js'
 import { adamUpdateM, adamUpdateV, adamUpdateP } from './ops.js'
 
 export interface AdamConfig {
-  lr: number
+  /** Constant scalar (e.g., `0.005`) or a per-step schedule function
+   *  `(step) => lr`. Schedule fn lets the user implement linear/cosine decay
+   *  or warmup; first call passes `step=1`. Decay-shrink (AdamW) updates
+   *  per-step automatically when this is a function. */
+  lr: number | ((step: number) => number)
   b1?: number   // default 0.9
   b2?: number   // default 0.999
   eps?: number  // default 1e-8
@@ -42,14 +52,30 @@ export interface AdamConfig {
   decayFilter?: (paramName: string) => boolean
 }
 
+/** Resolved hyperparameters: lr is the schedule fn (constants are wrapped). */
+export interface AdamResolvedConfig {
+  lr: (step: number) => number
+  b1: number
+  b2: number
+  eps: number
+  weightDecay: number
+  decayFilter: (name: string) => boolean
+  /** True iff the user supplied an lr function (vs a constant). When false,
+   *  decayShrink is baked at compile time and never updated. */
+  lrIsScheduled: boolean
+}
+
 export interface AdamResult {
   /** Writebacks the buffer planner should wire into the runtime. */
   writebacks: WritebackDecl[]
   /** Name of the per-step scalar tensor_input. The runtime fills this each call
    * with `lr * sqrt(1-b2^t)/(1-b1^t)` (Adam's bias-corrected effective LR). */
   lrtInputName: string
-  /** Hyperparameters as captured (so the runtime can compute lrt). */
-  config: Required<Omit<AdamConfig, 'decayFilter'>> & { decayFilter: (name: string) => boolean }
+  /** Name of the per-step decayShrink scalar tensor_input, or null when lr is
+   *  static (decayShrink baked into the kernel) or no params are decayed. */
+  decayShrinkInputName: string | null
+  /** Hyperparameters as captured (so the runtime can compute lrt and decayShrink). */
+  config: AdamResolvedConfig
 }
 
 /**
@@ -70,22 +96,40 @@ export function appendAdam(
   paramTensors: Record<string, Tensor>,
   config: AdamConfig,
 ): AdamResult {
-  const fullConfig = {
-    lr: config.lr,
+  const lrIsScheduled = typeof config.lr === 'function'
+  const lrFn = lrIsScheduled
+    ? config.lr as (step: number) => number
+    : (() => config.lr as number)
+  const initialLr = lrFn(1)
+  const fullConfig: AdamResolvedConfig = {
+    lr: lrFn,
     b1: config.b1 ?? 0.9,
     b2: config.b2 ?? 0.999,
     eps: config.eps ?? 1e-8,
     weightDecay: config.weightDecay ?? 0,
     decayFilter: config.decayFilter ?? (() => true),
+    lrIsScheduled,
   }
   const writebacks: WritebackDecl[] = []
   const lrtInputName = '_adam_lrt'
+  // Tensor input for runtime-updated decayShrink (only created when lr is a
+  // schedule fn AND at least one param will receive weight decay).
+  let decayShrinkInputName: string | null = null
 
   return traceInto(graph, () => {
-    // One scalar lrt input shared by every adam_update_p call. Runtime supplies
-    // it per step as `lr * sqrt(1-b2^t) / (1-b1^t)`.
     const lrt = tensorInput(lrtInputName, [], 'f32')
 
+    // Decide up-front whether we need a runtime decayShrink scalar. Only does
+    // something when both (a) lr varies per step and (b) some param is decayed.
+    const needsDynamicShrink = lrIsScheduled
+      && fullConfig.weightDecay > 0
+      && Object.keys(paramGrads).some(name => fullConfig.decayFilter(name))
+    let decayShrinkScalar: Tensor | null = null
+    if (needsDynamicShrink) {
+      decayShrinkInputName = '_adam_decay_shrink'
+      decayShrinkScalar = tensorInput(decayShrinkInputName, [], 'f32')
+    }
+
     for (const name of Object.keys(paramGrads)) {
       const p = paramTensors[name]
       const g = paramGrads[name]
@@ -95,12 +139,19 @@ export function appendAdam(
       const mState = stateInput(`adam_m_${name}`, p.shape, 'f32', 0)
       const vState = stateInput(`adam_v_${name}`, p.shape, 'f32', 0)
 
-      // decayShrink baked at compile time. 1.0 for plain Adam (no extra cost
-      // — the WGSL compiler folds the constant multiply); 1 - lr * weightDecay
-      // for the params the filter selects.
-      const decayShrink = (fullConfig.weightDecay > 0 && fullConfig.decayFilter(name))
-        ? 1 - fullConfig.lr * fullConfig.weightDecay
-        : 1
+      // Choose the decayShrink form per param:
+      //   - non-decayed params: literal 1 (kernel multiply folds out).
+      //   - decayed + static lr: literal `1 - lr * wd` baked at compile.
+      //   - decayed + scheduled lr: tensor input updated per step.
+      const isDecayed = fullConfig.weightDecay > 0 && fullConfig.decayFilter(name)
+      let decayShrink: number | Tensor
+      if (!isDecayed) {
+        decayShrink = 1
+      } else if (decayShrinkScalar !== null) {
+        decayShrink = decayShrinkScalar
+      } else {
+        decayShrink = 1 - initialLr * fullConfig.weightDecay
+      }
 
       // Three fused kernels per parameter — one for each of m_new / v_new / p_new.
       const newM = adamUpdateM(mState, g, fullConfig.b1)
@@ -111,6 +162,6 @@ export function appendAdam(
       writebacks.push({ source: newV, destName: `adam_v_${name}`, destKind: 'state' })
       writebacks.push({ source: newP, destName: name,             destKind: 'param' })
     }
-    return { writebacks, lrtInputName, config: fullConfig }
+    return { writebacks, lrtInputName, decayShrinkInputName, config: fullConfig }
   })
 }

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/codegen.ts

@@ -557,23 +557,34 @@ fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
     case 'adam_update_p': {
       // p_new = decayShrink * p - lrt[0] * m_new / (sqrt(v_new) + eps).
       // lrt is supplied per-step from CPU (already includes bias correction).
-      // decayShrink encodes AdamW's decoupled weight decay; when no decay is
-      // requested it's exactly 1.0 and the WGSL compiler folds the multiply away.
+      // decayShrink is either baked as a literal (no schedule, fixed lr) or
+      // bound as a per-step scalar input (when the user supplies an lr
+      // schedule via `adam: { lr: (step) => ... }`). When literal=1 the WGSL
+      // compiler folds the multiply away.
       const out = tof(op.out)
       const total = shapeSize(out.shape)
+      const dynamicShrink = op.decayShrinkTensor !== null
+      const shrinkExpr = dynamicShrink ? 'decayShrink[0]' : wgslLiteral(op.decayShrink, 'f32')
+      const shrinkBinding = dynamicShrink
+        ? `@group(0) @binding(4) var<storage, read> decayShrink : array<f32>;\n` +
+          `@group(0) @binding(5) var<storage, read_write> out : array<f32>;`
+        : `@group(0) @binding(4) var<storage, read_write> out : array<f32>;`
       const wgsl = `
 @group(0) @binding(0) var<storage, read> p : array<f32>;
 @group(0) @binding(1) var<storage, read> mNew : array<f32>;
 @group(0) @binding(2) var<storage, read> vNew : array<f32>;
 @group(0) @binding(3) var<storage, read> lrt : array<f32>;
-@group(0) @binding(4) var<storage, read_write> out : array<f32>;
+${shrinkBinding}
 @compute @workgroup_size(${WG_SIZE})
 fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
   let i = gid.x + gid.y * 16776960u;
   if (i >= ${total}u) { return; }
-  out[i] = ${wgslLiteral(op.decayShrink, 'f32')} * p[i] - lrt[0] * mNew[i] / (sqrt(vNew[i]) + ${wgslLiteral(op.eps, 'f32')});
+  out[i] = ${shrinkExpr} * p[i] - lrt[0] * mNew[i] / (sqrt(vNew[i]) + ${wgslLiteral(op.eps, 'f32')});
 }`.trim()
-      return { opIndex, opKind: op.kind, wgsl, bindings: [buf(op.p), buf(op.mNew), buf(op.vNew), buf(op.lrt), buf(op.out)], threads: total, workgroupSize: WG_SIZE }
+      const bindings = dynamicShrink
+        ? [buf(op.p), buf(op.mNew), buf(op.vNew), buf(op.lrt), buf(op.decayShrinkTensor!), buf(op.out)]
+        : [buf(op.p), buf(op.mNew), buf(op.vNew), buf(op.lrt), buf(op.out)]
+      return { opIndex, opKind: op.kind, wgsl, bindings, threads: total, workgroupSize: WG_SIZE }
     }
 
     case 'sum_to_shape': {

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 ?? [])
@@ -102,19 +112,115 @@ export async function compileModule<M extends Module>(
   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.
+  // If Adam is enabled, wrap step() to track the step count and supply lrt
+  // (and optionally decayShrink, when the user passed a per-step lr schedule).
+  // 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
+    const { lrtInputName, decayShrinkInputName, config } = adamResult
     let t = 0
     const lrtBuf = new Float32Array(1)
-    const innerStep = runtime.step.bind(runtime)
-    runtime.step = async (inputs) => {
+    const decayShrinkBuf = decayShrinkInputName ? new Float32Array(1) : null
+    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 lrNow = config.lr(t)
+      lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(config.b2, t)) / (1 - Math.pow(config.b1, t))
+      const merged: Record<string, Int32Array | Float32Array> = { ...inputs, [lrtInputName]: lrtBuf }
+      if (decayShrinkBuf && decayShrinkInputName) {
+        decayShrinkBuf[0] = 1 - lrNow * config.weightDecay
+        merged[decayShrinkInputName] = decayShrinkBuf
+      }
+      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

@@ -113,9 +113,21 @@ export type OpNode =
   // `lrt` is a scalar tensor (provided as a tensor_input updated per step) that
   // already includes Adam's bias-correction factor: lrt = lr * sqrt(1-b2^t) / (1-b1^t).
   // `decayShrink` is the decoupled-weight-decay factor (Loshchilov & Hutter,
-  // "AdamW") baked at compile time: 1 - lr * weightDecay when the param is being
-  // decayed, 1 otherwise. eps and decayShrink are both baked into the kernel.
-  | { kind: 'adam_update_p'; out: number; p: number; mNew: number; vNew: number; lrt: number; eps: number; decayShrink: number }
+  // "AdamW"): 1 - lr * weightDecay when the param is being decayed, 1 otherwise.
+  // It can be either a compile-time literal (number) for fixed-lr training, or a
+  // tensor id pointing at a scalar input that the runtime updates per step (used
+  // when the user supplies an lr schedule via `adam: { lr: (step) => ... }`).
+  | {
+      kind: 'adam_update_p'
+      out: number
+      p: number
+      mNew: number
+      vNew: number
+      lrt: number
+      eps: number
+      decayShrink: number               // literal (used when decayShrinkTensor is null)
+      decayShrinkTensor: number | null  // tensor id of a scalar input; takes precedence when set
+    }
 
   // ---- Slicing / broadcasting / autograd infrastructure -------------------
   // Slice [start, end) along the last axis. Output shape: input shape with
@@ -141,10 +153,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 }
 }
 
 // ----------------------------------------------------------------------------