tensorgrad 0.0.12 → 0.0.14

package/src/grad.ts

@@ -25,6 +25,7 @@ import {
   sumLast, where,
 } from './ops.js'
 import { traceInto } from './trace.js'
+import { shapesEqual } from './shape.js'
 
 // ============================================================================
 // Public API
@@ -121,11 +122,6 @@ function unbroadcast(cotan: Tensor, toShape: Shape): Tensor {
   return sumToShape(cotan, toShape)
 }
 
-function shapesEqual(a: Shape, b: Shape): boolean {
-  if (a.length !== b.length) return false
-  for (let i = 0; i < a.length; i++) if (a[i] !== b[i]) return false
-  return true
-}
 
 // ============================================================================
 // Transpose rules
@@ -435,12 +431,6 @@ function runTransposeRule(
 // Helpers
 // ============================================================================
 
-function identityPerm(rank: number): number[] {
-  const p: number[] = new Array(rank)
-  for (let i = 0; i < rank; i++) p[i] = i
-  return p
-}
-
 function invertPerm(perm: readonly number[]): number[] {
   const inv: number[] = new Array(perm.length)
   for (let i = 0; i < perm.length; i++) inv[perm[i]!] = i

package/src/runtime.ts

@@ -159,11 +159,14 @@ export async function createRuntime(
   // Param buffers may be supplied externally via opts.sharedParams; in that
   // case we reuse the provided GPUBuffer instead of allocating, and the
   // sibling compile that owns it is responsible for upload + lifetime.
+  // ownedBufferIds tracks which buffers we allocated ourselves (and so must
+  // destroy on .destroy()) vs which were handed in by a sibling compile.
   const buffers = new Map<number, GPUBuffer>()
+  const ownedBufferIds = new Set<number>()
   const sharedParams = opts.sharedParams
   for (const spec of plan.buffers) {
-    if (spec.kind === 'param' && sharedParams?.has(spec.name!)) {
-      const shared = sharedParams.get(spec.name!)!
+    const shared = spec.kind === 'param' ? sharedParams?.get(spec.name!) : undefined
+    if (shared) {
       if (shared.size !== spec.byteSize) {
         throw new Error(
           `sharedParams: size mismatch for '${spec.name}' — supplied ${shared.size} bytes, ` +
@@ -179,15 +182,9 @@ export async function createRuntime(
       label: spec.name ?? `t${spec.id}-${spec.kind}`,
     })
     buffers.set(spec.id, buf)
+    ownedBufferIds.add(spec.id)
     if (spec.kind === 'state') fillStateBuffer(spec, buf)
   }
-  // Track which params are externally owned — those are skipped on destroy().
-  const ownedBufferIds = new Set<number>()
-  for (const spec of plan.buffers) {
-    if (!(spec.kind === 'param' && sharedParams?.has(spec.name!))) {
-      ownedBufferIds.add(spec.id)
-    }
-  }
 
   // ---- Compile pipelines per kernel; cache by WGSL source -------------------
   // Push an error scope around each shader+pipeline creation so we can surface
@@ -253,20 +250,33 @@ export async function createRuntime(
   const outputSpec = plan.buffers[lossBufferId]!
   const outputReadback = device.createBuffer({ size: outputSpec.byteSize, usage: READBACK })
 
-  // ---- Capture readback staging buffers (lazy) ------------------------------
-  // Allocated on first `step({ withCaptures: true })` call and reused across
-  // subsequent calls. When the graph has no captures registered or when the
-  // caller never opts in, no extra GPU memory is allocated.
-  let captureStagings: Map<string, GPUBuffer> | null = null
-  function ensureCaptureStagings(): Map<string, GPUBuffer> {
-    if (captureStagings) return captureStagings
-    captureStagings = new Map()
+  // ---- Capture readback staging buffer (lazy, single concatenated) ---------
+  // One buffer for ALL captures, with each capture occupying a slice. Matters
+  // on mobile: each `mapAsync` round-trip on Android Chrome adds significant
+  // GPU-fence latency (~10–30 ms vs ~1 ms on desktop). With N captures, the
+  // per-call mobile cost is N × that latency on the main thread. Concatenating
+  // and reading back via one `mapAsync` collapses N stalls into one. Allocated
+  // on first `step({ withCaptures: true })` call.
+  type CaptureLayout = {
+    buffer: GPUBuffer
+    slices: { name: string; bufId: number; offset: number; byteSize: number }[]
+  }
+  let captureStaging: CaptureLayout | null = null
+  function ensureCaptureStaging(): CaptureLayout {
+    if (captureStaging) return captureStaging
+    let totalBytes = 0
+    const slices: CaptureLayout['slices'] = []
     for (const [name, bufId] of plan.capturesByName) {
       const spec = plan.buffers[bufId]!
-      const staging = device.createBuffer({ size: spec.byteSize, usage: READBACK, label: `cap-${name}` })
-      captureStagings.set(name, staging)
+      // copyBufferToBuffer offsets must be 4-aligned. Capture byteSizes are
+      // always shape-product × 4 (f32/i32/bool all 4 bytes), so cumulative
+      // offsets stay aligned.
+      slices.push({ name, bufId, offset: totalBytes, byteSize: spec.byteSize })
+      totalBytes += spec.byteSize
     }
-    return captureStagings
+    const buffer = device.createBuffer({ size: totalBytes, usage: READBACK, label: 'captures-staging' })
+    captureStaging = { buffer, slices }
+    return captureStaging
   }
 
   // ---- dispatch() — shared core for step() and run() -----------------------
@@ -339,14 +349,13 @@ export async function createRuntime(
       encoder.copyBufferToBuffer(buffers.get(wb.source)!, 0, buffers.get(wb.dest)!, 0, wb.bytes)
     }
     encoder.copyBufferToBuffer(buffers.get(lossBufferId)!, 0, outputReadback, 0, outputSpec.byteSize)
-    // Capture readbacks (only when opted in). Queued before submit so they
-    // observe the same kernel outputs as the main output.
-    let stagings: Map<string, GPUBuffer> | null = null
+    // Capture readbacks (only when opted in). All captures concatenate into
+    // a single staging buffer so we mapAsync once instead of N times.
+    let layout: CaptureLayout | null = null
     if (wantCaptures) {
-      stagings = ensureCaptureStagings()
-      for (const [name, bufId] of plan.capturesByName) {
-        const spec = plan.buffers[bufId]!
-        encoder.copyBufferToBuffer(buffers.get(bufId)!, 0, stagings.get(name)!, 0, spec.byteSize)
+      layout = ensureCaptureStaging()
+      for (const s of layout.slices) {
+        encoder.copyBufferToBuffer(buffers.get(s.bufId)!, 0, layout.buffer, s.offset, s.byteSize)
       }
     }
     queue.submit([encoder.finish()])
@@ -356,12 +365,15 @@ export async function createRuntime(
     outputReadback.unmap()
 
     const captures = new Map<string, Float32Array>()
-    if (wantCaptures) {
-      for (const [name, staging] of stagings!) {
-        await staging.mapAsync(GPUMapMode.READ)
-        captures.set(name, new Float32Array(staging.getMappedRange().slice(0)))
-        staging.unmap()
+    if (layout) {
+      await layout.buffer.mapAsync(GPUMapMode.READ)
+      const range = layout.buffer.getMappedRange()
+      for (const s of layout.slices) {
+        // Copy out (slice) before unmap — the underlying ArrayBuffer is
+        // detached when the buffer unmaps.
+        captures.set(s.name, new Float32Array(range, s.offset, s.byteSize / 4).slice())
       }
+      layout.buffer.unmap()
     }
     return { output, captures }
   }
@@ -483,7 +495,7 @@ export async function createRuntime(
       if (ownedBufferIds.has(id)) b.destroy()
     }
     outputReadback.destroy()
-    if (captureStagings) for (const b of captureStagings.values()) b.destroy()
+    if (captureStaging) captureStaging.buffer.destroy()
   }
 
   return {

package/src/trace.ts

@@ -84,20 +84,25 @@ export function traceInto<T>(g: Graph, fn: () => T): T {
 // Their .source on the Tensor points at that node so codegen knows where to
 // bind external data.
 
+// Param/tensor inputs share a namespace (a step() call passes both as keys in
+// the same dispatch object); state inputs have their own namespace.
+type NamedInputKind = 'param_input' | 'tensor_input' | 'state_input'
+function assertNameUnused(g: Graph, name: string, kinds: NamedInputKind[], label: string): void {
+  if (g.ops.some(op => kinds.includes(op.kind as NamedInputKind) && (op as { name?: string }).name === name)) {
+    throw new Error(`tensorgrad: ${label} name '${name}' already used in this trace`)
+  }
+}
+
 export function paramInput(name: string, shape: Shape, dtype: Dtype = 'f32'): Tensor {
   const g = currentGraph()
-  if (g.ops.some(op => (op.kind === 'param_input' || op.kind === 'tensor_input') && op.name === name)) {
-    throw new Error(`tensorgrad: input name '${name}' already used in this trace`)
-  }
+  assertNameUnused(g, name, ['param_input', 'tensor_input'], 'input')
   const site = captureSite('paramInput')
   return addOp(g, 'param_input', shape, dtype, site, { name } as any)
 }
 
 export function tensorInput(name: string, shape: Shape, dtype: Dtype = 'f32'): Tensor {
   const g = currentGraph()
-  if (g.ops.some(op => (op.kind === 'param_input' || op.kind === 'tensor_input') && op.name === name)) {
-    throw new Error(`tensorgrad: input name '${name}' already used in this trace`)
-  }
+  assertNameUnused(g, name, ['param_input', 'tensor_input'], 'input')
   const site = captureSite('tensorInput')
   return addOp(g, 'tensor_input', shape, dtype, site, { name } as any)
 }
@@ -106,9 +111,7 @@ export function tensorInput(name: string, shape: Shape, dtype: Dtype = 'f32'): T
 // and updated across step() calls via writebacks declared by the optimizer helper.
 export function stateInput(name: string, shape: Shape, dtype: Dtype = 'f32', initValue = 0): Tensor {
   const g = currentGraph()
-  if (g.ops.some(op => op.kind === 'state_input' && op.name === name)) {
-    throw new Error(`tensorgrad: state name '${name}' already used in this trace`)
-  }
+  assertNameUnused(g, name, ['state_input'], 'state')
   const site = captureSite('stateInput')
   return addOp(g, 'state_input', shape, dtype, site, { name, initValue } as any)
 }

package/dist/adam.d.ts

@@ -1,65 +0,0 @@
-import type { Tensor } from './ir.js';
-import type { Graph } from './ir.js';
-import type { WritebackDecl } from './buffers.js';
-export interface AdamConfig {
-    /** 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;
-    b2?: number;
-    eps?: number;
-    /** AdamW: decoupled weight decay coefficient. Default 0 (plain Adam).
-     *  When non-zero, every step shrinks each decayed param by a factor of
-     *  `1 - lr * weightDecay` before the gradient update. */
-    weightDecay?: number;
-    /** Filter deciding which params get weight decay. Only consulted when
-     *  weightDecay > 0. Default: decay every param. Override for the standard
-     *  transformer convention (decay weights/embeddings, skip biases + LN gains).
-     *  Example: `(name) => name.includes('.W') || name.endsWith('_emb')`. */
-    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;
-    /** 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;
-}
-/**
- * Append Adam update ops to `graph`. Must be called inside an active trace
- * context (or after a trace, since traceInto re-enters the graph).
- *
- * @param graph the graph (already containing forward + backward)
- * @param paramGrads param name -> gradient tensor (output of `appendGrad`)
- * @param paramTensors param name -> the param's leaf Tensor (the param_input).
- *                     Needed because the param_input lives in the graph but we
- *                     don't have a direct map by name in `Graph` — caller passes it.
- * @param config Adam hyperparameters. Set `weightDecay > 0` for AdamW; an
- *               optional `decayFilter` selects which params receive decay.
- */
-export declare function appendAdam(graph: Graph, paramGrads: Record<string, Tensor>, paramTensors: Record<string, Tensor>, config: AdamConfig, 
-/** Per-param decay flags from `materializeParams`. When supplied, overrides
- *  `config.decayFilter` for any name in the map; falls back to `decayFilter`
- *  for names not present (e.g., for low-level callers using `compile()`
- *  directly without a Module). */
-decayFlags?: Record<string, boolean>): AdamResult;
-//# sourceMappingURL=adam.d.ts.map

package/dist/adam.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"adam.d.ts","sourceRoot":"","sources":["../src/adam.ts"],"names":[],"mappings":"AA4BA,OAAO,KAAK,EAAE,MAAM,EAAE,MAAM,SAAS,CAAA;AACrC,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AACpC,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,cAAc,CAAA;AAIjD,MAAM,WAAW,UAAU;IACzB;;;0DAGsD;IACtD,EAAE,EAAE,MAAM,GAAG,CAAC,CAAC,IAAI,EAAE,MAAM,KAAK,MAAM,CAAC,CAAA;IACvC,EAAE,CAAC,EAAE,MAAM,CAAA;IACX,EAAE,CAAC,EAAE,MAAM,CAAA;IACX,GAAG,CAAC,EAAE,MAAM,CAAA;IACZ;;6DAEyD;IACzD,WAAW,CAAC,EAAE,MAAM,CAAA;IACpB;;;6EAGyE;IACzE,WAAW,CAAC,EAAE,CAAC,SAAS,EAAE,MAAM,KAAK,OAAO,CAAA;CAC7C;AAED,+EAA+E;AAC/E,MAAM,WAAW,kBAAkB;IACjC,EAAE,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,MAAM,CAAA;IAC5B,EAAE,EAAE,MAAM,CAAA;IACV,EAAE,EAAE,MAAM,CAAA;IACV,GAAG,EAAE,MAAM,CAAA;IACX,WAAW,EAAE,MAAM,CAAA;IACnB,WAAW,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,OAAO,CAAA;IACtC;kEAC8D;IAC9D,aAAa,EAAE,OAAO,CAAA;CACvB;AAED,MAAM,WAAW,UAAU;IACzB,kEAAkE;IAClE,UAAU,EAAE,aAAa,EAAE,CAAA;IAC3B;iFAC6E;IAC7E,YAAY,EAAE,MAAM,CAAA;IACpB;+EAC2E;IAC3E,oBAAoB,EAAE,MAAM,GAAG,IAAI,CAAA;IACnC,oFAAoF;IACpF,MAAM,EAAE,kBAAkB,CAAA;CAC3B;AAED;;;;;;;;;;;GAWG;AACH,wBAAgB,UAAU,CACxB,KAAK,EAAE,KAAK,EACZ,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAClC,YAAY,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EACpC,MAAM,EAAE,UAAU;AAClB;;;kCAGkC;AAClC,UAAU,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,GACnC,UAAU,CAyEZ"}

package/dist/buffers.d.ts

@@ -1,57 +0,0 @@
-import type { Graph, Tensor, Dtype, Shape } from './ir.js';
-export interface BufferSpec {
-    /** Matches tensor.id. */
-    id: number;
-    byteSize: number;
-    dtype: Dtype;
-    shape: Shape;
-    kind: 'param' | 'param_grad' | 'tensor_input' | 'state' | 'intermediate' | 'output';
-    /** External name for param/param_grad/tensor_input/state bindings. null otherwise. */
-    name: string | null;
-    /** For state buffers: the value to fill on initial allocation. 0 by default. */
-    initValue?: number;
-}
-/**
- * After step(), copy `source`'s buffer into `dest`'s buffer.
- * Used to write back updated optimizer state and updated parameters into
- * their persistent home buffers.
- */
-export interface Writeback {
-    source: number;
-    dest: number;
-    bytes: number;
-}
-export interface BufferPlan {
-    buffers: BufferSpec[];
-    /** Tensor id -> buffer id (currently 1:1 but kept opaque for future pooling). */
-    tensorToBuffer: Map<number, number>;
-    /** Easy lookup tables for the runtime. */
-    paramsByName: Map<string, number>;
-    inputsByName: Map<string, number>;
-    paramGradsByName: Map<string, number>;
-    statesByName: Map<string, number>;
-    capturesByName: Map<string, number>;
-    outputBufferIds: number[];
-    /** End-of-step writebacks (Adam updates for params, m, v, etc.) */
-    writebacks: Writeback[];
-}
-/**
- * Caller-supplied writeback declarations: "after each step, copy this Tensor's
- * buffer into the persistent home of this param/state."
- */
-export interface WritebackDecl {
-    /** The Tensor (output of some op) holding the new value to write back. */
-    source: Tensor;
-    /** Either a param name (writes to that param's home buffer) or a state name. */
-    destName: string;
-    destKind: 'param' | 'state';
-}
-/**
- * Build a BufferPlan from a graph + the param-grad map produced by appendGrad.
- * @param graph the full graph (forward + backward + any optimizer ops)
- * @param paramGrads map from param name -> the Tensor that holds its gradient
- * @param writebackDecls list of end-of-step writebacks (e.g. from appendAdam).
- *                       Empty when there's no optimizer in the graph.
- */
-export declare function planBuffers(graph: Graph, paramGrads: Record<string, Tensor>, writebackDecls?: WritebackDecl[]): BufferPlan;
-//# sourceMappingURL=buffers.d.ts.map

package/dist/buffers.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"buffers.d.ts","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAcA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAU,MAAM,SAAS,CAAA;AAElE,MAAM,WAAW,UAAU;IACzB,yBAAyB;IACzB,EAAE,EAAE,MAAM,CAAA;IACV,QAAQ,EAAE,MAAM,CAAA;IAChB,KAAK,EAAE,KAAK,CAAA;IACZ,KAAK,EAAE,KAAK,CAAA;IACZ,IAAI,EAAE,OAAO,GAAG,YAAY,GAAG,cAAc,GAAG,OAAO,GAAG,cAAc,GAAG,QAAQ,CAAA;IACnF,sFAAsF;IACtF,IAAI,EAAE,MAAM,GAAG,IAAI,CAAA;IACnB,gFAAgF;IAChF,SAAS,CAAC,EAAE,MAAM,CAAA;CACnB;AAED;;;;GAIG;AACH,MAAM,WAAW,SAAS;IACxB,MAAM,EAAE,MAAM,CAAA;IACd,IAAI,EAAE,MAAM,CAAA;IACZ,KAAK,EAAE,MAAM,CAAA;CACd;AAED,MAAM,WAAW,UAAU;IACzB,OAAO,EAAE,UAAU,EAAE,CAAA;IACrB,iFAAiF;IACjF,cAAc,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACnC,0CAA0C;IAC1C,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,gBAAgB,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACrC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,cAAc,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACnC,eAAe,EAAE,MAAM,EAAE,CAAA;IACzB,mEAAmE;IACnE,UAAU,EAAE,SAAS,EAAE,CAAA;CACxB;AAUD;;;GAGG;AACH,MAAM,WAAW,aAAa;IAC5B,0EAA0E;IAC1E,MAAM,EAAE,MAAM,CAAA;IACd,gFAAgF;IAChF,QAAQ,EAAE,MAAM,CAAA;IAChB,QAAQ,EAAE,OAAO,GAAG,OAAO,CAAA;CAC5B;AAED;;;;;;GAMG;AACH,wBAAgB,WAAW,CACzB,KAAK,EAAE,KAAK,EACZ,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAClC,cAAc,GAAE,aAAa,EAAO,GACnC,UAAU,CAmGZ"}

package/dist/capture.d.ts

@@ -1,3 +0,0 @@
-import type { Tensor } from './ir.js';
-export declare function capture<T extends Tensor>(name: string, t: T): T;
-//# sourceMappingURL=capture.d.ts.map

package/dist/capture.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"capture.d.ts","sourceRoot":"","sources":["../src/capture.ts"],"names":[],"mappings":"AAqBA,OAAO,KAAK,EAAE,MAAM,EAAE,MAAM,SAAS,CAAA;AAGrC,wBAAgB,OAAO,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,CAW/D"}

package/dist/codegen.d.ts

@@ -1,23 +0,0 @@
-import type { Graph, OpNode } from './ir.js';
-import type { BufferPlan } from './buffers.js';
-export interface KernelSpec {
-    /** Index into graph.ops. */
-    opIndex: number;
-    /** Op kind (for debugging / pipeline cache key). */
-    opKind: OpNode['kind'];
-    /** Generated WGSL source. Empty string for "logical" ops with no kernel. */
-    wgsl: string;
-    /**
-     * Buffer ids in binding-index order. The runtime creates a bind group with
-     * these in @binding(0..N) on @group(0). Inputs come first (read), output last
-     * (read_write).
-     */
-    bindings: number[];
-    /** Number of threads to dispatch (1-D). 0 means "skip" (e.g. reshape no-op). */
-    threads: number;
-    /** Workgroup size; usually WG_SIZE. */
-    workgroupSize: number;
-}
-/** Generate a KernelSpec per compute op in graph.ops (in dispatch order). */
-export declare function emitKernels(graph: Graph, plan: BufferPlan): KernelSpec[];
-//# sourceMappingURL=codegen.d.ts.map

package/dist/codegen.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"codegen.d.ts","sourceRoot":"","sources":["../src/codegen.ts"],"names":[],"mappings":"AAYA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAiB,MAAM,SAAS,CAAA;AAC3D,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,cAAc,CAAA;AAO9C,MAAM,WAAW,UAAU;IACzB,4BAA4B;IAC5B,OAAO,EAAE,MAAM,CAAA;IACf,oDAAoD;IACpD,MAAM,EAAE,MAAM,CAAC,MAAM,CAAC,CAAA;IACtB,4EAA4E;IAC5E,IAAI,EAAE,MAAM,CAAA;IACZ;;;;OAIG;IACH,QAAQ,EAAE,MAAM,EAAE,CAAA;IAClB,gFAAgF;IAChF,OAAO,EAAE,MAAM,CAAA;IACf,uCAAuC;IACvC,aAAa,EAAE,MAAM,CAAA;CACtB;AAMD,6EAA6E;AAC7E,wBAAgB,WAAW,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,UAAU,GAAG,UAAU,EAAE,CAQxE"}

package/dist/compile.d.ts

@@ -1,130 +0,0 @@
-import type { Tensor, Shape, Dtype } from './ir.js';
-import { type GradResult } from './grad.js';
-import { type AdamConfig } from './adam.js';
-import { type BufferPlan } from './buffers.js';
-import { type KernelSpec } from './codegen.js';
-import { type CompiledRuntime, type CompiledForward, type RuntimeOpts } from './runtime.js';
-import { Module } from './module.js';
-/** 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 {
-    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'];
-    loss: Tensor;
-    plan: BufferPlan;
-    kernels: KernelSpec[];
-}
-/** Trace + autograd + buffer-plan + codegen, without touching WebGPU. */
-export declare function compileToIR(traceFn: () => Tensor): CompiledIR;
-/** Full compile pipeline. Browser-only because it creates a GPUDevice. */
-export declare function compile(traceFn: () => Tensor, opts?: RuntimeOpts): Promise<CompiledRuntime & {
-    ir: CompiledIR;
-}>;
-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<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;
-}
-/**
- * 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 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. 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 declare function compileModule<M extends Module, I extends InputDecls = InputDecls>(modelFactory: () => M, forward: ForwardFn<M, I>, opts?: CompileModuleOptions<I>): Promise<CompiledModule<M>>;
-/**
- * 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`.
- *
- * **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.
- *
- * Initial param values are uploaded automatically for params *not* covered
- * by `sharedParams` (those are owned by the sibling compile).
- */
-export declare function compileForward<M extends Module, I extends InputDecls = InputDecls>(modelFactory: () => M, forward: ForwardFn<M, I>, opts?: CompileForwardOptions<I>): Promise<CompiledForwardModule>;
-//# sourceMappingURL=compile.d.ts.map

package/dist/compile.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"compile.d.ts","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAUA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAEnD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAuC,KAAK,eAAe,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,MAAM,cAAc,CAAA;AAChI,OAAO,EAAE,MAAM,EAAqB,MAAM,aAAa,CAAA;AAEvD;;sCAEsC;AACtC,MAAM,WAAW,SAAS;IACxB,KAAK,EAAE,KAAK,CAAA;IACZ,KAAK,CAAC,EAAE,KAAK,CAAA;CACd;AAED;;+CAE+C;AAC/C,MAAM,MAAM,UAAU,GAAG,MAAM,CAAC,MAAM,EAAE,SAAS,CAAC,CAAA;AAElD;;wDAEwD;AACxD,MAAM,MAAM,aAAa,CAAC,CAAC,SAAS,UAAU,IAAI;KAAG,CAAC,IAAI,MAAM,CAAC,GAAG,MAAM;CAAE,CAAA;AAE5E;;;;oCAIoC;AACpC,MAAM,MAAM,SAAS,CAAC,CAAC,SAAS,MAAM,EAAE,CAAC,SAAS,UAAU,GAAG,UAAU,IACvE,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,EAAE,aAAa,CAAC,CAAC,CAAC,KAAK,MAAM,CAAA;AAE5C,MAAM,WAAW,UAAU;IACzB,KAAK,EAAE,UAAU,CAAC,OAAO,CAAC,CAAA;IAC1B,UAAU,EAAE,UAAU,CAAC,YAAY,CAAC,CAAA;IACpC,IAAI,EAAE,MAAM,CAAA;IACZ,IAAI,EAAE,UAAU,CAAA;IAChB,OAAO,EAAE,UAAU,EAAE,CAAA;CACtB;AAED,yEAAyE;AACzE,wBAAgB,WAAW,CAAC,OAAO,EAAE,MAAM,MAAM,GAAG,UAAU,CAM7D;AAED,0EAA0E;AAC1E,wBAAsB,OAAO,CAAC,OAAO,EAAE,MAAM,MAAM,EAAE,IAAI,GAAE,WAAgB,GAAG,OAAO,CAAC,eAAe,GAAG;IAAE,EAAE,EAAE,UAAU,CAAA;CAAE,CAAC,CAK1H;AAMD,MAAM,WAAW,oBAAoB,CAAC,CAAC,SAAS,UAAU,GAAG,UAAU,CAAE,SAAQ,WAAW;IAC1F;;oEAEgE;IAChE,MAAM,CAAC,EAAE,CAAC,CAAA;IACV,iFAAiF;IACjF,IAAI,CAAC,EAAE,UAAU,CAAA;CAClB;AAED,MAAM,WAAW,qBAAqB,CAAC,CAAC,SAAS,UAAU,GAAG,UAAU,CAAE,SAAQ,WAAW;IAC3F,mEAAmE;IACnE,MAAM,CAAC,EAAE,CAAC,CAAA;CACX;AAED;;qDAEqD;AACrD,MAAM,WAAW,2BAA2B,CAAC,CAAC,SAAS,UAAU,GAAG,UAAU;IAC5E,MAAM,CAAC,EAAE,CAAC,CAAA;CACX;AAED;yDACyD;AACzD,MAAM,WAAW,cAAc,CAAC,CAAC,SAAS,MAAM,CAAE,SAAQ,eAAe;IACvE,EAAE,EAAE,UAAU,CAAA;IACd,6DAA6D;IAC7D,WAAW,EAAE,MAAM,CAAA;IACnB;2EACuE;IACvE,KAAK,IAAI,IAAI,CAAA;IACb;;;;kEAI8D;IAC9D,cAAc,CAAC,CAAC,SAAS,UAAU,EACjC,OAAO,EAAE,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,EACxB,IAAI,CAAC,EAAE,2BAA2B,CAAC,CAAC,CAAC,GACpC,OAAO,CAAC,qBAAqB,CAAC,CAAA;CAClC;AAED,yEAAyE;AACzE,MAAM,WAAW,qBAAsB,SAAQ,eAAe;IAC5D,EAAE,EAAE,UAAU,CAAA;IACd,6DAA6D;IAC7D,WAAW,EAAE,MAAM,CAAA;CACpB;AAED;;;;;;;;;;;;;;;;;;;;;;;;GAwBG;AACH,wBAAsB,aAAa,CAAC,CAAC,SAAS,MAAM,EAAE,CAAC,SAAS,UAAU,GAAG,UAAU,EACrF,YAAY,EAAE,MAAM,CAAC,EACrB,OAAO,EAAE,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,EACxB,IAAI,GAAE,oBAAoB,CAAC,CAAC,CAAM,GACjC,OAAO,CAAC,cAAc,CAAC,CAAC,CAAC,CAAC,CAyC5B;AAMD;;;;;;;;;;;;;;;;;;GAkBG;AACH,wBAAsB,cAAc,CAAC,CAAC,SAAS,MAAM,EAAE,CAAC,SAAS,UAAU,GAAG,UAAU,EACtF,YAAY,EAAE,MAAM,CAAC,EACrB,OAAO,EAAE,SAAS,CAAC,CAAC,EAAE,CAAC,CAAC,EACxB,IAAI,GAAE,qBAAqB,CAAC,CAAC,CAAM,GAClC,OAAO,CAAC,qBAAqB,CAAC,CAYhC"}

package/dist/grad.d.ts

@@ -1,8 +0,0 @@
-import type { Graph, Tensor } from './ir.js';
-export interface GradResult {
-    readonly graph: Graph;
-    readonly paramGrads: Record<string, Tensor>;
-    readonly loss: Tensor;
-}
-export declare function appendGrad(graph: Graph): GradResult;
-//# sourceMappingURL=grad.d.ts.map

package/dist/grad.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"grad.d.ts","sourceRoot":"","sources":["../src/grad.ts"],"names":[],"mappings":"AAiBA,OAAO,KAAK,EAAE,KAAK,EAAU,MAAM,EAAS,MAAM,SAAS,CAAA;AAe3D,MAAM,WAAW,UAAU;IAEzB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IAErB,QAAQ,CAAC,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IAE3C,QAAQ,CAAC,IAAI,EAAE,MAAM,CAAA;CACtB;AASD,wBAAgB,UAAU,CAAC,KAAK,EAAE,KAAK,GAAG,UAAU,CAmDnD"}

package/dist/index.d.ts.map

@@ -1 +0,0 @@
-{"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,EAAE,MAAM,EAEzB,OAAO,EAAE,SAAS,EAAE,QAAQ,EAE5B,MAAM,EAAE,aAAa,EAErB,MAAM,EAAE,MAAM,EAAE,SAAS,EAEzB,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,QAAQ,EAAE,KAAK,eAAe,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,SAAS,EAAE,MAAM,cAAc,CAAA;AAC5L,OAAO,EACL,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,cAAc,EACnD,KAAK,UAAU,EAAE,KAAK,oBAAoB,EAAE,KAAK,qBAAqB,EAAE,KAAK,2BAA2B,EACxG,KAAK,cAAc,EAAE,KAAK,qBAAqB,EAC/C,KAAK,SAAS,EAAE,KAAK,UAAU,EAAE,KAAK,aAAa,EAAE,KAAK,SAAS,GACpE,MAAM,cAAc,CAAA;AACrB,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"}