tensorgrad 0.0.16 → 0.0.17

package/README.md

@@ -1,119 +1,253 @@
-# tensorgrad
-
-A tiny TypeScript-native tensor library with autograd that compiles directly
-to WebGPU. Designed for training small models in the browser — without
-hand-writing WGSL kernels and without dragging in a 5 MB ML framework.
-
-```sh
-npm i tensorgrad
-```
-
-Roughly **3000 lines of zero-dependency TypeScript**, ~10 KB gzipped after
-build. Targets WebGPU only. Static shapes only. Forward + reverse-mode
-autograd; Adam optimizer; the whole training pipeline runs as compiled WGSL.
-
-## Quick example
-
-A 2-layer MLP fitting `y = sin(x)`:
-
-```ts
-import {
-  Module, compileModule,
-  add, mul, sub, sumLast, reshape, matmul, relu,
-  type Tensor,
-} from 'tensorgrad'
-
-class Linear extends Module {
-  W: Tensor; b: Tensor
-  constructor(public inDim: number, public outDim: number) {
-    super()
-    this.W = this.param([inDim, outDim])              // randn, scale 0.02
-    this.b = this.param([outDim], { init: 'zeros' })
-  }
-}
-
-class MLP extends Module {
-  l1 = new Linear(1, 64)
-  l2 = new Linear(64, 64)
-  l3 = new Linear(64, 1)
-}
-
-const linear = (p: Linear, x: Tensor) => add(matmul(x, p.W), p.b)
-
-function forward(m: MLP, x: Tensor): Tensor {
-  return linear(m.l3, relu(linear(m.l2, relu(linear(m.l1, x)))))
-}
-
-function loss(m: MLP, { x, y }: { x: Tensor; y: Tensor }): Tensor {
-  const diff = sub(forward(m, x), y)
-  return mul(sumLast(reshape(mul(diff, diff), [B])), 1 / B)
-}
-
-const B = 256
-const compiled = await compileModule(() => new MLP(), loss, {
-  adam: { lr: 0.005 },
-  inputs: {
-    x: { shape: [B, 1], dtype: 'f32' },
-    y: { shape: [B, 1], dtype: 'f32' },
-  },
-})
-
-// Initial params are uploaded automatically — no manual step needed.
-
-for (let step = 0; step < 1000; step++) {
-  const { x, y } = generateBatch()
-  const lossVal = await compiled.step({ x, y })
-  if (step % 100 === 0) console.log('step', step, 'loss', lossVal)
-}
-```
-
-That's the whole user-facing surface for this model: `Module` for parameter
-storage, plain functions for the forward pass, `compileModule` to JIT-compile
-to WGSL with autograd + Adam wired in. No decorators, no `tf.GradientTape`,
-no `register_pytree_node`.
-
-For a more involved example — a 3-layer transformer trained from scratch on
-2-digit addition — see the [`samples/`](./samples) workspace
-(`pnpm --filter samples dev`).
-
-## What this library is for
-
-Small browser-side ML where you want to *train* the model, not just run
-inference of a pretrained model. Educational artifacts, interactive
-demos, on-device personalization, "transformer from scratch in your browser"
-blog posts. Roughly the niche where the model is small enough to fit
-comfortably in a browser tab but where you still want autograd and a real
-optimizer.
-
-If you want to ship inference of a pretrained model, use
-[ONNX Runtime Web](https://github.com/microsoft/onnxruntime) or
-[transformers.js](https://github.com/xenova/transformers.js).
-If you need full JAX (vmap / pmap / dynamic shapes / multi-backend), use
-[jax-js](https://github.com/jax-js/jax).
-
-## Scope (deliberately small)
-
-The library only does what it does because of what it doesn't do. The
-load-bearing "out of scope" decisions are:
-
-- **WebGPU only** — no Wasm or WebGL fallback.
-- **Static shapes only** — every shape is fixed at compile time. This is
-  what lets us bake constants into the WGSL instead of carrying shape
-  uniforms.
-- **`grad` is the only transformation** — no `vmap`, `pmap`, `jvp`,
-  `custom_vjp`. Batch your data explicitly.
-- **`f32` only** — no dtype promotion, no mixed precision.
-- **Closed op set** — about 25 ops, listed in `SPEC.md`. Compositions of
-  those handle most needs (GELU, RMS norm, etc. are a few lines on top).
-- **Adam lives in the IR** — bias correction included; no CPU↔GPU
-  round-trip per step.
-
-## Status
-
-Alpha. Two real working models (a transformer training to <0.1 loss on
-addition, an MLP fitting `sin`). API may change before 1.0. Filing issues
-welcome.
-
-## License
-
-MIT
+# tensorgrad
+
+A tiny TypeScript-native tensor library with autograd that compiles to WebGPU.
+For training small models in the browser without hand-writing WGSL kernels and
+without dragging in a multi-megabyte ML framework.
+
+```sh
+npm i tensorgrad
+```
+
+Roughly 3000 lines of zero-dependency TypeScript. Static shapes, `f32`, Adam
+optimizer, ~25 ops, forward + reverse-mode autograd. Browser-only (uses
+WebGPU). All training/inference work runs in a library-internal Web Worker —
+every method on a compiled module returns a `Promise`.
+
+## Minimal example
+
+A 2-layer MLP fitting `y = sin(x)`:
+
+```ts
+import {
+  Module, compileModule, init,
+  add, mul, sub, sumLast, reshape, matmul, relu,
+  type Tensor,
+} from 'tensorgrad'
+
+const B = 256
+
+class Linear extends Module {
+  W: Tensor; b: Tensor
+  constructor(public inDim: number, public outDim: number) {
+    super()
+    this.W = this.param([inDim, outDim], { init: init.kaiming() })
+    this.b = this.param([outDim], { init: 'zeros' })
+  }
+}
+
+class MLP extends Module {
+  l1 = new Linear(1, 64)
+  l2 = new Linear(64, 64)
+  l3 = new Linear(64, 1)
+}
+
+const linearFwd = (p: Linear, x: Tensor) => add(matmul(x, p.W), p.b)
+
+function modelFwd(m: MLP, x: Tensor): Tensor {
+  return linearFwd(m.l3, relu(linearFwd(m.l2, relu(linearFwd(m.l1, x)))))
+}
+
+function lossFn(m: MLP, { x, y }: { x: Tensor; y: Tensor }): Tensor {
+  const diff = sub(modelFwd(m, x), y)
+  return mul(sumLast(reshape(mul(diff, diff), [B])), 1 / B)
+}
+
+const compiled = await compileModule(() => new MLP(), lossFn, {
+  adam: { lr: 0.005 },
+  inputs: {
+    x: { shape: [B, 1], dtype: 'f32' },
+    y: { shape: [B, 1], dtype: 'f32' },
+  },
+})
+
+for (let step = 0; step < 1000; step++) {
+  const { x, y } = generateBatch()
+  const lossVal = await compiled.step({ x, y })
+  if (step % 100 === 0) console.log('step', step, 'loss', lossVal)
+}
+```
+
+## Mental model
+
+- A `Module` subclass declares parameters via `this.param([shape], opts)` and
+  composes child modules as plain fields. The class is a tree of params.
+- A *forward function* takes the materialized module + a record of named
+  input tensors and returns a tensor (the loss for `compileModule`, or any
+  output for `compileForward`).
+- `compileModule(factory, forward, opts)` traces the forward, derives
+  gradients, wires Adam, plans buffers, generates WGSL, spawns a worker, and
+  returns a `CompiledModule`. The factory `() => new Model()` is invoked
+  once during compile; the model instance is consumed (its param sentinels
+  are mutated into `Tensor`s).
+- Every method on the compiled module is async. `await compiled.step(...)`
+  resolves with the loss after the worker's GPU work finishes.
+
+## Public API
+
+### Compile entry points
+
+```ts
+compileModule(factory, forward, { adam?, inputs? }): Promise<CompiledModule>
+compileForward(factory, forward, { inputs? }): Promise<CompiledForwardModule>
+```
+
+`compileForward` produces a forward-only graph in its own worker. To share
+params with an existing training graph, use the sibling method:
+
+```ts
+const train  = await compileModule(() => new Model(), lossFn, { ... })
+const infer  = await train.compileForward(predictFn, {
+  inputs: { tokens: { shape: [1, T], dtype: 'i32' } },
+})
+// infer runs in train's worker — every step's param updates are visible.
+```
+
+### CompiledModule methods (all `Promise`-returning)
+
+```ts
+compiled.step(inputs)                           // → loss: number
+compiled.step(inputs, { withCaptures: true })   // → { loss, captures }
+compiled.run(inputs)                            // → Float32Array
+compiled.run(inputs, { withCaptures: true })    // → { output, captures }
+compiled.uploadParams(record, { partial? })
+compiled.downloadParams()                       // → Record<name, Float32Array>
+compiled.downloadParamGrads()                   // → Record<name, Float32Array>
+compiled.reset()                                // re-init params + zero Adam state
+compiled.resetOptimizerState()
+compiled.compileForward(forward, { inputs? })   // sibling forward graph
+compiled.destroy()                              // tear down worker + GPU
+```
+
+`compiled.kernelCount`, `compiled.outputShape`, `compiled.paramNames`, and
+`compiled.ir` are sync properties for inspection.
+
+### Operators
+
+Imported from `'tensorgrad'`:
+
+- Element-wise: `add`, `sub`, `mul`, `div`, `sqrt`, `rsqrt`, `log`, `exp`, `relu`
+- Comparisons / select: `less`, `greater`, `where`
+- Reductions (last axis): `meanLast`, `sumLast`, `sumAll`
+- Shape: `reshape`, `transpose`, `swapAxes`
+- Linear algebra: `matmul`, `matmulBatched`
+- Indexing / casting: `oneHot`, `arange`, `embedding`
+- Slicing: `sliceLastRange`
+- Fused ML primitives: `softmaxCausalLast`, `logSoftmaxLast`, `whereCausal`
+
+`add`, `sub`, `mul`, `div` accept `(Tensor, Tensor)` or `(Tensor, number)`.
+
+### `nn` namespace
+
+```ts
+import { nn } from 'tensorgrad'
+
+nn.Linear(inDim, outDim, { bias? }) // .fwd(x)
+nn.LayerNorm(dim)                    // .fwd(x)
+nn.splitHeads(x, nHeads)             // [B, T, D] → [B, H, T, D/H]
+nn.mergeHeads(x)                     // inverse of splitHeads
+nn.unsplitHeads(captures, name)      // pull per-head slices off a capture
+nn.crossEntropyLast(logits, targets) // standard CE
+```
+
+Convention: leaf modules (`Linear`, `LayerNorm`) expose `.fwd(x)` for ergonomic
+chaining. Composite modules you write yourself are typically free functions
+taking `(p: ModuleType, x: Tensor)`.
+
+### LR schedules (`lr` namespace)
+
+```ts
+import { lr } from 'tensorgrad'
+
+adam: { lr: 0.005 }                                                 // constant
+adam: { lr: lr.linearDecay({ peak: 0.005, final: 0.0005, steps: 1500 }) }
+adam: { lr: lr.cosineDecay({ peak: 0.005, final: 0.0001, steps: 5000 }) }
+adam: { lr: lr.warmup({ peakLr: 0.001, warmupSteps: 200, after: lr.constant(0.001) }) }
+```
+
+LR schedules are serializable shapes, not closures (they cross the worker
+boundary). Use a `number` for constant LR, or one of the constructors above.
+
+### Param init (`init` namespace)
+
+```ts
+import { init } from 'tensorgrad'
+
+this.param([D, D], { init: init.kaiming() })           // gain=sqrt(2), fan_in=D
+this.param([D, D], { init: init.kaiming({ gain: 1 }) })
+this.param([D],    { init: 'zeros' })
+this.param([D],    { init: 'ones' })
+this.param([D, D], { init: init.randn({ scale: 0.02 }) })
+this.param([D],    { init: init.literal(myFloat32Array) })
+```
+
+Defaults: `'randn'` (std 0.02). AdamW weight decay defaults to `true` for
+randn/kaiming/literal init, `false` for zeros/ones — override per-param with
+`{ decay: true | false }`.
+
+### Captures (debugging / mech-interp)
+
+Wrap any tensor inside a forward to expose its activation post-run:
+
+```ts
+import { capture } from 'tensorgrad'
+
+const attn = capture(`attn.${i}`, softmaxCausalLast(scores))
+```
+
+```ts
+const { output, captures } = await compiled.run(inputs, { withCaptures: true })
+const attn0 = captures.get('attn.0')        // Float32Array
+captures.shapeOf('attn.0')                  // readonly number[]
+```
+
+Captures are zero-overhead unless `{ withCaptures: true }` is passed; they
+add a single batched mapAsync on the readback.
+
+## Constraints
+
+The library is small because of what it doesn't do. Plan accordingly:
+
+- **WebGPU only.** No Wasm, WebGL, or native fallback.
+- **Static shapes.** Every shape is fixed at compile time. Changing a batch
+  size means recompiling.
+- **`f32` only.** No mixed precision. Inputs may be `i32` for indices.
+- **One transformation: `grad`.** No `vmap`, `pmap`, `jvp`, `custom_vjp`.
+  Batch your data explicitly.
+- **Loss must be a scalar.** `compileModule`'s forward returns a rank-0 tensor.
+- **Closures don't cross the worker boundary.** LR schedules and inits are
+  serializable shapes, not functions. Anything per-step you write into a
+  user-defined optimizer (see *Extending* below) follows the same rule.
+- **One model per `compileModule` call.** Sibling forward graphs share params
+  via the method form; otherwise each compile spawns its own worker.
+
+## Extending
+
+The IR is open. Adam is built in only because it's the most common starting
+point — other optimizers, custom losses, or extra ops are user code following
+the same pattern as `appendAdam`:
+
+```ts
+import { appendAdam, appendGrad, compileToIR } from 'tensorgrad'
+```
+
+A custom optimizer is a function that takes the autograd output (graph +
+`paramGrads`) and the materialized param tensors, appends its update ops
+to the graph, and returns writeback declarations the buffer planner uses
+to wire each new value back into its persistent home. SGD, Lion, RMSProp
+all fit this shape; see `src/adam.ts` for the canonical example.
+
+The same applies to ops: anything missing from the built-in set can be
+expressed as a composition of existing ops (GELU, RMSNorm, etc. are a few
+lines), or — if you need a new primitive — added to the IR with a
+forward + backward + WGSL emit.
+
+## When not to use this
+
+- **Inference of pretrained models** → use ONNX Runtime Web or
+  transformers.js.
+- **Full JAX surface** (vmap, dynamic shapes, multi-backend) → use jax-js.
+- **Server-side training** → use PyTorch or JAX.
+
+## License
+
+MIT

package/dist/index.js

@@ -1805,7 +1805,7 @@ async function compileModule(modelFactory, forward, opts = {}) {
   const kernels = emitKernels(graph, plan);
   const ir = { graph, paramGrads, loss, plan, kernels };
   const initialParams = buildInitialParams(plan, materialized.initFns);
-  const proxy = new WorkerProxy('// src/runtime.ts\nvar Captures = class {\n  constructor(shapes, data) {\n    this.shapes = shapes;\n    this.data = data;\n  }\n  shapes;\n  data;\n  get(name) {\n    const d = this.data.get(name);\n    if (!d) {\n      const known = [...this.data.keys()].sort().join(", ");\n      const detail = known ? `Known this call: ${known}` : `(call run/step with { withCaptures: true } to populate)`;\n      throw new Error(`Captures.get: \'${name}\' not present. ${detail}`);\n    }\n    return d;\n  }\n  shapeOf(name) {\n    const s = this.shapes[name];\n    if (!s) {\n      const known = Object.keys(this.shapes).sort().join(", ") || "(none registered)";\n      throw new Error(`Captures.shapeOf: \'${name}\' not registered. Known: ${known}`);\n    }\n    return s;\n  }\n  has(name) {\n    return this.data.has(name);\n  }\n  names() {\n    return [...this.data.keys()].sort();\n  }\n};\nvar STORAGE_RW = 128 | 8 | 4;\nvar READBACK = 1 | 8;\nasync function createRuntime(plan, kernels, lossBufferId, opts = {}) {\n  const device2 = opts.device ?? await acquireDevice();\n  const queue = device2.queue;\n  const buffers = /* @__PURE__ */ new Map();\n  const ownedBufferIds = /* @__PURE__ */ new Set();\n  const sharedParams = opts.sharedParams;\n  for (const spec of plan.buffers) {\n    const shared = spec.kind === "param" ? sharedParams?.get(spec.name) : void 0;\n    if (shared) {\n      if (shared.size !== spec.byteSize) {\n        throw new Error(\n          `sharedParams: size mismatch for \'${spec.name}\' \\u2014 supplied ${shared.size} bytes, compiled graph expects ${spec.byteSize}.`\n        );\n      }\n      buffers.set(spec.id, shared);\n      continue;\n    }\n    const buf = device2.createBuffer({\n      size: spec.byteSize,\n      usage: STORAGE_RW,\n      label: spec.name ?? `t${spec.id}-${spec.kind}`\n    });\n    buffers.set(spec.id, buf);\n    ownedBufferIds.add(spec.id);\n    if (spec.kind === "state") fillStateBuffer(spec, buf);\n  }\n  const moduleCache = /* @__PURE__ */ new Map();\n  const pipelines = [];\n  const probes = [];\n  for (const k of kernels) {\n    if (!k.wgsl) {\n      pipelines.push(null);\n      continue;\n    }\n    let module = moduleCache.get(k.wgsl);\n    if (!module) {\n      module = device2.createShaderModule({ code: k.wgsl, label: k.opKind });\n      moduleCache.set(k.wgsl, module);\n    }\n    device2.pushErrorScope("validation");\n    const pipeline = device2.createComputePipeline({\n      layout: "auto",\n      compute: { module, entryPoint: "main" },\n      label: k.opKind\n    });\n    pipelines.push(pipeline);\n    probes.push(device2.popErrorScope().then((err) => err ? { k, module, err } : null));\n  }\n  const probeResults = await Promise.all(probes);\n  const failures = probeResults.filter((p) => p != null);\n  if (failures.length > 0) {\n    const reports = [];\n    for (const { k, module, err } of failures) {\n      const info = await module.getCompilationInfo();\n      const messages = info.messages.map((m) => `  L${m.lineNum}:${m.linePos} [${m.type}] ${m.message}`).join("\\n");\n      reports.push(\n        `[shader compile error] ${k.opKind} (op #${k.opIndex}): ${err.message}\n` + (messages || "  (no compilation messages)") + `\n--- WGSL ---\n${k.wgsl}\n-----------`\n      );\n    }\n    console.error(reports.join("\\n\\n"));\n    throw new Error(`tensorgrad: ${failures.length} shader(s) failed to compile (see console).`);\n  }\n  const bindGroups = kernels.map((k, i) => {\n    const pipeline = pipelines[i];\n    if (!pipeline) return null;\n    return device2.createBindGroup({\n      layout: pipeline.getBindGroupLayout(0),\n      entries: k.bindings.map((bufId, idx) => ({\n        binding: idx,\n        resource: { buffer: buffers.get(bufId) }\n      }))\n    });\n  });\n  const outputSpec = plan.buffers[lossBufferId];\n  const outputReadback = device2.createBuffer({ size: outputSpec.byteSize, usage: READBACK });\n  let captureStaging = null;\n  function ensureCaptureStaging() {\n    if (captureStaging) return captureStaging;\n    let totalBytes = 0;\n    const slices = [];\n    for (const [name, bufId] of plan.capturesByName) {\n      const spec = plan.buffers[bufId];\n      slices.push({ name, bufId, offset: totalBytes, byteSize: spec.byteSize });\n      totalBytes += spec.byteSize;\n    }\n    const buffer = device2.createBuffer({ size: totalBytes, usage: READBACK, label: "captures-staging" });\n    captureStaging = { buffer, slices };\n    return captureStaging;\n  }\n  let pending = Promise.resolve();\n  async function dispatch(inputs, opts2) {\n    const turn = pending.catch(() => {\n    }).then(() => dispatchUnsynchronized(inputs, opts2));\n    pending = turn;\n    return turn;\n  }\n  async function dispatchUnsynchronized(inputs, opts2) {\n    const wantCaptures = opts2.wantCaptures;\n    if (wantCaptures && plan.capturesByName.size === 0) {\n      throw new Error(\n        `withCaptures=true but no capture(...) calls were registered during the trace. Add capture(\'name\', tensor) inside your forward pass for the intermediates you want read back.`\n      );\n    }\n    for (const [name, bufId] of plan.inputsByName) {\n      const data = inputs[name];\n      if (!data) throw new Error(`tensorgrad: missing input \'${name}\'`);\n      const expectedBytes = plan.buffers[bufId].byteSize;\n      if (data.byteLength !== expectedBytes) {\n        throw new Error(`tensorgrad: input \'${name}\' has ${data.byteLength} bytes, expected ${expectedBytes}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-step" });\n    for (let i = 0; i < kernels.length; i++) {\n      const k = kernels[i];\n      if (!k.wgsl || k.threads === 0) continue;\n      const pipeline = pipelines[i];\n      const bindGroup = bindGroups[i];\n      const pass = encoder.beginComputePass({ label: k.opKind });\n      pass.setPipeline(pipeline);\n      pass.setBindGroup(0, bindGroup);\n      const wgCount = Math.max(1, Math.ceil(k.threads / k.workgroupSize));\n      const MAX_X = 65535;\n      const wgX = Math.min(wgCount, MAX_X);\n      const wgY = Math.ceil(wgCount / MAX_X);\n      pass.dispatchWorkgroups(wgX, wgY, 1);\n      pass.end();\n    }\n    for (const wb of plan.writebacks) {\n      encoder.copyBufferToBuffer(buffers.get(wb.source), 0, buffers.get(wb.dest), 0, wb.bytes);\n    }\n    encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, outputReadback, 0, outputSpec.byteSize);\n    let layout = null;\n    if (wantCaptures) {\n      layout = ensureCaptureStaging();\n      for (const s of layout.slices) {\n        encoder.copyBufferToBuffer(buffers.get(s.bufId), 0, layout.buffer, s.offset, s.byteSize);\n      }\n    }\n    queue.submit([encoder.finish()]);\n    if (!opts2.readback) return null;\n    await outputReadback.mapAsync(GPUMapMode.READ);\n    const output = new Float32Array(outputReadback.getMappedRange().slice(0));\n    outputReadback.unmap();\n    const captures = /* @__PURE__ */ new Map();\n    if (layout) {\n      await layout.buffer.mapAsync(GPUMapMode.READ);\n      const range = layout.buffer.getMappedRange();\n      for (const s of layout.slices) {\n        captures.set(s.name, new Float32Array(range, s.offset, s.byteSize / 4).slice());\n      }\n      layout.buffer.unmap();\n    }\n    return { output, captures };\n  }\n  async function step(inputs, opts2) {\n    if (opts2?.readLoss === false) {\n      await dispatch(inputs, { wantCaptures: false, readback: false });\n      return;\n    }\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { loss: r.output[0], captures: new Captures(captureShapes, r.captures) };\n    return r.output[0];\n  }\n  async function readLoss() {\n    const turn = pending.catch(() => {\n    }).then(async () => {\n      await outputReadback.mapAsync(GPUMapMode.READ);\n      const v = new Float32Array(outputReadback.getMappedRange())[0];\n      outputReadback.unmap();\n      return v;\n    });\n    pending = turn;\n    return turn;\n  }\n  async function run(inputs, opts2) {\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { output: r.output, captures: new Captures(captureShapes, r.captures) };\n    return r.output;\n  }\n  function uploadParams(params2, opts2) {\n    const partial = opts2?.partial ?? false;\n    for (const name of Object.keys(params2)) {\n      if (!plan.paramsByName.has(name)) {\n        throw new Error(\n          `uploadParams: unknown param \'${name}\'. Known: ${[...plan.paramsByName.keys()].sort().join(", ")}`\n        );\n      }\n    }\n    if (!partial) {\n      for (const name of plan.paramsByName.keys()) {\n        if (!(name in params2)) {\n          throw new Error(\n            `uploadParams: missing param \'${name}\'. Pass { partial: true } if you mean to update only some params.`\n          );\n        }\n      }\n    }\n    for (const [name, bufId] of plan.paramsByName) {\n      const data = params2[name];\n      if (!data) continue;\n      const expected = plan.buffers[bufId].byteSize / 4;\n      if (data.length !== expected) {\n        throw new Error(`uploadParams: \'${name}\' has ${data.length} elements, expected ${expected}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n  }\n  async function downloadFromMap(map) {\n    const stagings = [];\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-download" });\n    for (const [name, bufId] of map) {\n      const spec = plan.buffers[bufId];\n      const staging = device2.createBuffer({ size: spec.byteSize, usage: READBACK });\n      encoder.copyBufferToBuffer(buffers.get(bufId), 0, staging, 0, spec.byteSize);\n      stagings.push({ name, buf: staging, bytes: spec.byteSize });\n    }\n    queue.submit([encoder.finish()]);\n    const out = {};\n    for (const s of stagings) {\n      await s.buf.mapAsync(GPUMapMode.READ);\n      out[s.name] = new Float32Array(s.buf.getMappedRange().slice(0));\n      s.buf.unmap();\n      s.buf.destroy();\n    }\n    return out;\n  }\n  function fillStateBuffer(spec, target) {\n    const elements = spec.byteSize / 4;\n    const init = spec.dtype === "f32" ? new Float32Array(elements).fill(spec.initValue ?? 0) : new Int32Array(elements).fill(Math.trunc(spec.initValue ?? 0));\n    queue.writeBuffer(target, 0, init);\n  }\n  function resetOptimizerState() {\n    for (const spec of plan.buffers) {\n      if (spec.kind === "state") fillStateBuffer(spec, buffers.get(spec.id));\n    }\n  }\n  const params = /* @__PURE__ */ new Map();\n  for (const [name, bufId] of plan.paramsByName) {\n    params.set(name, buffers.get(bufId));\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const outputShape = [...plan.buffers[lossBufferId].shape];\n  const destroy = () => {\n    for (const [id, b] of buffers) {\n      if (ownedBufferIds.has(id)) b.destroy();\n    }\n    outputReadback.destroy();\n    if (captureStaging) captureStaging.buffer.destroy();\n  };\n  return {\n    device: device2,\n    params,\n    outputShape,\n    uploadParams,\n    downloadParams: () => downloadFromMap(plan.paramsByName),\n    downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),\n    step,\n    run,\n    readLoss,\n    resetOptimizerState,\n    destroy\n  };\n}\nasync function acquireDevice() {\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad: no WebGPU adapter");\n  return await adapter.requestDevice();\n}\n\n// src/adam.ts\nfunction resolveLR(schedule, step) {\n  if (typeof schedule === "number") return schedule;\n  switch (schedule.kind) {\n    case "constant":\n      return schedule.value;\n    case "linearDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.peak + (schedule.final - schedule.peak) * f;\n    }\n    case "cosineDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.final + 0.5 * (schedule.peak - schedule.final) * (1 + Math.cos(Math.PI * f));\n    }\n    case "warmup": {\n      if (step <= schedule.warmupSteps) return schedule.peakLr * (step / schedule.warmupSteps);\n      return resolveLR(schedule.after, step - schedule.warmupSteps);\n    }\n  }\n}\n\n// src/worker-protocol.ts\nfunction wireError(e) {\n  if (e instanceof Error) {\n    return { name: e.name, message: e.message, stack: e.stack ?? "" };\n  }\n  return { name: "Error", message: String(e), stack: "" };\n}\n\n// src/worker.ts\nvar graphs = /* @__PURE__ */ new Map();\nvar device = null;\nasync function ensureDevice() {\n  if (device) return device;\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad worker: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad worker: no WebGPU adapter");\n  device = await adapter.requestDevice();\n  return device;\n}\nasync function handleCreateRuntime(payload) {\n  const dev = await ensureDevice();\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  if (Object.keys(payload.initialParams).length > 0) {\n    runtime.uploadParams(payload.initialParams);\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: payload.adam ? createAdamState(payload.adam) : null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nasync function handleCompileForward(payload) {\n  const dev = await ensureDevice();\n  const parent = graphs.get(payload.parentGraphId);\n  if (!parent) throw new Error(`compileForward: parent graph ${payload.parentGraphId} not found`);\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev, sharedParams: parent.runtime.params };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nfunction createAdamState(cfg) {\n  return {\n    config: cfg,\n    t: 0,\n    lrtBuf: new Float32Array(1),\n    decayShrinkBuf: cfg.decayShrinkInputName ? new Float32Array(1) : null\n  };\n}\nfunction injectAdamScalars(slot, inputs) {\n  const a = slot.adam;\n  if (!a) return inputs;\n  a.t++;\n  const lrNow = resolveLR(a.config.lr, a.t);\n  a.lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(a.config.b2, a.t)) / (1 - Math.pow(a.config.b1, a.t));\n  const merged = { ...inputs, [a.config.lrtInputName]: a.lrtBuf };\n  if (a.decayShrinkBuf && a.config.decayShrinkInputName) {\n    a.decayShrinkBuf[0] = 1 - lrNow * a.config.weightDecay;\n    merged[a.config.decayShrinkInputName] = a.decayShrinkBuf;\n  }\n  return merged;\n}\nasync function handleStep(payload) {\n  const slot = mustGet(payload.graphId);\n  const merged = injectAdamScalars(slot, payload.inputs);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.step(merged, { withCaptures: true });\n    return { loss: r.loss, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const loss = await slot.runtime.step(merged);\n  return { loss, captures: null };\n}\nasync function handleRun(payload) {\n  const slot = mustGet(payload.graphId);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.run(payload.inputs, { withCaptures: true });\n    return { output: r.output, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const output = await slot.runtime.run(payload.inputs);\n  return { output, captures: null };\n}\nfunction capturesToRecord(captures, shapes) {\n  const out = {};\n  for (const name of Object.keys(shapes)) {\n    if (captures.has(name)) out[name] = captures.get(name);\n  }\n  return out;\n}\nfunction handleUploadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.uploadParams(payload.params, { partial: payload.partial });\n}\nasync function handleDownloadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParams() };\n}\nasync function handleDownloadParamGrads(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParamGrads() };\n}\nfunction handleResetOptimizer(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.resetOptimizerState();\n  if (slot.adam) slot.adam.t = 0;\n}\nfunction handleDestroy(payload) {\n  const slot = graphs.get(payload.graphId);\n  if (!slot) return;\n  slot.runtime.destroy();\n  graphs.delete(payload.graphId);\n}\nfunction mustGet(graphId) {\n  const slot = graphs.get(graphId);\n  if (!slot) throw new Error(`tensorgrad worker: graph ${graphId} not found`);\n  return slot;\n}\nself.onmessage = async (ev) => {\n  const req = ev.data;\n  try {\n    let result;\n    let transferList = [];\n    switch (req.kind) {\n      case "createRuntime":\n        result = await handleCreateRuntime(req.payload);\n        break;\n      case "compileForward":\n        result = await handleCompileForward(req.payload);\n        break;\n      case "step":\n        result = await handleStep(req.payload);\n        transferList = collectTransfers(result.captures);\n        break;\n      case "run": {\n        const r = await handleRun(req.payload);\n        result = r;\n        transferList = [r.output.buffer, ...collectTransfers(r.captures)];\n        break;\n      }\n      case "uploadParams":\n        handleUploadParams(req.payload);\n        result = null;\n        break;\n      case "downloadParams": {\n        const r = await handleDownloadParams(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "downloadParamGrads": {\n        const r = await handleDownloadParamGrads(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "resetOptimizer":\n        handleResetOptimizer(req.payload);\n        result = null;\n        break;\n      case "destroy":\n        handleDestroy(req.payload);\n        result = null;\n        break;\n      default:\n        throw new Error(`unknown request kind: ${req.kind}`);\n    }\n    const reply = { id: req.id, ok: true, result };\n    self.postMessage(reply, { transfer: transferList });\n  } catch (e) {\n    const error = wireError(e);\n    const reply = { id: req.id, ok: false, error };\n    self.postMessage(reply);\n  }\n};\nfunction collectTransfers(rec) {\n  if (!rec) return [];\n  const out = [];\n  for (const v of Object.values(rec)) out.push(v.buffer);\n  return out;\n}\n');
+  const proxy = new WorkerProxy('// src/runtime.ts\nvar Captures = class {\n  constructor(shapes, data) {\n    this.shapes = shapes;\n    this.data = data;\n  }\n  shapes;\n  data;\n  get(name) {\n    const d = this.data.get(name);\n    if (!d) {\n      const known = [...this.data.keys()].sort().join(", ");\n      const detail = known ? `Known this call: ${known}` : `(call run/step with { withCaptures: true } to populate)`;\n      throw new Error(`Captures.get: \'${name}\' not present. ${detail}`);\n    }\n    return d;\n  }\n  shapeOf(name) {\n    const s = this.shapes[name];\n    if (!s) {\n      const known = Object.keys(this.shapes).sort().join(", ") || "(none registered)";\n      throw new Error(`Captures.shapeOf: \'${name}\' not registered. Known: ${known}`);\n    }\n    return s;\n  }\n  has(name) {\n    return this.data.has(name);\n  }\n  names() {\n    return [...this.data.keys()].sort();\n  }\n};\nvar STORAGE_RW = 128 | 8 | 4;\nvar READBACK = 1 | 8;\nasync function createRuntime(plan, kernels, lossBufferId, opts = {}) {\n  const device2 = opts.device ?? await acquireDevice();\n  const queue = device2.queue;\n  const buffers = /* @__PURE__ */ new Map();\n  const ownedBufferIds = /* @__PURE__ */ new Set();\n  const sharedParams = opts.sharedParams;\n  for (const spec of plan.buffers) {\n    const shared = spec.kind === "param" ? sharedParams?.get(spec.name) : void 0;\n    if (shared) {\n      if (shared.size !== spec.byteSize) {\n        throw new Error(\n          `sharedParams: size mismatch for \'${spec.name}\' \\u2014 supplied ${shared.size} bytes, compiled graph expects ${spec.byteSize}.`\n        );\n      }\n      buffers.set(spec.id, shared);\n      continue;\n    }\n    const buf = device2.createBuffer({\n      size: spec.byteSize,\n      usage: STORAGE_RW,\n      label: spec.name ?? `t${spec.id}-${spec.kind}`\n    });\n    buffers.set(spec.id, buf);\n    ownedBufferIds.add(spec.id);\n    if (spec.kind === "state") fillStateBuffer(spec, buf);\n  }\n  const moduleCache = /* @__PURE__ */ new Map();\n  const pipelines = [];\n  const probes = [];\n  for (const k of kernels) {\n    if (!k.wgsl) {\n      pipelines.push(null);\n      continue;\n    }\n    let module = moduleCache.get(k.wgsl);\n    if (!module) {\n      module = device2.createShaderModule({ code: k.wgsl, label: k.opKind });\n      moduleCache.set(k.wgsl, module);\n    }\n    device2.pushErrorScope("validation");\n    const pipeline = device2.createComputePipeline({\n      layout: "auto",\n      compute: { module, entryPoint: "main" },\n      label: k.opKind\n    });\n    pipelines.push(pipeline);\n    probes.push(device2.popErrorScope().then((err) => err ? { k, module, err } : null));\n  }\n  const probeResults = await Promise.all(probes);\n  const failures = probeResults.filter((p) => p != null);\n  if (failures.length > 0) {\n    const reports = [];\n    for (const { k, module, err } of failures) {\n      const info = await module.getCompilationInfo();\n      const messages = info.messages.map((m) => `  L${m.lineNum}:${m.linePos} [${m.type}] ${m.message}`).join("\\n");\n      reports.push(\n        `[shader compile error] ${k.opKind} (op #${k.opIndex}): ${err.message}\n` + (messages || "  (no compilation messages)") + `\n--- WGSL ---\n${k.wgsl}\n-----------`\n      );\n    }\n    console.error(reports.join("\\n\\n"));\n    throw new Error(`tensorgrad: ${failures.length} shader(s) failed to compile (see console).`);\n  }\n  const bindGroups = kernels.map((k, i) => {\n    const pipeline = pipelines[i];\n    if (!pipeline) return null;\n    return device2.createBindGroup({\n      layout: pipeline.getBindGroupLayout(0),\n      entries: k.bindings.map((bufId, idx) => ({\n        binding: idx,\n        resource: { buffer: buffers.get(bufId) }\n      }))\n    });\n  });\n  const outputSpec = plan.buffers[lossBufferId];\n  const outputReadback = device2.createBuffer({ size: outputSpec.byteSize, usage: READBACK });\n  let captureStaging = null;\n  function ensureCaptureStaging() {\n    if (captureStaging) return captureStaging;\n    let totalBytes = 0;\n    const slices = [];\n    for (const [name, bufId] of plan.capturesByName) {\n      const spec = plan.buffers[bufId];\n      slices.push({ name, bufId, offset: totalBytes, byteSize: spec.byteSize });\n      totalBytes += spec.byteSize;\n    }\n    const buffer = device2.createBuffer({ size: totalBytes, usage: READBACK, label: "captures-staging" });\n    captureStaging = { buffer, slices };\n    return captureStaging;\n  }\n  let pending = Promise.resolve();\n  async function dispatch(inputs, opts2) {\n    const turn = pending.catch(() => {\n    }).then(() => dispatchUnsynchronized(inputs, opts2));\n    pending = turn;\n    return turn;\n  }\n  async function dispatchUnsynchronized(inputs, opts2) {\n    const wantCaptures = opts2.wantCaptures;\n    if (wantCaptures && plan.capturesByName.size === 0) {\n      throw new Error(\n        `withCaptures=true but no capture(...) calls were registered during the trace. Add capture(\'name\', tensor) inside your forward pass for the intermediates you want read back.`\n      );\n    }\n    for (const [name, bufId] of plan.inputsByName) {\n      const data = inputs[name];\n      if (!data) throw new Error(`tensorgrad: missing input \'${name}\'`);\n      const expectedBytes = plan.buffers[bufId].byteSize;\n      if (data.byteLength !== expectedBytes) {\n        throw new Error(`tensorgrad: input \'${name}\' has ${data.byteLength} bytes, expected ${expectedBytes}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n    const CHUNK_SIZE = 32;\n    let layout = null;\n    if (wantCaptures) {\n      layout = ensureCaptureStaging();\n    }\n    let kernelIdx = 0;\n    while (kernelIdx < kernels.length) {\n      const chunkEnd = Math.min(kernelIdx + CHUNK_SIZE, kernels.length);\n      const isLast = chunkEnd === kernels.length;\n      const encoder = device2.createCommandEncoder({\n        label: kernels.length > CHUNK_SIZE ? `tensorgrad-chunk-${kernelIdx}` : "tensorgrad-step"\n      });\n      for (let i = kernelIdx; i < chunkEnd; i++) {\n        const k = kernels[i];\n        if (!k.wgsl || k.threads === 0) continue;\n        const pipeline = pipelines[i];\n        const bindGroup = bindGroups[i];\n        const pass = encoder.beginComputePass({ label: k.opKind });\n        pass.setPipeline(pipeline);\n        pass.setBindGroup(0, bindGroup);\n        const wgCount = Math.max(1, Math.ceil(k.threads / k.workgroupSize));\n        const MAX_X = 65535;\n        const wgX = Math.min(wgCount, MAX_X);\n        const wgY = Math.ceil(wgCount / MAX_X);\n        pass.dispatchWorkgroups(wgX, wgY, 1);\n        pass.end();\n      }\n      if (isLast) {\n        for (const wb of plan.writebacks) {\n          encoder.copyBufferToBuffer(buffers.get(wb.source), 0, buffers.get(wb.dest), 0, wb.bytes);\n        }\n        encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, outputReadback, 0, outputSpec.byteSize);\n        if (layout) {\n          for (const s of layout.slices) {\n            encoder.copyBufferToBuffer(buffers.get(s.bufId), 0, layout.buffer, s.offset, s.byteSize);\n          }\n        }\n      }\n      queue.submit([encoder.finish()]);\n      if (!isLast) {\n        await queue.onSubmittedWorkDone();\n      }\n      kernelIdx = chunkEnd;\n    }\n    if (!opts2.readback) return null;\n    await outputReadback.mapAsync(GPUMapMode.READ);\n    const output = new Float32Array(outputReadback.getMappedRange().slice(0));\n    outputReadback.unmap();\n    const captures = /* @__PURE__ */ new Map();\n    if (layout) {\n      await layout.buffer.mapAsync(GPUMapMode.READ);\n      const range = layout.buffer.getMappedRange();\n      for (const s of layout.slices) {\n        captures.set(s.name, new Float32Array(range, s.offset, s.byteSize / 4).slice());\n      }\n      layout.buffer.unmap();\n    }\n    return { output, captures };\n  }\n  async function step(inputs, opts2) {\n    if (opts2?.readLoss === false) {\n      await dispatch(inputs, { wantCaptures: false, readback: false });\n      return;\n    }\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { loss: r.output[0], captures: new Captures(captureShapes, r.captures) };\n    return r.output[0];\n  }\n  async function readLoss() {\n    const turn = pending.catch(() => {\n    }).then(async () => {\n      await outputReadback.mapAsync(GPUMapMode.READ);\n      const v = new Float32Array(outputReadback.getMappedRange())[0];\n      outputReadback.unmap();\n      return v;\n    });\n    pending = turn;\n    return turn;\n  }\n  async function run(inputs, opts2) {\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { output: r.output, captures: new Captures(captureShapes, r.captures) };\n    return r.output;\n  }\n  function uploadParams(params2, opts2) {\n    const partial = opts2?.partial ?? false;\n    for (const name of Object.keys(params2)) {\n      if (!plan.paramsByName.has(name)) {\n        throw new Error(\n          `uploadParams: unknown param \'${name}\'. Known: ${[...plan.paramsByName.keys()].sort().join(", ")}`\n        );\n      }\n    }\n    if (!partial) {\n      for (const name of plan.paramsByName.keys()) {\n        if (!(name in params2)) {\n          throw new Error(\n            `uploadParams: missing param \'${name}\'. Pass { partial: true } if you mean to update only some params.`\n          );\n        }\n      }\n    }\n    for (const [name, bufId] of plan.paramsByName) {\n      const data = params2[name];\n      if (!data) continue;\n      const expected = plan.buffers[bufId].byteSize / 4;\n      if (data.length !== expected) {\n        throw new Error(`uploadParams: \'${name}\' has ${data.length} elements, expected ${expected}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n  }\n  async function downloadFromMap(map) {\n    const stagings = [];\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-download" });\n    for (const [name, bufId] of map) {\n      const spec = plan.buffers[bufId];\n      const staging = device2.createBuffer({ size: spec.byteSize, usage: READBACK });\n      encoder.copyBufferToBuffer(buffers.get(bufId), 0, staging, 0, spec.byteSize);\n      stagings.push({ name, buf: staging, bytes: spec.byteSize });\n    }\n    queue.submit([encoder.finish()]);\n    const out = {};\n    for (const s of stagings) {\n      await s.buf.mapAsync(GPUMapMode.READ);\n      out[s.name] = new Float32Array(s.buf.getMappedRange().slice(0));\n      s.buf.unmap();\n      s.buf.destroy();\n    }\n    return out;\n  }\n  function fillStateBuffer(spec, target) {\n    const elements = spec.byteSize / 4;\n    const init = spec.dtype === "f32" ? new Float32Array(elements).fill(spec.initValue ?? 0) : new Int32Array(elements).fill(Math.trunc(spec.initValue ?? 0));\n    queue.writeBuffer(target, 0, init);\n  }\n  function resetOptimizerState() {\n    for (const spec of plan.buffers) {\n      if (spec.kind === "state") fillStateBuffer(spec, buffers.get(spec.id));\n    }\n  }\n  const params = /* @__PURE__ */ new Map();\n  for (const [name, bufId] of plan.paramsByName) {\n    params.set(name, buffers.get(bufId));\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const outputShape = [...plan.buffers[lossBufferId].shape];\n  const destroy = () => {\n    for (const [id, b] of buffers) {\n      if (ownedBufferIds.has(id)) b.destroy();\n    }\n    outputReadback.destroy();\n    if (captureStaging) captureStaging.buffer.destroy();\n  };\n  return {\n    device: device2,\n    params,\n    outputShape,\n    uploadParams,\n    downloadParams: () => downloadFromMap(plan.paramsByName),\n    downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),\n    step,\n    run,\n    readLoss,\n    resetOptimizerState,\n    destroy\n  };\n}\nasync function acquireDevice() {\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad: no WebGPU adapter");\n  return await adapter.requestDevice();\n}\n\n// src/adam.ts\nfunction resolveLR(schedule, step) {\n  if (typeof schedule === "number") return schedule;\n  switch (schedule.kind) {\n    case "constant":\n      return schedule.value;\n    case "linearDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.peak + (schedule.final - schedule.peak) * f;\n    }\n    case "cosineDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.final + 0.5 * (schedule.peak - schedule.final) * (1 + Math.cos(Math.PI * f));\n    }\n    case "warmup": {\n      if (step <= schedule.warmupSteps) return schedule.peakLr * (step / schedule.warmupSteps);\n      return resolveLR(schedule.after, step - schedule.warmupSteps);\n    }\n  }\n}\n\n// src/worker-protocol.ts\nfunction wireError(e) {\n  if (e instanceof Error) {\n    return { name: e.name, message: e.message, stack: e.stack ?? "" };\n  }\n  return { name: "Error", message: String(e), stack: "" };\n}\n\n// src/worker.ts\nvar graphs = /* @__PURE__ */ new Map();\nvar device = null;\nasync function ensureDevice() {\n  if (device) return device;\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad worker: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad worker: no WebGPU adapter");\n  device = await adapter.requestDevice();\n  return device;\n}\nasync function handleCreateRuntime(payload) {\n  const dev = await ensureDevice();\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  if (Object.keys(payload.initialParams).length > 0) {\n    runtime.uploadParams(payload.initialParams);\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: payload.adam ? createAdamState(payload.adam) : null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nasync function handleCompileForward(payload) {\n  const dev = await ensureDevice();\n  const parent = graphs.get(payload.parentGraphId);\n  if (!parent) throw new Error(`compileForward: parent graph ${payload.parentGraphId} not found`);\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev, sharedParams: parent.runtime.params };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nfunction createAdamState(cfg) {\n  return {\n    config: cfg,\n    t: 0,\n    lrtBuf: new Float32Array(1),\n    decayShrinkBuf: cfg.decayShrinkInputName ? new Float32Array(1) : null\n  };\n}\nfunction injectAdamScalars(slot, inputs) {\n  const a = slot.adam;\n  if (!a) return inputs;\n  a.t++;\n  const lrNow = resolveLR(a.config.lr, a.t);\n  a.lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(a.config.b2, a.t)) / (1 - Math.pow(a.config.b1, a.t));\n  const merged = { ...inputs, [a.config.lrtInputName]: a.lrtBuf };\n  if (a.decayShrinkBuf && a.config.decayShrinkInputName) {\n    a.decayShrinkBuf[0] = 1 - lrNow * a.config.weightDecay;\n    merged[a.config.decayShrinkInputName] = a.decayShrinkBuf;\n  }\n  return merged;\n}\nasync function handleStep(payload) {\n  const slot = mustGet(payload.graphId);\n  const merged = injectAdamScalars(slot, payload.inputs);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.step(merged, { withCaptures: true });\n    return { loss: r.loss, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const loss = await slot.runtime.step(merged);\n  return { loss, captures: null };\n}\nasync function handleRun(payload) {\n  const slot = mustGet(payload.graphId);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.run(payload.inputs, { withCaptures: true });\n    return { output: r.output, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const output = await slot.runtime.run(payload.inputs);\n  return { output, captures: null };\n}\nfunction capturesToRecord(captures, shapes) {\n  const out = {};\n  for (const name of Object.keys(shapes)) {\n    if (captures.has(name)) out[name] = captures.get(name);\n  }\n  return out;\n}\nfunction handleUploadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.uploadParams(payload.params, { partial: payload.partial });\n}\nasync function handleDownloadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParams() };\n}\nasync function handleDownloadParamGrads(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParamGrads() };\n}\nfunction handleResetOptimizer(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.resetOptimizerState();\n  if (slot.adam) slot.adam.t = 0;\n}\nfunction handleDestroy(payload) {\n  const slot = graphs.get(payload.graphId);\n  if (!slot) return;\n  slot.runtime.destroy();\n  graphs.delete(payload.graphId);\n}\nfunction mustGet(graphId) {\n  const slot = graphs.get(graphId);\n  if (!slot) throw new Error(`tensorgrad worker: graph ${graphId} not found`);\n  return slot;\n}\nself.onmessage = async (ev) => {\n  const req = ev.data;\n  try {\n    let result;\n    let transferList = [];\n    switch (req.kind) {\n      case "createRuntime":\n        result = await handleCreateRuntime(req.payload);\n        break;\n      case "compileForward":\n        result = await handleCompileForward(req.payload);\n        break;\n      case "step":\n        result = await handleStep(req.payload);\n        transferList = collectTransfers(result.captures);\n        break;\n      case "run": {\n        const r = await handleRun(req.payload);\n        result = r;\n        transferList = [r.output.buffer, ...collectTransfers(r.captures)];\n        break;\n      }\n      case "uploadParams":\n        handleUploadParams(req.payload);\n        result = null;\n        break;\n      case "downloadParams": {\n        const r = await handleDownloadParams(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "downloadParamGrads": {\n        const r = await handleDownloadParamGrads(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "resetOptimizer":\n        handleResetOptimizer(req.payload);\n        result = null;\n        break;\n      case "destroy":\n        handleDestroy(req.payload);\n        result = null;\n        break;\n      default:\n        throw new Error(`unknown request kind: ${req.kind}`);\n    }\n    const reply = { id: req.id, ok: true, result };\n    self.postMessage(reply, { transfer: transferList });\n  } catch (e) {\n    const error = wireError(e);\n    const reply = { id: req.id, ok: false, error };\n    self.postMessage(reply);\n  }\n};\nfunction collectTransfers(rec) {\n  if (!rec) return [];\n  const out = [];\n  for (const v of Object.values(rec)) out.push(v.buffer);\n  return out;\n}\n');
   const wireIR = { graph, plan, kernels };
   const wireAdam = adamResult ? wireAdamConfig(adamResult) : null;
   const transfers = transferablesOfRecord(initialParams);
@@ -1843,7 +1843,7 @@ async function compileForward(modelFactory, forward, opts = {}) {
   const kernels = emitKernels(graph, plan);
   const ir = { graph, paramGrads: {}, loss: outputTensor, plan, kernels };
   const initialParams = buildInitialParams(plan, materialized.initFns);
-  const proxy = new WorkerProxy('// src/runtime.ts\nvar Captures = class {\n  constructor(shapes, data) {\n    this.shapes = shapes;\n    this.data = data;\n  }\n  shapes;\n  data;\n  get(name) {\n    const d = this.data.get(name);\n    if (!d) {\n      const known = [...this.data.keys()].sort().join(", ");\n      const detail = known ? `Known this call: ${known}` : `(call run/step with { withCaptures: true } to populate)`;\n      throw new Error(`Captures.get: \'${name}\' not present. ${detail}`);\n    }\n    return d;\n  }\n  shapeOf(name) {\n    const s = this.shapes[name];\n    if (!s) {\n      const known = Object.keys(this.shapes).sort().join(", ") || "(none registered)";\n      throw new Error(`Captures.shapeOf: \'${name}\' not registered. Known: ${known}`);\n    }\n    return s;\n  }\n  has(name) {\n    return this.data.has(name);\n  }\n  names() {\n    return [...this.data.keys()].sort();\n  }\n};\nvar STORAGE_RW = 128 | 8 | 4;\nvar READBACK = 1 | 8;\nasync function createRuntime(plan, kernels, lossBufferId, opts = {}) {\n  const device2 = opts.device ?? await acquireDevice();\n  const queue = device2.queue;\n  const buffers = /* @__PURE__ */ new Map();\n  const ownedBufferIds = /* @__PURE__ */ new Set();\n  const sharedParams = opts.sharedParams;\n  for (const spec of plan.buffers) {\n    const shared = spec.kind === "param" ? sharedParams?.get(spec.name) : void 0;\n    if (shared) {\n      if (shared.size !== spec.byteSize) {\n        throw new Error(\n          `sharedParams: size mismatch for \'${spec.name}\' \\u2014 supplied ${shared.size} bytes, compiled graph expects ${spec.byteSize}.`\n        );\n      }\n      buffers.set(spec.id, shared);\n      continue;\n    }\n    const buf = device2.createBuffer({\n      size: spec.byteSize,\n      usage: STORAGE_RW,\n      label: spec.name ?? `t${spec.id}-${spec.kind}`\n    });\n    buffers.set(spec.id, buf);\n    ownedBufferIds.add(spec.id);\n    if (spec.kind === "state") fillStateBuffer(spec, buf);\n  }\n  const moduleCache = /* @__PURE__ */ new Map();\n  const pipelines = [];\n  const probes = [];\n  for (const k of kernels) {\n    if (!k.wgsl) {\n      pipelines.push(null);\n      continue;\n    }\n    let module = moduleCache.get(k.wgsl);\n    if (!module) {\n      module = device2.createShaderModule({ code: k.wgsl, label: k.opKind });\n      moduleCache.set(k.wgsl, module);\n    }\n    device2.pushErrorScope("validation");\n    const pipeline = device2.createComputePipeline({\n      layout: "auto",\n      compute: { module, entryPoint: "main" },\n      label: k.opKind\n    });\n    pipelines.push(pipeline);\n    probes.push(device2.popErrorScope().then((err) => err ? { k, module, err } : null));\n  }\n  const probeResults = await Promise.all(probes);\n  const failures = probeResults.filter((p) => p != null);\n  if (failures.length > 0) {\n    const reports = [];\n    for (const { k, module, err } of failures) {\n      const info = await module.getCompilationInfo();\n      const messages = info.messages.map((m) => `  L${m.lineNum}:${m.linePos} [${m.type}] ${m.message}`).join("\\n");\n      reports.push(\n        `[shader compile error] ${k.opKind} (op #${k.opIndex}): ${err.message}\n` + (messages || "  (no compilation messages)") + `\n--- WGSL ---\n${k.wgsl}\n-----------`\n      );\n    }\n    console.error(reports.join("\\n\\n"));\n    throw new Error(`tensorgrad: ${failures.length} shader(s) failed to compile (see console).`);\n  }\n  const bindGroups = kernels.map((k, i) => {\n    const pipeline = pipelines[i];\n    if (!pipeline) return null;\n    return device2.createBindGroup({\n      layout: pipeline.getBindGroupLayout(0),\n      entries: k.bindings.map((bufId, idx) => ({\n        binding: idx,\n        resource: { buffer: buffers.get(bufId) }\n      }))\n    });\n  });\n  const outputSpec = plan.buffers[lossBufferId];\n  const outputReadback = device2.createBuffer({ size: outputSpec.byteSize, usage: READBACK });\n  let captureStaging = null;\n  function ensureCaptureStaging() {\n    if (captureStaging) return captureStaging;\n    let totalBytes = 0;\n    const slices = [];\n    for (const [name, bufId] of plan.capturesByName) {\n      const spec = plan.buffers[bufId];\n      slices.push({ name, bufId, offset: totalBytes, byteSize: spec.byteSize });\n      totalBytes += spec.byteSize;\n    }\n    const buffer = device2.createBuffer({ size: totalBytes, usage: READBACK, label: "captures-staging" });\n    captureStaging = { buffer, slices };\n    return captureStaging;\n  }\n  let pending = Promise.resolve();\n  async function dispatch(inputs, opts2) {\n    const turn = pending.catch(() => {\n    }).then(() => dispatchUnsynchronized(inputs, opts2));\n    pending = turn;\n    return turn;\n  }\n  async function dispatchUnsynchronized(inputs, opts2) {\n    const wantCaptures = opts2.wantCaptures;\n    if (wantCaptures && plan.capturesByName.size === 0) {\n      throw new Error(\n        `withCaptures=true but no capture(...) calls were registered during the trace. Add capture(\'name\', tensor) inside your forward pass for the intermediates you want read back.`\n      );\n    }\n    for (const [name, bufId] of plan.inputsByName) {\n      const data = inputs[name];\n      if (!data) throw new Error(`tensorgrad: missing input \'${name}\'`);\n      const expectedBytes = plan.buffers[bufId].byteSize;\n      if (data.byteLength !== expectedBytes) {\n        throw new Error(`tensorgrad: input \'${name}\' has ${data.byteLength} bytes, expected ${expectedBytes}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-step" });\n    for (let i = 0; i < kernels.length; i++) {\n      const k = kernels[i];\n      if (!k.wgsl || k.threads === 0) continue;\n      const pipeline = pipelines[i];\n      const bindGroup = bindGroups[i];\n      const pass = encoder.beginComputePass({ label: k.opKind });\n      pass.setPipeline(pipeline);\n      pass.setBindGroup(0, bindGroup);\n      const wgCount = Math.max(1, Math.ceil(k.threads / k.workgroupSize));\n      const MAX_X = 65535;\n      const wgX = Math.min(wgCount, MAX_X);\n      const wgY = Math.ceil(wgCount / MAX_X);\n      pass.dispatchWorkgroups(wgX, wgY, 1);\n      pass.end();\n    }\n    for (const wb of plan.writebacks) {\n      encoder.copyBufferToBuffer(buffers.get(wb.source), 0, buffers.get(wb.dest), 0, wb.bytes);\n    }\n    encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, outputReadback, 0, outputSpec.byteSize);\n    let layout = null;\n    if (wantCaptures) {\n      layout = ensureCaptureStaging();\n      for (const s of layout.slices) {\n        encoder.copyBufferToBuffer(buffers.get(s.bufId), 0, layout.buffer, s.offset, s.byteSize);\n      }\n    }\n    queue.submit([encoder.finish()]);\n    if (!opts2.readback) return null;\n    await outputReadback.mapAsync(GPUMapMode.READ);\n    const output = new Float32Array(outputReadback.getMappedRange().slice(0));\n    outputReadback.unmap();\n    const captures = /* @__PURE__ */ new Map();\n    if (layout) {\n      await layout.buffer.mapAsync(GPUMapMode.READ);\n      const range = layout.buffer.getMappedRange();\n      for (const s of layout.slices) {\n        captures.set(s.name, new Float32Array(range, s.offset, s.byteSize / 4).slice());\n      }\n      layout.buffer.unmap();\n    }\n    return { output, captures };\n  }\n  async function step(inputs, opts2) {\n    if (opts2?.readLoss === false) {\n      await dispatch(inputs, { wantCaptures: false, readback: false });\n      return;\n    }\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { loss: r.output[0], captures: new Captures(captureShapes, r.captures) };\n    return r.output[0];\n  }\n  async function readLoss() {\n    const turn = pending.catch(() => {\n    }).then(async () => {\n      await outputReadback.mapAsync(GPUMapMode.READ);\n      const v = new Float32Array(outputReadback.getMappedRange())[0];\n      outputReadback.unmap();\n      return v;\n    });\n    pending = turn;\n    return turn;\n  }\n  async function run(inputs, opts2) {\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { output: r.output, captures: new Captures(captureShapes, r.captures) };\n    return r.output;\n  }\n  function uploadParams(params2, opts2) {\n    const partial = opts2?.partial ?? false;\n    for (const name of Object.keys(params2)) {\n      if (!plan.paramsByName.has(name)) {\n        throw new Error(\n          `uploadParams: unknown param \'${name}\'. Known: ${[...plan.paramsByName.keys()].sort().join(", ")}`\n        );\n      }\n    }\n    if (!partial) {\n      for (const name of plan.paramsByName.keys()) {\n        if (!(name in params2)) {\n          throw new Error(\n            `uploadParams: missing param \'${name}\'. Pass { partial: true } if you mean to update only some params.`\n          );\n        }\n      }\n    }\n    for (const [name, bufId] of plan.paramsByName) {\n      const data = params2[name];\n      if (!data) continue;\n      const expected = plan.buffers[bufId].byteSize / 4;\n      if (data.length !== expected) {\n        throw new Error(`uploadParams: \'${name}\' has ${data.length} elements, expected ${expected}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n  }\n  async function downloadFromMap(map) {\n    const stagings = [];\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-download" });\n    for (const [name, bufId] of map) {\n      const spec = plan.buffers[bufId];\n      const staging = device2.createBuffer({ size: spec.byteSize, usage: READBACK });\n      encoder.copyBufferToBuffer(buffers.get(bufId), 0, staging, 0, spec.byteSize);\n      stagings.push({ name, buf: staging, bytes: spec.byteSize });\n    }\n    queue.submit([encoder.finish()]);\n    const out = {};\n    for (const s of stagings) {\n      await s.buf.mapAsync(GPUMapMode.READ);\n      out[s.name] = new Float32Array(s.buf.getMappedRange().slice(0));\n      s.buf.unmap();\n      s.buf.destroy();\n    }\n    return out;\n  }\n  function fillStateBuffer(spec, target) {\n    const elements = spec.byteSize / 4;\n    const init = spec.dtype === "f32" ? new Float32Array(elements).fill(spec.initValue ?? 0) : new Int32Array(elements).fill(Math.trunc(spec.initValue ?? 0));\n    queue.writeBuffer(target, 0, init);\n  }\n  function resetOptimizerState() {\n    for (const spec of plan.buffers) {\n      if (spec.kind === "state") fillStateBuffer(spec, buffers.get(spec.id));\n    }\n  }\n  const params = /* @__PURE__ */ new Map();\n  for (const [name, bufId] of plan.paramsByName) {\n    params.set(name, buffers.get(bufId));\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const outputShape = [...plan.buffers[lossBufferId].shape];\n  const destroy = () => {\n    for (const [id, b] of buffers) {\n      if (ownedBufferIds.has(id)) b.destroy();\n    }\n    outputReadback.destroy();\n    if (captureStaging) captureStaging.buffer.destroy();\n  };\n  return {\n    device: device2,\n    params,\n    outputShape,\n    uploadParams,\n    downloadParams: () => downloadFromMap(plan.paramsByName),\n    downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),\n    step,\n    run,\n    readLoss,\n    resetOptimizerState,\n    destroy\n  };\n}\nasync function acquireDevice() {\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad: no WebGPU adapter");\n  return await adapter.requestDevice();\n}\n\n// src/adam.ts\nfunction resolveLR(schedule, step) {\n  if (typeof schedule === "number") return schedule;\n  switch (schedule.kind) {\n    case "constant":\n      return schedule.value;\n    case "linearDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.peak + (schedule.final - schedule.peak) * f;\n    }\n    case "cosineDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.final + 0.5 * (schedule.peak - schedule.final) * (1 + Math.cos(Math.PI * f));\n    }\n    case "warmup": {\n      if (step <= schedule.warmupSteps) return schedule.peakLr * (step / schedule.warmupSteps);\n      return resolveLR(schedule.after, step - schedule.warmupSteps);\n    }\n  }\n}\n\n// src/worker-protocol.ts\nfunction wireError(e) {\n  if (e instanceof Error) {\n    return { name: e.name, message: e.message, stack: e.stack ?? "" };\n  }\n  return { name: "Error", message: String(e), stack: "" };\n}\n\n// src/worker.ts\nvar graphs = /* @__PURE__ */ new Map();\nvar device = null;\nasync function ensureDevice() {\n  if (device) return device;\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad worker: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad worker: no WebGPU adapter");\n  device = await adapter.requestDevice();\n  return device;\n}\nasync function handleCreateRuntime(payload) {\n  const dev = await ensureDevice();\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  if (Object.keys(payload.initialParams).length > 0) {\n    runtime.uploadParams(payload.initialParams);\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: payload.adam ? createAdamState(payload.adam) : null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nasync function handleCompileForward(payload) {\n  const dev = await ensureDevice();\n  const parent = graphs.get(payload.parentGraphId);\n  if (!parent) throw new Error(`compileForward: parent graph ${payload.parentGraphId} not found`);\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev, sharedParams: parent.runtime.params };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nfunction createAdamState(cfg) {\n  return {\n    config: cfg,\n    t: 0,\n    lrtBuf: new Float32Array(1),\n    decayShrinkBuf: cfg.decayShrinkInputName ? new Float32Array(1) : null\n  };\n}\nfunction injectAdamScalars(slot, inputs) {\n  const a = slot.adam;\n  if (!a) return inputs;\n  a.t++;\n  const lrNow = resolveLR(a.config.lr, a.t);\n  a.lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(a.config.b2, a.t)) / (1 - Math.pow(a.config.b1, a.t));\n  const merged = { ...inputs, [a.config.lrtInputName]: a.lrtBuf };\n  if (a.decayShrinkBuf && a.config.decayShrinkInputName) {\n    a.decayShrinkBuf[0] = 1 - lrNow * a.config.weightDecay;\n    merged[a.config.decayShrinkInputName] = a.decayShrinkBuf;\n  }\n  return merged;\n}\nasync function handleStep(payload) {\n  const slot = mustGet(payload.graphId);\n  const merged = injectAdamScalars(slot, payload.inputs);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.step(merged, { withCaptures: true });\n    return { loss: r.loss, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const loss = await slot.runtime.step(merged);\n  return { loss, captures: null };\n}\nasync function handleRun(payload) {\n  const slot = mustGet(payload.graphId);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.run(payload.inputs, { withCaptures: true });\n    return { output: r.output, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const output = await slot.runtime.run(payload.inputs);\n  return { output, captures: null };\n}\nfunction capturesToRecord(captures, shapes) {\n  const out = {};\n  for (const name of Object.keys(shapes)) {\n    if (captures.has(name)) out[name] = captures.get(name);\n  }\n  return out;\n}\nfunction handleUploadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.uploadParams(payload.params, { partial: payload.partial });\n}\nasync function handleDownloadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParams() };\n}\nasync function handleDownloadParamGrads(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParamGrads() };\n}\nfunction handleResetOptimizer(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.resetOptimizerState();\n  if (slot.adam) slot.adam.t = 0;\n}\nfunction handleDestroy(payload) {\n  const slot = graphs.get(payload.graphId);\n  if (!slot) return;\n  slot.runtime.destroy();\n  graphs.delete(payload.graphId);\n}\nfunction mustGet(graphId) {\n  const slot = graphs.get(graphId);\n  if (!slot) throw new Error(`tensorgrad worker: graph ${graphId} not found`);\n  return slot;\n}\nself.onmessage = async (ev) => {\n  const req = ev.data;\n  try {\n    let result;\n    let transferList = [];\n    switch (req.kind) {\n      case "createRuntime":\n        result = await handleCreateRuntime(req.payload);\n        break;\n      case "compileForward":\n        result = await handleCompileForward(req.payload);\n        break;\n      case "step":\n        result = await handleStep(req.payload);\n        transferList = collectTransfers(result.captures);\n        break;\n      case "run": {\n        const r = await handleRun(req.payload);\n        result = r;\n        transferList = [r.output.buffer, ...collectTransfers(r.captures)];\n        break;\n      }\n      case "uploadParams":\n        handleUploadParams(req.payload);\n        result = null;\n        break;\n      case "downloadParams": {\n        const r = await handleDownloadParams(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "downloadParamGrads": {\n        const r = await handleDownloadParamGrads(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "resetOptimizer":\n        handleResetOptimizer(req.payload);\n        result = null;\n        break;\n      case "destroy":\n        handleDestroy(req.payload);\n        result = null;\n        break;\n      default:\n        throw new Error(`unknown request kind: ${req.kind}`);\n    }\n    const reply = { id: req.id, ok: true, result };\n    self.postMessage(reply, { transfer: transferList });\n  } catch (e) {\n    const error = wireError(e);\n    const reply = { id: req.id, ok: false, error };\n    self.postMessage(reply);\n  }\n};\nfunction collectTransfers(rec) {\n  if (!rec) return [];\n  const out = [];\n  for (const v of Object.values(rec)) out.push(v.buffer);\n  return out;\n}\n');
+  const proxy = new WorkerProxy('// src/runtime.ts\nvar Captures = class {\n  constructor(shapes, data) {\n    this.shapes = shapes;\n    this.data = data;\n  }\n  shapes;\n  data;\n  get(name) {\n    const d = this.data.get(name);\n    if (!d) {\n      const known = [...this.data.keys()].sort().join(", ");\n      const detail = known ? `Known this call: ${known}` : `(call run/step with { withCaptures: true } to populate)`;\n      throw new Error(`Captures.get: \'${name}\' not present. ${detail}`);\n    }\n    return d;\n  }\n  shapeOf(name) {\n    const s = this.shapes[name];\n    if (!s) {\n      const known = Object.keys(this.shapes).sort().join(", ") || "(none registered)";\n      throw new Error(`Captures.shapeOf: \'${name}\' not registered. Known: ${known}`);\n    }\n    return s;\n  }\n  has(name) {\n    return this.data.has(name);\n  }\n  names() {\n    return [...this.data.keys()].sort();\n  }\n};\nvar STORAGE_RW = 128 | 8 | 4;\nvar READBACK = 1 | 8;\nasync function createRuntime(plan, kernels, lossBufferId, opts = {}) {\n  const device2 = opts.device ?? await acquireDevice();\n  const queue = device2.queue;\n  const buffers = /* @__PURE__ */ new Map();\n  const ownedBufferIds = /* @__PURE__ */ new Set();\n  const sharedParams = opts.sharedParams;\n  for (const spec of plan.buffers) {\n    const shared = spec.kind === "param" ? sharedParams?.get(spec.name) : void 0;\n    if (shared) {\n      if (shared.size !== spec.byteSize) {\n        throw new Error(\n          `sharedParams: size mismatch for \'${spec.name}\' \\u2014 supplied ${shared.size} bytes, compiled graph expects ${spec.byteSize}.`\n        );\n      }\n      buffers.set(spec.id, shared);\n      continue;\n    }\n    const buf = device2.createBuffer({\n      size: spec.byteSize,\n      usage: STORAGE_RW,\n      label: spec.name ?? `t${spec.id}-${spec.kind}`\n    });\n    buffers.set(spec.id, buf);\n    ownedBufferIds.add(spec.id);\n    if (spec.kind === "state") fillStateBuffer(spec, buf);\n  }\n  const moduleCache = /* @__PURE__ */ new Map();\n  const pipelines = [];\n  const probes = [];\n  for (const k of kernels) {\n    if (!k.wgsl) {\n      pipelines.push(null);\n      continue;\n    }\n    let module = moduleCache.get(k.wgsl);\n    if (!module) {\n      module = device2.createShaderModule({ code: k.wgsl, label: k.opKind });\n      moduleCache.set(k.wgsl, module);\n    }\n    device2.pushErrorScope("validation");\n    const pipeline = device2.createComputePipeline({\n      layout: "auto",\n      compute: { module, entryPoint: "main" },\n      label: k.opKind\n    });\n    pipelines.push(pipeline);\n    probes.push(device2.popErrorScope().then((err) => err ? { k, module, err } : null));\n  }\n  const probeResults = await Promise.all(probes);\n  const failures = probeResults.filter((p) => p != null);\n  if (failures.length > 0) {\n    const reports = [];\n    for (const { k, module, err } of failures) {\n      const info = await module.getCompilationInfo();\n      const messages = info.messages.map((m) => `  L${m.lineNum}:${m.linePos} [${m.type}] ${m.message}`).join("\\n");\n      reports.push(\n        `[shader compile error] ${k.opKind} (op #${k.opIndex}): ${err.message}\n` + (messages || "  (no compilation messages)") + `\n--- WGSL ---\n${k.wgsl}\n-----------`\n      );\n    }\n    console.error(reports.join("\\n\\n"));\n    throw new Error(`tensorgrad: ${failures.length} shader(s) failed to compile (see console).`);\n  }\n  const bindGroups = kernels.map((k, i) => {\n    const pipeline = pipelines[i];\n    if (!pipeline) return null;\n    return device2.createBindGroup({\n      layout: pipeline.getBindGroupLayout(0),\n      entries: k.bindings.map((bufId, idx) => ({\n        binding: idx,\n        resource: { buffer: buffers.get(bufId) }\n      }))\n    });\n  });\n  const outputSpec = plan.buffers[lossBufferId];\n  const outputReadback = device2.createBuffer({ size: outputSpec.byteSize, usage: READBACK });\n  let captureStaging = null;\n  function ensureCaptureStaging() {\n    if (captureStaging) return captureStaging;\n    let totalBytes = 0;\n    const slices = [];\n    for (const [name, bufId] of plan.capturesByName) {\n      const spec = plan.buffers[bufId];\n      slices.push({ name, bufId, offset: totalBytes, byteSize: spec.byteSize });\n      totalBytes += spec.byteSize;\n    }\n    const buffer = device2.createBuffer({ size: totalBytes, usage: READBACK, label: "captures-staging" });\n    captureStaging = { buffer, slices };\n    return captureStaging;\n  }\n  let pending = Promise.resolve();\n  async function dispatch(inputs, opts2) {\n    const turn = pending.catch(() => {\n    }).then(() => dispatchUnsynchronized(inputs, opts2));\n    pending = turn;\n    return turn;\n  }\n  async function dispatchUnsynchronized(inputs, opts2) {\n    const wantCaptures = opts2.wantCaptures;\n    if (wantCaptures && plan.capturesByName.size === 0) {\n      throw new Error(\n        `withCaptures=true but no capture(...) calls were registered during the trace. Add capture(\'name\', tensor) inside your forward pass for the intermediates you want read back.`\n      );\n    }\n    for (const [name, bufId] of plan.inputsByName) {\n      const data = inputs[name];\n      if (!data) throw new Error(`tensorgrad: missing input \'${name}\'`);\n      const expectedBytes = plan.buffers[bufId].byteSize;\n      if (data.byteLength !== expectedBytes) {\n        throw new Error(`tensorgrad: input \'${name}\' has ${data.byteLength} bytes, expected ${expectedBytes}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n    const CHUNK_SIZE = 32;\n    let layout = null;\n    if (wantCaptures) {\n      layout = ensureCaptureStaging();\n    }\n    let kernelIdx = 0;\n    while (kernelIdx < kernels.length) {\n      const chunkEnd = Math.min(kernelIdx + CHUNK_SIZE, kernels.length);\n      const isLast = chunkEnd === kernels.length;\n      const encoder = device2.createCommandEncoder({\n        label: kernels.length > CHUNK_SIZE ? `tensorgrad-chunk-${kernelIdx}` : "tensorgrad-step"\n      });\n      for (let i = kernelIdx; i < chunkEnd; i++) {\n        const k = kernels[i];\n        if (!k.wgsl || k.threads === 0) continue;\n        const pipeline = pipelines[i];\n        const bindGroup = bindGroups[i];\n        const pass = encoder.beginComputePass({ label: k.opKind });\n        pass.setPipeline(pipeline);\n        pass.setBindGroup(0, bindGroup);\n        const wgCount = Math.max(1, Math.ceil(k.threads / k.workgroupSize));\n        const MAX_X = 65535;\n        const wgX = Math.min(wgCount, MAX_X);\n        const wgY = Math.ceil(wgCount / MAX_X);\n        pass.dispatchWorkgroups(wgX, wgY, 1);\n        pass.end();\n      }\n      if (isLast) {\n        for (const wb of plan.writebacks) {\n          encoder.copyBufferToBuffer(buffers.get(wb.source), 0, buffers.get(wb.dest), 0, wb.bytes);\n        }\n        encoder.copyBufferToBuffer(buffers.get(lossBufferId), 0, outputReadback, 0, outputSpec.byteSize);\n        if (layout) {\n          for (const s of layout.slices) {\n            encoder.copyBufferToBuffer(buffers.get(s.bufId), 0, layout.buffer, s.offset, s.byteSize);\n          }\n        }\n      }\n      queue.submit([encoder.finish()]);\n      if (!isLast) {\n        await queue.onSubmittedWorkDone();\n      }\n      kernelIdx = chunkEnd;\n    }\n    if (!opts2.readback) return null;\n    await outputReadback.mapAsync(GPUMapMode.READ);\n    const output = new Float32Array(outputReadback.getMappedRange().slice(0));\n    outputReadback.unmap();\n    const captures = /* @__PURE__ */ new Map();\n    if (layout) {\n      await layout.buffer.mapAsync(GPUMapMode.READ);\n      const range = layout.buffer.getMappedRange();\n      for (const s of layout.slices) {\n        captures.set(s.name, new Float32Array(range, s.offset, s.byteSize / 4).slice());\n      }\n      layout.buffer.unmap();\n    }\n    return { output, captures };\n  }\n  async function step(inputs, opts2) {\n    if (opts2?.readLoss === false) {\n      await dispatch(inputs, { wantCaptures: false, readback: false });\n      return;\n    }\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { loss: r.output[0], captures: new Captures(captureShapes, r.captures) };\n    return r.output[0];\n  }\n  async function readLoss() {\n    const turn = pending.catch(() => {\n    }).then(async () => {\n      await outputReadback.mapAsync(GPUMapMode.READ);\n      const v = new Float32Array(outputReadback.getMappedRange())[0];\n      outputReadback.unmap();\n      return v;\n    });\n    pending = turn;\n    return turn;\n  }\n  async function run(inputs, opts2) {\n    const r = await dispatch(inputs, { wantCaptures: opts2?.withCaptures === true, readback: true });\n    if (opts2?.withCaptures) return { output: r.output, captures: new Captures(captureShapes, r.captures) };\n    return r.output;\n  }\n  function uploadParams(params2, opts2) {\n    const partial = opts2?.partial ?? false;\n    for (const name of Object.keys(params2)) {\n      if (!plan.paramsByName.has(name)) {\n        throw new Error(\n          `uploadParams: unknown param \'${name}\'. Known: ${[...plan.paramsByName.keys()].sort().join(", ")}`\n        );\n      }\n    }\n    if (!partial) {\n      for (const name of plan.paramsByName.keys()) {\n        if (!(name in params2)) {\n          throw new Error(\n            `uploadParams: missing param \'${name}\'. Pass { partial: true } if you mean to update only some params.`\n          );\n        }\n      }\n    }\n    for (const [name, bufId] of plan.paramsByName) {\n      const data = params2[name];\n      if (!data) continue;\n      const expected = plan.buffers[bufId].byteSize / 4;\n      if (data.length !== expected) {\n        throw new Error(`uploadParams: \'${name}\' has ${data.length} elements, expected ${expected}`);\n      }\n      queue.writeBuffer(buffers.get(bufId), 0, data);\n    }\n  }\n  async function downloadFromMap(map) {\n    const stagings = [];\n    const encoder = device2.createCommandEncoder({ label: "tensorgrad-download" });\n    for (const [name, bufId] of map) {\n      const spec = plan.buffers[bufId];\n      const staging = device2.createBuffer({ size: spec.byteSize, usage: READBACK });\n      encoder.copyBufferToBuffer(buffers.get(bufId), 0, staging, 0, spec.byteSize);\n      stagings.push({ name, buf: staging, bytes: spec.byteSize });\n    }\n    queue.submit([encoder.finish()]);\n    const out = {};\n    for (const s of stagings) {\n      await s.buf.mapAsync(GPUMapMode.READ);\n      out[s.name] = new Float32Array(s.buf.getMappedRange().slice(0));\n      s.buf.unmap();\n      s.buf.destroy();\n    }\n    return out;\n  }\n  function fillStateBuffer(spec, target) {\n    const elements = spec.byteSize / 4;\n    const init = spec.dtype === "f32" ? new Float32Array(elements).fill(spec.initValue ?? 0) : new Int32Array(elements).fill(Math.trunc(spec.initValue ?? 0));\n    queue.writeBuffer(target, 0, init);\n  }\n  function resetOptimizerState() {\n    for (const spec of plan.buffers) {\n      if (spec.kind === "state") fillStateBuffer(spec, buffers.get(spec.id));\n    }\n  }\n  const params = /* @__PURE__ */ new Map();\n  for (const [name, bufId] of plan.paramsByName) {\n    params.set(name, buffers.get(bufId));\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const outputShape = [...plan.buffers[lossBufferId].shape];\n  const destroy = () => {\n    for (const [id, b] of buffers) {\n      if (ownedBufferIds.has(id)) b.destroy();\n    }\n    outputReadback.destroy();\n    if (captureStaging) captureStaging.buffer.destroy();\n  };\n  return {\n    device: device2,\n    params,\n    outputShape,\n    uploadParams,\n    downloadParams: () => downloadFromMap(plan.paramsByName),\n    downloadParamGrads: () => downloadFromMap(plan.paramGradsByName),\n    step,\n    run,\n    readLoss,\n    resetOptimizerState,\n    destroy\n  };\n}\nasync function acquireDevice() {\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad: no WebGPU adapter");\n  return await adapter.requestDevice();\n}\n\n// src/adam.ts\nfunction resolveLR(schedule, step) {\n  if (typeof schedule === "number") return schedule;\n  switch (schedule.kind) {\n    case "constant":\n      return schedule.value;\n    case "linearDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.peak + (schedule.final - schedule.peak) * f;\n    }\n    case "cosineDecay": {\n      const f = Math.min(step / schedule.steps, 1);\n      return schedule.final + 0.5 * (schedule.peak - schedule.final) * (1 + Math.cos(Math.PI * f));\n    }\n    case "warmup": {\n      if (step <= schedule.warmupSteps) return schedule.peakLr * (step / schedule.warmupSteps);\n      return resolveLR(schedule.after, step - schedule.warmupSteps);\n    }\n  }\n}\n\n// src/worker-protocol.ts\nfunction wireError(e) {\n  if (e instanceof Error) {\n    return { name: e.name, message: e.message, stack: e.stack ?? "" };\n  }\n  return { name: "Error", message: String(e), stack: "" };\n}\n\n// src/worker.ts\nvar graphs = /* @__PURE__ */ new Map();\nvar device = null;\nasync function ensureDevice() {\n  if (device) return device;\n  if (typeof navigator === "undefined" || !navigator.gpu) {\n    throw new Error("tensorgrad worker: WebGPU not available in this environment");\n  }\n  const adapter = await navigator.gpu.requestAdapter();\n  if (!adapter) throw new Error("tensorgrad worker: no WebGPU adapter");\n  device = await adapter.requestDevice();\n  return device;\n}\nasync function handleCreateRuntime(payload) {\n  const dev = await ensureDevice();\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  if (Object.keys(payload.initialParams).length > 0) {\n    runtime.uploadParams(payload.initialParams);\n  }\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: payload.adam ? createAdamState(payload.adam) : null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nasync function handleCompileForward(payload) {\n  const dev = await ensureDevice();\n  const parent = graphs.get(payload.parentGraphId);\n  if (!parent) throw new Error(`compileForward: parent graph ${payload.parentGraphId} not found`);\n  const { graph, plan, kernels } = payload.ir;\n  const outputTensorId = graph.outputs[0];\n  const outputBufferId = plan.tensorToBuffer.get(outputTensorId);\n  const opts = { device: dev, sharedParams: parent.runtime.params };\n  const runtime = await createRuntime(plan, kernels, outputBufferId, opts);\n  const captureShapes = {};\n  for (const [name, bufId] of plan.capturesByName) {\n    captureShapes[name] = [...plan.buffers[bufId].shape];\n  }\n  const slot = {\n    runtime,\n    paramNames: [...plan.paramsByName.keys()],\n    outputShape: [...runtime.outputShape],\n    kernelCount: kernels.filter((k) => k.wgsl).length,\n    captureShapes,\n    adam: null\n  };\n  graphs.set(payload.graphId, slot);\n  return {\n    paramNames: [...slot.paramNames],\n    outputShape: slot.outputShape,\n    kernelCount: slot.kernelCount,\n    captureShapes: slot.captureShapes\n  };\n}\nfunction createAdamState(cfg) {\n  return {\n    config: cfg,\n    t: 0,\n    lrtBuf: new Float32Array(1),\n    decayShrinkBuf: cfg.decayShrinkInputName ? new Float32Array(1) : null\n  };\n}\nfunction injectAdamScalars(slot, inputs) {\n  const a = slot.adam;\n  if (!a) return inputs;\n  a.t++;\n  const lrNow = resolveLR(a.config.lr, a.t);\n  a.lrtBuf[0] = lrNow * Math.sqrt(1 - Math.pow(a.config.b2, a.t)) / (1 - Math.pow(a.config.b1, a.t));\n  const merged = { ...inputs, [a.config.lrtInputName]: a.lrtBuf };\n  if (a.decayShrinkBuf && a.config.decayShrinkInputName) {\n    a.decayShrinkBuf[0] = 1 - lrNow * a.config.weightDecay;\n    merged[a.config.decayShrinkInputName] = a.decayShrinkBuf;\n  }\n  return merged;\n}\nasync function handleStep(payload) {\n  const slot = mustGet(payload.graphId);\n  const merged = injectAdamScalars(slot, payload.inputs);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.step(merged, { withCaptures: true });\n    return { loss: r.loss, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const loss = await slot.runtime.step(merged);\n  return { loss, captures: null };\n}\nasync function handleRun(payload) {\n  const slot = mustGet(payload.graphId);\n  if (payload.withCaptures) {\n    const r = await slot.runtime.run(payload.inputs, { withCaptures: true });\n    return { output: r.output, captures: capturesToRecord(r.captures, slot.captureShapes) };\n  }\n  const output = await slot.runtime.run(payload.inputs);\n  return { output, captures: null };\n}\nfunction capturesToRecord(captures, shapes) {\n  const out = {};\n  for (const name of Object.keys(shapes)) {\n    if (captures.has(name)) out[name] = captures.get(name);\n  }\n  return out;\n}\nfunction handleUploadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.uploadParams(payload.params, { partial: payload.partial });\n}\nasync function handleDownloadParams(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParams() };\n}\nasync function handleDownloadParamGrads(payload) {\n  const slot = mustGet(payload.graphId);\n  return { params: await slot.runtime.downloadParamGrads() };\n}\nfunction handleResetOptimizer(payload) {\n  const slot = mustGet(payload.graphId);\n  slot.runtime.resetOptimizerState();\n  if (slot.adam) slot.adam.t = 0;\n}\nfunction handleDestroy(payload) {\n  const slot = graphs.get(payload.graphId);\n  if (!slot) return;\n  slot.runtime.destroy();\n  graphs.delete(payload.graphId);\n}\nfunction mustGet(graphId) {\n  const slot = graphs.get(graphId);\n  if (!slot) throw new Error(`tensorgrad worker: graph ${graphId} not found`);\n  return slot;\n}\nself.onmessage = async (ev) => {\n  const req = ev.data;\n  try {\n    let result;\n    let transferList = [];\n    switch (req.kind) {\n      case "createRuntime":\n        result = await handleCreateRuntime(req.payload);\n        break;\n      case "compileForward":\n        result = await handleCompileForward(req.payload);\n        break;\n      case "step":\n        result = await handleStep(req.payload);\n        transferList = collectTransfers(result.captures);\n        break;\n      case "run": {\n        const r = await handleRun(req.payload);\n        result = r;\n        transferList = [r.output.buffer, ...collectTransfers(r.captures)];\n        break;\n      }\n      case "uploadParams":\n        handleUploadParams(req.payload);\n        result = null;\n        break;\n      case "downloadParams": {\n        const r = await handleDownloadParams(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "downloadParamGrads": {\n        const r = await handleDownloadParamGrads(req.payload);\n        result = r;\n        transferList = collectTransfers(r.params);\n        break;\n      }\n      case "resetOptimizer":\n        handleResetOptimizer(req.payload);\n        result = null;\n        break;\n      case "destroy":\n        handleDestroy(req.payload);\n        result = null;\n        break;\n      default:\n        throw new Error(`unknown request kind: ${req.kind}`);\n    }\n    const reply = { id: req.id, ok: true, result };\n    self.postMessage(reply, { transfer: transferList });\n  } catch (e) {\n    const error = wireError(e);\n    const reply = { id: req.id, ok: false, error };\n    self.postMessage(reply);\n  }\n};\nfunction collectTransfers(rec) {\n  if (!rec) return [];\n  const out = [];\n  for (const v of Object.values(rec)) out.push(v.buffer);\n  return out;\n}\n');
   const wireIR = { graph, plan, kernels };
   const transfers = transferablesOfRecord(initialParams);
   let meta;