tensorgrad 0.0.14 → 0.0.16

package/src/worker.ts

@@ -0,0 +1,281 @@
+// Worker entry point. Holds the GPUDevice + CompiledRuntime for one or more
+// graphs and proxies main-thread requests via postMessage. See
+// specs/WorkerArchitecture.md for the rationale.
+//
+// Keep this file dependency-free of anything DOM-y: it bundles into a Blob
+// URL and runs in a Web Worker context where `window`/`document` don't
+// exist. WebGPU IS available in workers (Chrome 113+, Safari 17.4+).
+
+import { createRuntime, type CompiledRuntime, type RuntimeOpts } from './runtime.js'
+import { resolveLR, type LRSchedule } from './adam.js'
+import type { Req, Res, WireIR, WireAdamConfig, WireError } from './worker-protocol.js'
+import { wireError } from './worker-protocol.js'
+
+// ----------------------------------------------------------------------------
+// Per-graph state
+// ----------------------------------------------------------------------------
+
+interface GraphSlot {
+  runtime: CompiledRuntime
+  paramNames: readonly string[]
+  outputShape: number[]
+  kernelCount: number
+  captureShapes: Record<string, number[]>
+  /** Adam state for this graph, if it's a training graph. The wrapped step
+   *  uses these to populate the per-step lrt and decayShrink scalars. */
+  adam: AdamState | null
+}
+
+interface AdamState {
+  config: WireAdamConfig
+  t: number
+  lrtBuf: Float32Array
+  decayShrinkBuf: Float32Array | null
+}
+
+const graphs = new Map<number, GraphSlot>()
+
+// Worker holds one device shared across all graphs (sibling forward graphs
+// must share param GPUBuffers, which means sharing a device).
+let device: GPUDevice | null = null
+
+async function ensureDevice(): Promise<GPUDevice> {
+  if (device) return device
+  if (typeof navigator === 'undefined' || !navigator.gpu) {
+    throw new Error('tensorgrad worker: WebGPU not available in this environment')
+  }
+  const adapter = await navigator.gpu.requestAdapter()
+  if (!adapter) throw new Error('tensorgrad worker: no WebGPU adapter')
+  device = await adapter.requestDevice()
+  return device
+}
+
+// ----------------------------------------------------------------------------
+// Request handlers
+// ----------------------------------------------------------------------------
+
+async function handleCreateRuntime(payload: {
+  graphId: number
+  ir: WireIR
+  initialParams: Record<string, Float32Array>
+  adam: WireAdamConfig | null
+}): Promise<{ paramNames: string[]; outputShape: number[]; kernelCount: number; captureShapes: Record<string, number[]> }> {
+  const dev = await ensureDevice()
+  const { graph, plan, kernels } = payload.ir
+  const outputTensorId = graph.outputs[0]!
+  const outputBufferId = plan.tensorToBuffer.get(outputTensorId)!
+  const opts: RuntimeOpts = { device: dev }
+  const runtime = await createRuntime(plan, kernels, outputBufferId, opts)
+
+  // Upload initial params.
+  if (Object.keys(payload.initialParams).length > 0) {
+    runtime.uploadParams(payload.initialParams)
+  }
+
+  // Capture shape metadata for return.
+  const captureShapes: Record<string, number[]> = {}
+  for (const [name, bufId] of plan.capturesByName) {
+    captureShapes[name] = [...plan.buffers[bufId]!.shape]
+  }
+
+  const slot: GraphSlot = {
+    runtime,
+    paramNames: [...plan.paramsByName.keys()],
+    outputShape: [...runtime.outputShape],
+    kernelCount: kernels.filter(k => k.wgsl).length,
+    captureShapes,
+    adam: payload.adam ? createAdamState(payload.adam) : null,
+  }
+  graphs.set(payload.graphId, slot)
+
+  return {
+    paramNames: [...slot.paramNames],
+    outputShape: slot.outputShape,
+    kernelCount: slot.kernelCount,
+    captureShapes: slot.captureShapes,
+  }
+}
+
+async function handleCompileForward(payload: {
+  graphId: number
+  parentGraphId: number
+  ir: WireIR
+}): Promise<{ paramNames: string[]; outputShape: number[]; kernelCount: number; captureShapes: Record<string, number[]> }> {
+  const dev = await ensureDevice()
+  const parent = graphs.get(payload.parentGraphId)
+  if (!parent) throw new Error(`compileForward: parent graph ${payload.parentGraphId} not found`)
+
+  const { graph, plan, kernels } = payload.ir
+  const outputTensorId = graph.outputs[0]!
+  const outputBufferId = plan.tensorToBuffer.get(outputTensorId)!
+  const opts: RuntimeOpts = { device: dev, sharedParams: parent.runtime.params }
+  const runtime = await createRuntime(plan, kernels, outputBufferId, opts)
+  // No initial-param upload — sharedParams covers everything.
+
+  const captureShapes: Record<string, number[]> = {}
+  for (const [name, bufId] of plan.capturesByName) {
+    captureShapes[name] = [...plan.buffers[bufId]!.shape]
+  }
+
+  const slot: GraphSlot = {
+    runtime,
+    paramNames: [...plan.paramsByName.keys()],
+    outputShape: [...runtime.outputShape],
+    kernelCount: kernels.filter(k => k.wgsl).length,
+    captureShapes,
+    adam: null,
+  }
+  graphs.set(payload.graphId, slot)
+
+  return {
+    paramNames: [...slot.paramNames],
+    outputShape: slot.outputShape,
+    kernelCount: slot.kernelCount,
+    captureShapes: slot.captureShapes,
+  }
+}
+
+function createAdamState(cfg: WireAdamConfig): AdamState {
+  return {
+    config: cfg,
+    t: 0,
+    lrtBuf: new Float32Array(1),
+    decayShrinkBuf: cfg.decayShrinkInputName ? new Float32Array(1) : null,
+  }
+}
+
+/** Inject Adam's per-step lrt + decayShrink scalars into the inputs map.
+ *  Called before every step on a training graph. The buffers are reused
+ *  across steps to avoid allocation. */
+function injectAdamScalars(slot: GraphSlot, inputs: Record<string, Int32Array | Float32Array>): Record<string, Int32Array | Float32Array> {
+  const a = slot.adam
+  if (!a) return inputs
+  a.t++
+  const lrNow = resolveLR(a.config.lr as LRSchedule, a.t)
+  a.lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(a.config.b2, a.t)) / (1 - Math.pow(a.config.b1, a.t))
+  const merged: Record<string, Int32Array | Float32Array> = { ...inputs, [a.config.lrtInputName]: a.lrtBuf }
+  if (a.decayShrinkBuf && a.config.decayShrinkInputName) {
+    a.decayShrinkBuf[0] = 1 - lrNow * a.config.weightDecay
+    merged[a.config.decayShrinkInputName] = a.decayShrinkBuf
+  }
+  return merged
+}
+
+async function handleStep(payload: {
+  graphId: number
+  inputs: Record<string, Int32Array | Float32Array>
+  withCaptures: boolean
+}): Promise<{ loss: number; captures: Record<string, Float32Array> | null }> {
+  const slot = mustGet(payload.graphId)
+  const merged = injectAdamScalars(slot, payload.inputs)
+  if (payload.withCaptures) {
+    const r = await slot.runtime.step(merged, { withCaptures: true })
+    return { loss: r.loss, captures: capturesToRecord(r.captures, slot.captureShapes) }
+  }
+  const loss = await slot.runtime.step(merged)
+  return { loss, captures: null }
+}
+
+async function handleRun(payload: {
+  graphId: number
+  inputs: Record<string, Int32Array | Float32Array>
+  withCaptures: boolean
+}): Promise<{ output: Float32Array; captures: Record<string, Float32Array> | null }> {
+  const slot = mustGet(payload.graphId)
+  if (payload.withCaptures) {
+    const r = await slot.runtime.run(payload.inputs, { withCaptures: true })
+    return { output: r.output, captures: capturesToRecord(r.captures, slot.captureShapes) }
+  }
+  const output = await slot.runtime.run(payload.inputs)
+  return { output, captures: null }
+}
+
+/** Captures (a class instance with a private Map) → a plain Record so the
+ *  worker can transfer Float32Arrays back without serializing the class. */
+function capturesToRecord(
+  captures: { get(name: string): Float32Array; has(name: string): boolean; names(): string[] },
+  // captureShapes available but not used directly — capture names from
+  // shapes in case captures.names() is filtered (it isn't, but be safe).
+  shapes: Record<string, number[]>,
+): Record<string, Float32Array> {
+  const out: Record<string, Float32Array> = {}
+  for (const name of Object.keys(shapes)) {
+    if (captures.has(name)) out[name] = captures.get(name)
+  }
+  return out
+}
+
+function handleUploadParams(payload: {
+  graphId: number
+  params: Record<string, Float32Array>
+  partial: boolean
+}): void {
+  const slot = mustGet(payload.graphId)
+  slot.runtime.uploadParams(payload.params, { partial: payload.partial })
+}
+
+async function handleDownloadParams(payload: { graphId: number }): Promise<{ params: Record<string, Float32Array> }> {
+  const slot = mustGet(payload.graphId)
+  return { params: await slot.runtime.downloadParams() }
+}
+
+async function handleDownloadParamGrads(payload: { graphId: number }): Promise<{ params: Record<string, Float32Array> }> {
+  const slot = mustGet(payload.graphId)
+  return { params: await slot.runtime.downloadParamGrads() }
+}
+
+function handleResetOptimizer(payload: { graphId: number }): void {
+  const slot = mustGet(payload.graphId)
+  slot.runtime.resetOptimizerState()
+  if (slot.adam) slot.adam.t = 0
+}
+
+function handleDestroy(payload: { graphId: number }): void {
+  const slot = graphs.get(payload.graphId)
+  if (!slot) return
+  slot.runtime.destroy()
+  graphs.delete(payload.graphId)
+}
+
+function mustGet(graphId: number): GraphSlot {
+  const slot = graphs.get(graphId)
+  if (!slot) throw new Error(`tensorgrad worker: graph ${graphId} not found`)
+  return slot
+}
+
+// ----------------------------------------------------------------------------
+// Message dispatch
+// ----------------------------------------------------------------------------
+
+self.onmessage = async (ev: MessageEvent<Req>) => {
+  const req = ev.data
+  try {
+    let result: unknown
+    let transferList: ArrayBuffer[] = []
+    switch (req.kind) {
+      case 'createRuntime':     result = await handleCreateRuntime(req.payload); break
+      case 'compileForward':    result = await handleCompileForward(req.payload); break
+      case 'step':              result = await handleStep(req.payload); transferList = collectTransfers((result as any).captures); break
+      case 'run':               { const r = await handleRun(req.payload); result = r; transferList = [r.output.buffer as ArrayBuffer, ...collectTransfers(r.captures)]; break }
+      case 'uploadParams':      handleUploadParams(req.payload); result = null; break
+      case 'downloadParams':    { const r = await handleDownloadParams(req.payload); result = r; transferList = collectTransfers(r.params); break }
+      case 'downloadParamGrads':{ const r = await handleDownloadParamGrads(req.payload); result = r; transferList = collectTransfers(r.params); break }
+      case 'resetOptimizer':    handleResetOptimizer(req.payload); result = null; break
+      case 'destroy':           handleDestroy(req.payload); result = null; break
+      default: throw new Error(`unknown request kind: ${(req as { kind: string }).kind}`)
+    }
+    const reply: Res = { id: req.id, ok: true, result }
+    self.postMessage(reply, { transfer: transferList })
+  } catch (e) {
+    const error: WireError = wireError(e)
+    const reply: Res = { id: req.id, ok: false, error }
+    self.postMessage(reply)
+  }
+}
+
+function collectTransfers(rec: Record<string, Float32Array> | null | undefined): ArrayBuffer[] {
+  if (!rec) return []
+  const out: ArrayBuffer[] = []
+  for (const v of Object.values(rec)) out.push(v.buffer as ArrayBuffer)
+  return out
+}

package/dist/adam.js

@@ -1,111 +0,0 @@
-// Adam / AdamW optimizer, in-graph.
-//
-// `appendAdam` extends a graph that already has a forward pass + autograd-emitted
-// backward (i.e., has paramGrads from `appendGrad`) with the Adam update math.
-//
-// Per parameter P with gradient g:
-//   m_new = b1 * m + (1 - b1) * g
-//   v_new = b2 * v + (1 - b2) * g²
-//   p_new = decayShrink * p - lrt * m_new / (sqrt(v_new) + eps)
-//
-// `decayShrink = 1 - lr * weightDecay` when the param is being decayed
-// (Loshchilov & Hutter, "AdamW") and 1 otherwise — at which point the
-// multiply folds out and you're left with plain Adam. `lrt` is supplied
-// per-step from CPU and includes the bias-correction factor
-// `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
-// updates are aliased back to the param buffers.
-import { traceInto, stateInput, tensorInput } from './trace.js';
-import { adamUpdateM, adamUpdateV, adamUpdateP } from './ops.js';
-/**
- * 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 function appendAdam(graph, paramGrads, paramTensors, config, 
-/** 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) {
-    const lrIsScheduled = typeof config.lr === 'function';
-    const lrFn = lrIsScheduled
-        ? config.lr
-        : (() => config.lr);
-    const initialLr = lrFn(1);
-    const fullConfig = {
-        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 = [];
-    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 = null;
-    return traceInto(graph, () => {
-        const lrt = tensorInput(lrtInputName, [], 'f32');
-        // Up-front: which params receive weight decay? Per-param decayFlags (set
-        // by Module.param's options) wins; falls back to decayFilter for names
-        // not in the map. Empty when weightDecay = 0 so the rest of the function
-        // can just ask "is this name in the set?".
-        const decayedNames = new Set(fullConfig.weightDecay > 0
-            ? Object.keys(paramGrads).filter(name => (decayFlags && name in decayFlags) ? decayFlags[name] : fullConfig.decayFilter(name))
-            : []);
-        // We only need a runtime decayShrink scalar when lr varies per step AND
-        // at least one param is being decayed. Otherwise the value is constant
-        // and bakes into the kernel as a literal.
-        let decayShrinkScalar = null;
-        if (lrIsScheduled && decayedNames.size > 0) {
-            decayShrinkInputName = '_adam_decay_shrink';
-            decayShrinkScalar = tensorInput(decayShrinkInputName, [], 'f32');
-        }
-        for (const name of Object.keys(paramGrads)) {
-            const p = paramTensors[name];
-            const g = paramGrads[name];
-            if (!p)
-                throw new Error(`appendAdam: missing param tensor for '${name}'`);
-            if (!g)
-                throw new Error(`appendAdam: missing gradient for '${name}'`);
-            const mState = stateInput(`adam_m_${name}`, p.shape, 'f32', 0);
-            const vState = stateInput(`adam_v_${name}`, p.shape, 'f32', 0);
-            // Choose the decayShrink form per param:
-            //   - non-decayed params: literal 1 (kernel multiply folds out).
-            //   - decayed + scheduled lr: tensor input updated per step.
-            //   - decayed + static lr: literal `1 - lr * wd` baked at compile.
-            const decayShrink = !decayedNames.has(name) ? 1
-                : decayShrinkScalar !== null ? decayShrinkScalar
-                    : 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);
-            const newV = adamUpdateV(vState, g, fullConfig.b2);
-            const newP = adamUpdateP(p, newM, newV, lrt, fullConfig.eps, decayShrink);
-            writebacks.push({ source: newM, destName: `adam_m_${name}`, destKind: 'state' });
-            writebacks.push({ source: newV, destName: `adam_v_${name}`, destKind: 'state' });
-            writebacks.push({ source: newP, destName: name, destKind: 'param' });
-        }
-        return { writebacks, lrtInputName, decayShrinkInputName, config: fullConfig };
-    });
-}
-//# sourceMappingURL=adam.js.map

package/dist/adam.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"adam.js","sourceRoot":"","sources":["../src/adam.ts"],"names":[],"mappings":"AAAA,oCAAoC;AACpC,EAAE;AACF,kFAAkF;AAClF,+EAA+E;AAC/E,EAAE;AACF,mCAAmC;AACnC,kCAAkC;AAClC,mCAAmC;AACnC,gEAAgE;AAChE,EAAE;AACF,uEAAuE;AACvE,sEAAsE;AACtE,wEAAwE;AACxE,4DAA4D;AAC5D,wEAAwE;AACxE,4DAA4D;AAC5D,EAAE;AACF,2EAA2E;AAC3E,2EAA2E;AAC3E,0EAA0E;AAC1E,0EAA0E;AAC1E,6EAA6E;AAC7E,EAAE;AACF,wEAAwE;AACxE,0EAA0E;AAC1E,0EAA0E;AAC1E,iDAAiD;AAKjD,OAAO,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,MAAM,YAAY,CAAA;AAC/D,OAAO,EAAE,WAAW,EAAE,WAAW,EAAE,WAAW,EAAE,MAAM,UAAU,CAAA;AAgDhE;;;;;;;;;;;GAWG;AACH,MAAM,UAAU,UAAU,CACxB,KAAY,EACZ,UAAkC,EAClC,YAAoC,EACpC,MAAkB;AAClB;;;kCAGkC;AAClC,UAAoC;IAEpC,MAAM,aAAa,GAAG,OAAO,MAAM,CAAC,EAAE,KAAK,UAAU,CAAA;IACrD,MAAM,IAAI,GAAG,aAAa;QACxB,CAAC,CAAC,MAAM,CAAC,EAA8B;QACvC,CAAC,CAAC,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,EAAY,CAAC,CAAA;IAC/B,MAAM,SAAS,GAAG,IAAI,CAAC,CAAC,CAAC,CAAA;IACzB,MAAM,UAAU,GAAuB;QACrC,EAAE,EAAE,IAAI;QACR,EAAE,EAAE,MAAM,CAAC,EAAE,IAAI,GAAG;QACpB,EAAE,EAAE,MAAM,CAAC,EAAE,IAAI,KAAK;QACtB,GAAG,EAAE,MAAM,CAAC,GAAG,IAAI,IAAI;QACvB,WAAW,EAAE,MAAM,CAAC,WAAW,IAAI,CAAC;QACpC,WAAW,EAAE,MAAM,CAAC,WAAW,IAAI,CAAC,GAAG,EAAE,CAAC,IAAI,CAAC;QAC/C,aAAa;KACd,CAAA;IACD,MAAM,UAAU,GAAoB,EAAE,CAAA;IACtC,MAAM,YAAY,GAAG,WAAW,CAAA;IAChC,0EAA0E;IAC1E,iEAAiE;IACjE,IAAI,oBAAoB,GAAkB,IAAI,CAAA;IAE9C,OAAO,SAAS,CAAC,KAAK,EAAE,GAAG,EAAE;QAC3B,MAAM,GAAG,GAAG,WAAW,CAAC,YAAY,EAAE,EAAE,EAAE,KAAK,CAAC,CAAA;QAEhD,yEAAyE;QACzE,uEAAuE;QACvE,yEAAyE;QACzE,2CAA2C;QAC3C,MAAM,YAAY,GAAG,IAAI,GAAG,CAC1B,UAAU,CAAC,WAAW,GAAG,CAAC;YACxB,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,UAAU,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,EAAE,CACpC,CAAC,UAAU,IAAI,IAAI,IAAI,UAAU,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC,IAAI,CAAE,CAAC,CAAC,CAAC,UAAU,CAAC,WAAW,CAAC,IAAI,CAAC,CAAC;YAC1F,CAAC,CAAC,EAAE,CACP,CAAA;QAED,wEAAwE;QACxE,uEAAuE;QACvE,0CAA0C;QAC1C,IAAI,iBAAiB,GAAkB,IAAI,CAAA;QAC3C,IAAI,aAAa,IAAI,YAAY,CAAC,IAAI,GAAG,CAAC,EAAE,CAAC;YAC3C,oBAAoB,GAAG,oBAAoB,CAAA;YAC3C,iBAAiB,GAAG,WAAW,CAAC,oBAAoB,EAAE,EAAE,EAAE,KAAK,CAAC,CAAA;QAClE,CAAC;QAED,KAAK,MAAM,IAAI,IAAI,MAAM,CAAC,IAAI,CAAC,UAAU,CAAC,EAAE,CAAC;YAC3C,MAAM,CAAC,GAAG,YAAY,CAAC,IAAI,CAAC,CAAA;YAC5B,MAAM,CAAC,GAAG,UAAU,CAAC,IAAI,CAAC,CAAA;YAC1B,IAAI,CAAC,CAAC;gBAAE,MAAM,IAAI,KAAK,CAAC,yCAAyC,IAAI,GAAG,CAAC,CAAA;YACzE,IAAI,CAAC,CAAC;gBAAE,MAAM,IAAI,KAAK,CAAC,qCAAqC,IAAI,GAAG,CAAC,CAAA;YAErE,MAAM,MAAM,GAAG,UAAU,CAAC,UAAU,IAAI,EAAE,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC,CAAA;YAC9D,MAAM,MAAM,GAAG,UAAU,CAAC,UAAU,IAAI,EAAE,EAAE,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,CAAC,CAAC,CAAA;YAE9D,yCAAyC;YACzC,iEAAiE;YACjE,6DAA6D;YAC7D,mEAAmE;YACnE,MAAM,WAAW,GACf,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,CAAC;gBAC3B,CAAC,CAAC,iBAAiB,KAAK,IAAI,CAAC,CAAC,CAAC,iBAAiB;oBAChD,CAAC,CAAC,CAAC,GAAG,SAAS,GAAG,UAAU,CAAC,WAAW,CAAA;YAE1C,6EAA6E;YAC7E,MAAM,IAAI,GAAG,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,UAAU,CAAC,EAAE,CAAC,CAAA;YAClD,MAAM,IAAI,GAAG,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,UAAU,CAAC,EAAE,CAAC,CAAA;YAClD,MAAM,IAAI,GAAG,WAAW,CAAC,CAAC,EAAE,IAAI,EAAE,IAAI,EAAE,GAAG,EAAE,UAAU,CAAC,GAAG,EAAE,WAAW,CAAC,CAAA;YAEzE,UAAU,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,UAAU,IAAI,EAAE,EAAE,QAAQ,EAAE,OAAO,EAAE,CAAC,CAAA;YAChF,UAAU,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,UAAU,IAAI,EAAE,EAAE,QAAQ,EAAE,OAAO,EAAE,CAAC,CAAA;YAChF,UAAU,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,IAAI,EAAc,QAAQ,EAAE,OAAO,EAAE,CAAC,CAAA;QAClF,CAAC;QACD,OAAO,EAAE,UAAU,EAAE,YAAY,EAAE,oBAAoB,EAAE,MAAM,EAAE,UAAU,EAAE,CAAA;IAC/E,CAAC,CAAC,CAAA;AACJ,CAAC"}

package/dist/buffers.js

@@ -1,120 +0,0 @@
-// Buffer planning: walk a Graph and decide which GPU buffer each Tensor maps to.
-//
-// v1 strategy: one GPU buffer per IR Tensor. Static shapes mean every buffer's
-// size is known at compile time and lifetimes don't overlap between steps —
-// so no pooling needed. Total memory is the sum of every intermediate tensor.
-// For our transformer at B=256: ~30 MB of activations + grads. Easily fits.
-//
-// Categorization is what the runtime cares about:
-//   * param        — uploaded by user via uploadParams; persistent across steps
-//   * param_grad   — written each step by the backward pass; readable for inspection
-//   * tensor_input — uploaded each step (tokens, targets, masks)
-//   * intermediate — produced by an op; lifetime = within a single step
-//   * output       — special intermediate that should be made readable (loss)
-import { shapeSize } from './shape.js';
-const dtypeBytes = { f32: 4, i32: 4, bool: 4 };
-/**
- * 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 function planBuffers(graph, paramGrads, writebackDecls = []) {
-    const buffers = [];
-    const tensorToBuffer = new Map();
-    const paramsByName = new Map();
-    const inputsByName = new Map();
-    const paramGradsByName = new Map();
-    const statesByName = new Map();
-    // Build a quick reverse map: tensorId -> param name (for grads).
-    const gradTensorIdToName = new Map();
-    for (const [name, tensor] of Object.entries(paramGrads)) {
-        gradTensorIdToName.set(tensor.id, name);
-    }
-    // ...and tensorId -> param/input op (so we can name the buffer correctly).
-    const opByOutId = new Map();
-    for (const op of graph.ops)
-        opByOutId.set(op.out, op);
-    const outputSet = new Set(graph.outputs);
-    // Walk all tensors in id order. Categorize each.
-    for (const t of graph.tensors) {
-        const op = opByOutId.get(t.id);
-        let kind = 'intermediate';
-        let name = null;
-        let initValue;
-        if (op?.kind === 'param_input') {
-            kind = 'param';
-            name = op.name;
-        }
-        else if (op?.kind === 'tensor_input') {
-            kind = 'tensor_input';
-            name = op.name;
-        }
-        else if (op?.kind === 'state_input') {
-            kind = 'state';
-            name = op.name;
-            initValue = op.initValue;
-        }
-        else if (gradTensorIdToName.has(t.id)) {
-            kind = 'param_grad';
-            name = gradTensorIdToName.get(t.id);
-        }
-        else if (outputSet.has(t.id)) {
-            kind = 'output';
-        }
-        const spec = {
-            id: t.id,
-            byteSize: Math.max(4, shapeSize(t.shape) * dtypeBytes[t.dtype]),
-            dtype: t.dtype,
-            shape: t.shape,
-            kind,
-            name,
-            ...(initValue !== undefined ? { initValue } : {}),
-        };
-        buffers.push(spec);
-        tensorToBuffer.set(t.id, t.id); // 1:1 for v1
-        if (kind === 'param')
-            paramsByName.set(name, t.id);
-        if (kind === 'tensor_input')
-            inputsByName.set(name, t.id);
-        if (kind === 'param_grad')
-            paramGradsByName.set(name, t.id);
-        if (kind === 'state')
-            statesByName.set(name, t.id);
-    }
-    const outputBufferIds = graph.outputs.map(id => tensorToBuffer.get(id));
-    // Resolve writeback declarations to (source, dest) buffer-id pairs.
-    const writebacks = writebackDecls.map(decl => {
-        const sourceBufId = tensorToBuffer.get(decl.source.id);
-        if (sourceBufId === undefined) {
-            throw new Error(`planBuffers: writeback source tensor #${decl.source.id} not in graph`);
-        }
-        const destBufId = decl.destKind === 'param'
-            ? paramsByName.get(decl.destName)
-            : statesByName.get(decl.destName);
-        if (destBufId === undefined) {
-            throw new Error(`planBuffers: writeback dest ${decl.destKind}:'${decl.destName}' not found`);
-        }
-        const sourceSpec = buffers[sourceBufId];
-        const destSpec = buffers[destBufId];
-        if (sourceSpec.byteSize !== destSpec.byteSize) {
-            throw new Error(`planBuffers: writeback size mismatch for ${decl.destKind}:'${decl.destName}' ` +
-                `(source ${sourceSpec.byteSize} bytes vs dest ${destSpec.byteSize})`);
-        }
-        return { source: sourceBufId, dest: destBufId, bytes: sourceSpec.byteSize };
-    });
-    // 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();
-    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 };
-}
-//# sourceMappingURL=buffers.js.map

package/dist/buffers.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"buffers.js","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,+EAA+E;AAC/E,4EAA4E;AAC5E,8EAA8E;AAC9E,4EAA4E;AAC5E,EAAE;AACF,kDAAkD;AAClD,gFAAgF;AAChF,qFAAqF;AACrF,iEAAiE;AACjE,wEAAwE;AACxE,8EAA8E;AAG9E,OAAO,EAAE,SAAS,EAAE,MAAM,YAAY,CAAA;AAyCtC,MAAM,UAAU,GAA0B,EAAE,GAAG,EAAE,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,IAAI,EAAE,CAAC,EAAE,CAAA;AAcrE;;;;;;GAMG;AACH,MAAM,UAAU,WAAW,CACzB,KAAY,EACZ,UAAkC,EAClC,iBAAkC,EAAE;IAEpC,MAAM,OAAO,GAAiB,EAAE,CAAA;IAChC,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkB,CAAA;IAChD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,gBAAgB,GAAG,IAAI,GAAG,EAAkB,CAAA;IAClD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAE9C,iEAAiE;IACjE,MAAM,kBAAkB,GAAG,IAAI,GAAG,EAAkB,CAAA;IACpD,KAAK,MAAM,CAAC,IAAI,EAAE,MAAM,CAAC,IAAI,MAAM,CAAC,OAAO,CAAC,UAAU,CAAC,EAAE,CAAC;QACxD,kBAAkB,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACzC,CAAC;IACD,2EAA2E;IAC3E,MAAM,SAAS,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC3C,KAAK,MAAM,EAAE,IAAI,KAAK,CAAC,GAAG;QAAE,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,EAAE,EAAE,CAAC,CAAA;IAErD,MAAM,SAAS,GAAG,IAAI,GAAG,CAAC,KAAK,CAAC,OAAO,CAAC,CAAA;IAExC,iDAAiD;IACjD,KAAK,MAAM,CAAC,IAAI,KAAK,CAAC,OAAO,EAAE,CAAC;QAC9B,MAAM,EAAE,GAAG,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAA;QAC9B,IAAI,IAAI,GAAuB,cAAc,CAAA;QAC7C,IAAI,IAAI,GAAkB,IAAI,CAAA;QAC9B,IAAI,SAA6B,CAAA;QAEjC,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YAC/B,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,cAAc,EAAE,CAAC;YACvC,IAAI,GAAG,cAAc,CAAA;YACrB,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YACtC,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;YACd,SAAS,GAAG,EAAE,CAAC,SAAS,CAAA;QAC1B,CAAC;aAAM,IAAI,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YACxC,IAAI,GAAG,YAAY,CAAA;YACnB,IAAI,GAAG,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAE,CAAA;QACtC,CAAC;aAAM,IAAI,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YAC/B,IAAI,GAAG,QAAQ,CAAA;QACjB,CAAC;QAED,MAAM,IAAI,GAAe;YACvB,EAAE,EAAE,CAAC,CAAC,EAAE;YACR,QAAQ,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC,KAAK,CAAC,GAAG,UAAU,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;YAC/D,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,IAAI;YACJ,IAAI;YACJ,GAAG,CAAC,SAAS,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SAClD,CAAA;QACD,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;QAClB,cAAc,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA,CAAE,aAAa;QAE7C,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QACnD,IAAI,IAAI,KAAK,cAAc;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC1D,IAAI,IAAI,KAAK,YAAY;YAAE,gBAAgB,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC5D,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;IACrD,CAAC;IAED,MAAM,eAAe,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAE,CAAC,CAAA;IAExE,oEAAoE;IACpE,MAAM,UAAU,GAAgB,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE;QACxD,MAAM,WAAW,GAAG,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAC,CAAA;QACtD,IAAI,WAAW,KAAK,SAAS,EAAE,CAAC;YAC9B,MAAM,IAAI,KAAK,CAAC,yCAAyC,IAAI,CAAC,MAAM,CAAC,EAAE,eAAe,CAAC,CAAA;QACzF,CAAC;QACD,MAAM,SAAS,GAAG,IAAI,CAAC,QAAQ,KAAK,OAAO;YACzC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC;YACjC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAA;QACnC,IAAI,SAAS,KAAK,SAAS,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,+BAA+B,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,aAAa,CAAC,CAAA;QAC9F,CAAC;QACD,MAAM,UAAU,GAAG,OAAO,CAAC,WAAW,CAAE,CAAA;QACxC,MAAM,QAAQ,GAAG,OAAO,CAAC,SAAS,CAAE,CAAA;QACpC,IAAI,UAAU,CAAC,QAAQ,KAAK,QAAQ,CAAC,QAAQ,EAAE,CAAC;YAC9C,MAAM,IAAI,KAAK,CACb,4CAA4C,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,IAAI;gBAC/E,WAAW,UAAU,CAAC,QAAQ,kBAAkB,QAAQ,CAAC,QAAQ,GAAG,CACrE,CAAA;QACH,CAAC;QACD,OAAO,EAAE,MAAM,EAAE,WAAW,EAAE,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,UAAU,CAAC,QAAQ,EAAE,CAAA;IAC7E,CAAC,CAAC,CAAA;IAEF,qEAAqE;IACrE,0EAA0E;IAC1E,8DAA8D;IAC9D,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkB,CAAA;IAChD,KAAK,MAAM,CAAC,IAAI,EAAE,QAAQ,CAAC,IAAI,KAAK,CAAC,QAAQ,EAAE,CAAC;QAC9C,MAAM,KAAK,GAAG,cAAc,CAAC,GAAG,CAAC,QAAQ,CAAC,CAAA;QAC1C,IAAI,KAAK,KAAK,SAAS,EAAE,CAAC;YACxB,MAAM,IAAI,KAAK,CAAC,yBAAyB,IAAI,gCAAgC,QAAQ,EAAE,CAAC,CAAA;QAC1F,CAAC;QACD,cAAc,CAAC,GAAG,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;IACjC,CAAC;IAED,OAAO,EAAE,OAAO,EAAE,cAAc,EAAE,YAAY,EAAE,YAAY,EAAE,gBAAgB,EAAE,YAAY,EAAE,cAAc,EAAE,eAAe,EAAE,UAAU,EAAE,CAAA;AAC7I,CAAC"}

package/dist/capture.js

@@ -1,33 +0,0 @@
-// 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 { currentGraph, isCaptureEnabled } from './trace.js';
-export function capture(name, 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;
-}
-//# sourceMappingURL=capture.js.map

package/dist/capture.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"capture.js","sourceRoot":"","sources":["../src/capture.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,0CAA0C;AAC1C,EAAE;AACF,yCAAyC;AACzC,EAAE;AACF,kCAAkC;AAClC,yDAAyD;AACzD,0EAA0E;AAC1E,oCAAoC;AACpC,MAAM;AACN,EAAE;AACF,2EAA2E;AAC3E,6EAA6E;AAC7E,8EAA8E;AAC9E,8CAA8C;AAC9C,EAAE;AACF,yEAAyE;AACzE,yEAAyE;AACzE,iEAAiE;AAGjE,OAAO,EAAE,YAAY,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAA;AAE3D,MAAM,UAAU,OAAO,CAAmB,IAAY,EAAE,CAAI;IAC1D,IAAI,CAAC,gBAAgB,EAAE;QAAE,OAAO,CAAC,CAAA;IACjC,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,IAAI,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC;QACzB,MAAM,IAAI,KAAK,CACb,kBAAkB,IAAI,yCAAyC;YAC/D,4DAA4D,CAC7D,CAAA;IACH,CAAC;IACD,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;IAC1B,OAAO,CAAC,CAAA;AACV,CAAC"}