tensorgrad 0.0.9 → 0.0.11

package/src/compile.ts

@@ -17,15 +17,32 @@ import { emitKernels, type KernelSpec } from './codegen.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
- *  the function's parameter list (after `model`). The `name` is used at
- *  runtime to upload data via `step({ [name]: data })`. */
+/** Declares one input tensor of the model's forward function. The name is the
+ *  key in the `inputs:` Record at compile time and the key on the `step()`/
+ *  `run()` data object at runtime. */
 export interface InputDecl {
-  name: string
   shape: Shape
   dtype?: Dtype
 }
 
+/** Inputs declaration: a Record from input name to its shape/dtype. The name
+ *  doubles as the key the forward fn destructures and the key the runtime
+ *  expects in `step({...})` / `run({...})`. */
+export type InputDecls = Record<string, InputDecl>
+
+/** Maps an `InputDecls` Record to its forward-time tensor counterpart —
+ *  same keys, each value is a Tensor. Used to type the forward function's
+ *  `inputs` argument from the declared shape Record. */
+export type InputsTensors<I extends InputDecls> = { [K in keyof I]: Tensor }
+
+/** Forward function shape: takes the materialized model and a Record of
+ *  named input tensors (matching the declared `inputs:` keys), returns the
+ *  output tensor (loss for compileModule; logits/etc. for compileForward).
+ *  The second generic flows from the inputs declaration so destructuring
+ *  the input record stays typed. */
+export type ForwardFn<M extends Module, I extends InputDecls = InputDecls> =
+  (m: M, inputs: InputsTensors<I>) => Tensor
+
 export interface CompiledIR {
   graph: GradResult['graph']
   paramGrads: GradResult['paramGrads']
@@ -55,19 +72,52 @@ export async function compile(traceFn: () => Tensor, opts: RuntimeOpts = {}): Pr
 // Module-aware compile
 // ============================================================================
 
-export interface CompileModuleOptions extends RuntimeOpts {
-  /** Per-step data inputs to the forward function. Order matches the forward
-   *  function's parameters (after the model). e.g. for
-   *  `(model, tokens, targets, mask) => loss`, inputs is
-   *  `[{name:'tokens',...}, {name:'targets',...}, {name:'mask',...}]`. */
-  inputs?: InputDecl[]
+export interface CompileModuleOptions<I extends InputDecls = InputDecls> extends RuntimeOpts {
+  /** Per-step data inputs to the forward function, keyed by name. The forward
+   *  fn destructures these out of its second argument; runtime calls to
+   *  `step()` / `run()` pass typed arrays under the same keys. */
+  inputs?: I
   /** Adam hyperparameters. If omitted, no optimizer is appended (forward-only). */
   adam?: AdamConfig
 }
 
-export interface CompileForwardOptions extends RuntimeOpts {
-  /** Per-step data inputs to the forward function. */
-  inputs?: InputDecl[]
+export interface CompileForwardOptions<I extends InputDecls = InputDecls> extends RuntimeOpts {
+  /** Per-step data inputs to the forward function, keyed by name. */
+  inputs?: I
+}
+
+/** Forward-only compile options as taken by the `compileForward` *method* on
+ *  a training runtime — no `device` (inherited) and no `sharedParams`
+ *  (auto-supplied from the train graph's params). */
+export interface CompileForwardMethodOptions<I extends InputDecls = InputDecls> {
+  inputs?: I
+}
+
+/** Returned by `compileModule`. Adds training-graph extras (auto-init, reset,
+ *  sibling-graph compile) on top of the base runtime. */
+export interface CompiledModule<M extends Module> extends CompiledRuntime {
+  ir: CompiledIR
+  /** Number of dispatchable kernels (excludes leaf no-ops). */
+  kernelCount: number
+  /** Re-initialize all params from their declared init specs and zero the
+   *  optimizer state. Use to start training over without recompiling. */
+  reset(): void
+  /** Compile a sibling forward-only graph (e.g., a B=1 inference graph or a
+   *  B=N held-out eval graph) that shares this runtime's device and param
+   *  buffers. Pass the forward fn (typically distinct from your loss fn —
+   *  it returns logits, not a scalar) and any shape changes via `inputs`.
+   *  Auto-initialization is a no-op since params are shared. */
+  compileForward<I extends InputDecls>(
+    forward: ForwardFn<M, I>,
+    opts?: CompileForwardMethodOptions<I>,
+  ): Promise<CompiledForwardModule>
+}
+
+/** Returned by `compileForward` (and by the `compileForward` method). */
+export interface CompiledForwardModule extends CompiledForward {
+  ir: CompiledIR
+  /** Number of dispatchable kernels (excludes leaf no-ops). */
+  kernelCount: number
 }
 
 /**
@@ -76,103 +126,70 @@ export interface CompileForwardOptions extends RuntimeOpts {
  * 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.
+ * The forward function takes the materialized model and a Record of named
+ * input tensors, returns the loss tensor. Inputs are matched by name with the
+ * `inputs:` declaration:
+ *
+ *   inputs: {
+ *     tokens:  { shape: [B, T], dtype: 'i32' },
+ *     targets: { shape: [B, T], dtype: 'i32' },
+ *   }
+ *   forward: (m, { tokens, targets }) => …
  *
  * Walks the module tree to materialize params with auto-derived names, then
- * runs trace → grad → adam → buffer plan → codegen → runtime.
+ * runs trace → grad → adam → buffer plan → codegen → runtime. Initial
+ * parameter values are uploaded automatically before this function returns;
+ * call `reset()` later to re-randomize.
  *
  * If `opts.adam` is set, the runtime's `step()` automatically tracks an
  * internal step count and injects the bias-corrected `lrt` scalar each call;
  * users don't need to provide it themselves.
  */
-export async function compileModule<M extends Module>(
+export async function compileModule<M extends Module, I extends InputDecls = InputDecls>(
   modelFactory: () => M,
-  forward: (m: M, ...inputs: Tensor[]) => Tensor,
-  opts: CompileModuleOptions = {},
-): Promise<CompiledRuntime & { ir: CompiledIR; uploadInitialParams: () => void }> {
-  const inputDecls = opts.inputs ?? []
-  const model = modelFactory()
-  let materialized: ReturnType<typeof materializeParams> = { tensors: {}, initFns: {}, decayFlags: {} }
-  const graph = trace(() => {
-    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: ReturnType<typeof appendAdam> | undefined
-  if (opts.adam) {
-    adamResult = appendAdam(graph, paramGrads, materialized.tensors, opts.adam, materialized.decayFlags)
-  }
-
-  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)
+  forward: ForwardFn<M, I>,
+  opts: CompileModuleOptions<I> = {},
+): Promise<CompiledModule<M>> {
+  const { runtime, materialized, plan, kernels, ir } = await buildModuleRuntime(
+    modelFactory, forward, opts, /* sharedParams */ undefined, /* withGrad */ true,
+  )
 
   // 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, decayShrinkInputName, config } = adamResult
-    let t = 0
-    const lrtBuf = new Float32Array(1)
-    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++
-      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()
-    }
+  if (opts.adam) {
+    wrapStepForAdam(runtime, opts.adam, ir)
   }
 
-  const uploadInitialParams = () => {
-    const out = buildInitialParamUploads(plan, materialized.initFns)
-    runtime.uploadParams(out)
-  }
+  // Auto-upload initial param values. Always wanted at this entry point —
+  // training runtimes own their params and need them randomized before step 1.
+  uploadInitialParams(plan, materialized.initFns, runtime, /* sharedParams */ undefined)
 
-  const ir: CompiledIR = { graph, paramGrads, loss, plan, kernels }
-  return Object.assign(runtime, { ir, uploadInitialParams })
-}
+  const kernelCount = kernels.filter(k => k.wgsl).length
 
-// Build a Record<paramName, Float32Array> by running each param's init
-// function against its shape. Shared by compileModule and compileForward.
-// `sharedParams`, when supplied, skips any name it covers (those are owned
-// by the sibling compile and already initialized there).
-type InitFn = (size: number, shape: readonly number[]) => Float32Array
-function buildInitialParamUploads(
-  plan: BufferPlan,
-  initFns: Record<string, InitFn>,
-  sharedParams?: Map<string, GPUBuffer>,
-): Record<string, Float32Array> {
-  const out: Record<string, Float32Array> = {}
-  for (const [name, bufId] of plan.paramsByName) {
-    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)
+  const reset = () => {
+    uploadInitialParams(plan, materialized.initFns, runtime, undefined)
+    runtime.resetOptimizerState()
+  }
+
+  const compileForwardMethod = async <J extends InputDecls>(
+    forwardFn: ForwardFn<M, J>,
+    fOpts: CompileForwardMethodOptions<J> = {},
+  ): Promise<CompiledForwardModule> => {
+    return compileForward<M, J>(modelFactory, forwardFn, {
+      ...fOpts,
+      device: runtime.device,
+      sharedParams: runtime.params,
+    })
   }
-  return out
+
+  return Object.assign(runtime, {
+    ir,
+    kernelCount,
+    reset,
+    compileForward: compileForwardMethod,
+  })
 }
 
 // ============================================================================
@@ -185,43 +202,157 @@ function buildInitialParamUploads(
  * scalar loss; runtime exposes `run(inputs)` returning the full output as a
  * `Float32Array`.
  *
+ * **Prefer the `compileForward` method on a training runtime** when both
+ * graphs use the same Module class — it auto-supplies `device` and
+ * `sharedParams`. This standalone form is for forward-only models with no
+ * training graph at all, or for sharing params across a different model.
+ *
  * **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`.
+ * visible to `run()` calls here, no copies.
  *
- * Typical use: a B=1 inference graph alongside a B=512 training graph,
- * built from the same `Module` factory.
+ * Initial param values are uploaded automatically for params *not* covered
+ * by `sharedParams` (those are owned by the sibling compile).
  */
-export async function compileForward<M extends Module>(
+export async function compileForward<M extends Module, I extends InputDecls = InputDecls>(
+  modelFactory: () => M,
+  forward: ForwardFn<M, I>,
+  opts: CompileForwardOptions<I> = {},
+): Promise<CompiledForwardModule> {
+  const sharedParams = opts.sharedParams
+  const { runtime, materialized, plan, kernels, ir } = await buildModuleRuntime(
+    modelFactory, forward, opts, sharedParams, /* withGrad */ false,
+  )
+
+  // Auto-upload initial values for any params this graph owns. With
+  // `sharedParams` covering everything, this is a no-op.
+  uploadInitialParams(plan, materialized.initFns, runtime, sharedParams)
+
+  const kernelCount = kernels.filter(k => k.wgsl).length
+  return Object.assign(runtime, { ir, kernelCount })
+}
+
+// ============================================================================
+// Internals
+// ============================================================================
+
+type InitFn = (size: number, shape: readonly number[]) => Float32Array
+
+interface BuiltRuntime {
+  runtime: CompiledRuntime
+  materialized: ReturnType<typeof materializeParams>
+  plan: BufferPlan
+  kernels: KernelSpec[]
+  ir: CompiledIR
+}
+
+/** Shared body of compileModule + compileForward. The training and forward
+ *  pipelines diverge only in (a) whether grad/Adam are appended and (b)
+ *  whether the output buffer is the loss scalar or the user's returned
+ *  tensor — both come out of the same trace and codegen path. */
+async function buildModuleRuntime<M extends Module, I extends InputDecls>(
   modelFactory: () => M,
-  forward: (m: M, ...inputs: Tensor[]) => Tensor,
-  opts: CompileForwardOptions = {},
-): Promise<CompiledForward & { ir: CompiledIR; uploadInitialParams: () => void }> {
-  const inputDecls = opts.inputs ?? []
+  forward: ForwardFn<M, I>,
+  opts: CompileModuleOptions<I> | CompileForwardOptions<I>,
+  sharedParams: Map<string, GPUBuffer> | undefined,
+  withGrad: boolean,
+): Promise<BuiltRuntime> {
+  const inputDecls: InputDecls = opts.inputs ?? {}
   const model = modelFactory()
   let materialized: ReturnType<typeof materializeParams> = { tensors: {}, initFns: {}, decayFlags: {} }
   const graph = trace(() => {
     materialized = materializeParams(model)
-    const inputTensors = inputDecls.map(d => tensorInput(d.name, d.shape, d.dtype ?? 'f32'))
-    return forward(model, ...inputTensors)
+    const inputTensors: Record<string, Tensor> = {}
+    for (const [name, decl] of Object.entries(inputDecls)) {
+      inputTensors[name] = tensorInput(name, decl.shape, decl.dtype ?? 'f32')
+    }
+    return forward(model, inputTensors as InputsTensors<I>)
   })
 
-  const plan = planBuffers(graph, /* paramGrads */ {})
+  let paramGrads: GradResult['paramGrads'] = {}
+  let outputTensor: Tensor
+  let adamWritebacks: ReturnType<typeof appendAdam>['writebacks'] = []
+
+  if (withGrad) {
+    const gradResult = appendGrad(graph)
+    paramGrads = gradResult.paramGrads
+    outputTensor = gradResult.loss
+    const adamCfg = (opts as CompileModuleOptions).adam
+    if (adamCfg) {
+      const adamResult = appendAdam(graph, paramGrads, materialized.tensors, adamCfg, materialized.decayFlags)
+      adamWritebacks = adamResult.writebacks
+      // Stash adam result on the graph so wrapStepForAdam can find it.
+      ;(graph as Graph & { __adam?: ReturnType<typeof appendAdam> }).__adam = adamResult
+    }
+  } else {
+    outputTensor = graph.tensors[graph.outputs[0]!]!
+  }
+
+  const plan = planBuffers(graph, paramGrads, adamWritebacks)
   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)
+  // exactOptionalPropertyTypes: only include sharedParams when defined.
+  const runtimeOpts: RuntimeOpts = sharedParams
+    ? { ...opts, sharedParams }
+    : { ...opts }
+  const runtime = withGrad
+    ? await createRuntime(plan, kernels, outputBufferId, runtimeOpts)
+    : await createForwardRuntime(plan, kernels, outputBufferId, runtimeOpts)
 
-  const sharedParams = opts.sharedParams
-  const uploadInitialParams = () => {
-    const out = buildInitialParamUploads(plan, materialized.initFns, sharedParams)
-    if (Object.keys(out).length > 0) runtime.uploadParams(out, { partial: !!sharedParams })
+  const ir: CompiledIR = { graph, paramGrads, loss: outputTensor, plan, kernels }
+  return { runtime: runtime as CompiledRuntime, materialized, plan, kernels, ir }
+}
+
+type Graph = ReturnType<typeof trace>
+
+function wrapStepForAdam(runtime: CompiledRuntime, adamCfg: AdamConfig, ir: CompiledIR): void {
+  const adamResult = (ir.graph as Graph & { __adam?: ReturnType<typeof appendAdam> }).__adam!
+  const { lrtInputName, decayShrinkInputName, config } = adamResult
+  let t = 0
+  const lrtBuf = new Float32Array(1)
+  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 },
+  ) => {
+    t++
+    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)
+  }) as CompiledRuntime['step']
+  runtime.step = wrappedStep
+  runtime.resetOptimizerState = () => {
+    t = 0
+    innerReset()
   }
+  void adamCfg
+}
 
-  // 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 })
+/** Build a Record<paramName, Float32Array> by running each param's init
+ *  function against its shape and uploading them to the runtime. Skips any
+ *  param covered by `sharedParams` (those are owned by a sibling compile). */
+function uploadInitialParams(
+  plan: BufferPlan,
+  initFns: Record<string, InitFn>,
+  runtime: CompiledRuntime | CompiledForward,
+  sharedParams: Map<string, GPUBuffer> | undefined,
+): void {
+  const out: Record<string, Float32Array> = {}
+  for (const [name, bufId] of plan.paramsByName) {
+    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(`compile: no init for param '${name}'`)
+    out[name] = initFn(size, shape)
+  }
+  if (Object.keys(out).length > 0) runtime.uploadParams(out, { partial: !!sharedParams })
 }

package/src/index.ts

@@ -36,7 +36,12 @@ 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, 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 { createRuntime, createForwardRuntime, Captures, type CompiledRuntime, type CompiledForward, type RuntimeOpts, type RunOptions, type StepResult, type RunResult } from './runtime.js'
+export {
+  compile, compileToIR, compileModule, compileForward,
+  type CompiledIR, type CompileModuleOptions, type CompileForwardOptions, type CompileForwardMethodOptions,
+  type CompiledModule, type CompiledForwardModule,
+  type InputDecl, type InputDecls, type InputsTensors, type ForwardFn,
+} from './compile.js'
 export { Module, materializeParams, type InitSpec, type ParamOptions, type MaterializedParams } from './module.js'
 export * as nn from './nn.js'

package/src/nn.ts

@@ -1,41 +1,44 @@
 // 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:
+// Each class declares its params and a `.fwd(x)` method that runs the forward
+// computation. Forward methods are pure tensorgrad ops — autograd traces
+// through them just like any other call.
 //
 //   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))
+//   const y = p.ffn.fwd(p.ln.fwd(x))
 
 import { Module } from './module.js'
 import type { Tensor } from './ir.js'
 import { add, matmul, sub, mul, div, sqrt, meanLast, sumLast, reshape, swapAxes, oneHot, logSoftmaxLast } from './ops.js'
 import { ShapeError } from './shape.js'
 import { captureSite } from './ir.js'
+import type { Captures } from './runtime.js'
 
 // ----------------------------------------------------------------------------
 // Linear: y = x @ W (+ b)
 // ----------------------------------------------------------------------------
 
+export interface LinearOptions {
+  /** Include a bias term (default true). */
+  bias?: boolean
+}
+
 export class Linear extends Module {
   W: Tensor
   b: Tensor | null
-  constructor(public readonly inDim: number, public readonly outDim: number, withBias = true) {
+  constructor(public readonly inDim: number, public readonly outDim: number, opts: LinearOptions = {}) {
     super()
     this.W = this.param([inDim, outDim])                      // randn, scale 0.02
-    this.b = withBias ? this.param([outDim], { init: 'zeros' }) : null
+    this.b = opts.bias === false ? null : this.param([outDim], { init: 'zeros' })
+  }
+  fwd(x: Tensor): Tensor {
+    const out = matmul(x, this.W)
+    return this.b ? add(out, this.b) : out
   }
-}
-
-export function linearFwd(p: Linear, x: Tensor): Tensor {
-  const out = matmul(x, p.W)
-  return p.b ? add(out, p.b) : out
 }
 
 // ----------------------------------------------------------------------------
@@ -50,14 +53,13 @@ export class LayerNorm extends Module {
     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)
+  fwd(x: Tensor): Tensor {
+    const m = meanLast(x)
+    const c = sub(x, m)
+    const v = meanLast(mul(c, c))
+    const stdev = sqrt(add(v, this.eps))
+    return add(mul(div(c, stdev), this.g), this.b)
+  }
 }
 
 // ----------------------------------------------------------------------------
@@ -97,26 +99,26 @@ export function mergeHeads(x: Tensor): Tensor {
   return reshape(swapped, [...lead, T, H * d])
 }
 
-/** Slice a flat capture readback of shape `[H, ..., ...]` into one
- *  Float32Array per head. The leading axis is treated as the head axis;
- *  pass the shape from `compiled.captureShapes[name]`. Result: `H` arrays,
- *  each holding the row-major data for that head (size = product of trailing
- *  axes). For B>1 graphs, prefix the result by the batch — this helper
- *  assumes the leading axis is heads, which matches how `splitHeads` lays
- *  out captures at B=1 (the typical capture-readback shape). */
-export function unsplitHeads(flat: Float32Array, shape: readonly number[]): Float32Array[] {
+/** Slice a captured tensor named `name` into one Float32Array per head, using
+ *  the static shape registered at compile time. The leading axis is treated as
+ *  heads (matching `splitHeads` layout at B=1); a leading singleton batch is
+ *  stripped if present so callers can pass capture names directly. Throws if
+ *  the capture isn't registered or wasn't read back this call. */
+export function unsplitHeads(captures: Captures, name: string): Float32Array[] {
+  const flat = captures.get(name)
+  const shape = captures.shapeOf(name)
   if (shape.length < 2) {
-    throw new Error(`unsplitHeads: shape needs >= 2 dims, got [${shape.join(', ')}]`)
+    throw new Error(`unsplitHeads: '${name}' shape needs >= 2 dims, got [${shape.join(', ')}]`)
   }
   // For inference graphs at B=1, captures have shape [1, H, ..., ...]. Strip
-  // the leading 1 if present so callers can pass captureShapes[name] directly.
+  // the leading 1 if present so the next axis is heads.
   const s = shape[0] === 1 ? shape.slice(1) : shape
   const H = s[0]!
   let stride = 1
   for (let i = 1; i < s.length; i++) stride *= s[i]!
   const expected = H * stride
   if (flat.length !== expected) {
-    throw new Error(`unsplitHeads: flat length ${flat.length} doesn't match shape product ${expected}`)
+    throw new Error(`unsplitHeads: '${name}' length ${flat.length} doesn't match shape product ${expected}`)
   }
   return Array.from({ length: H }, (_, h) => flat.slice(h * stride, (h + 1) * stride))
 }