tensorgrad 0.0.2 → 0.0.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +7 -9
- package/dist/buffers.d.ts +1 -0
- package/dist/buffers.d.ts.map +1 -1
- package/dist/buffers.js +12 -1
- package/dist/buffers.js.map +1 -1
- package/dist/capture.d.ts +3 -0
- package/dist/capture.d.ts.map +1 -0
- package/dist/capture.js +33 -0
- package/dist/capture.js.map +1 -0
- package/dist/compile.d.ts +33 -5
- package/dist/compile.d.ts.map +1 -1
- package/dist/compile.js +96 -11
- package/dist/compile.js.map +1 -1
- package/dist/index.d.ts +5 -3
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +4 -2
- package/dist/index.js.map +1 -1
- package/dist/ir.d.ts +1 -0
- package/dist/ir.d.ts.map +1 -1
- package/dist/ir.js +1 -1
- package/dist/ir.js.map +1 -1
- package/dist/module.d.ts +30 -4
- package/dist/module.d.ts.map +1 -1
- package/dist/module.js +39 -13
- package/dist/module.js.map +1 -1
- package/dist/nn.d.ts +19 -0
- package/dist/nn.d.ts.map +1 -0
- package/dist/nn.js +60 -0
- package/dist/nn.js.map +1 -0
- package/dist/runtime.d.ts +79 -4
- package/dist/runtime.d.ts.map +1 -1
- package/dist/runtime.js +153 -19
- package/dist/runtime.js.map +1 -1
- package/dist/trace.d.ts +1 -0
- package/dist/trace.d.ts.map +1 -1
- package/dist/trace.js +12 -0
- package/dist/trace.js.map +1 -1
- package/package.json +1 -2
- package/src/buffers.ts +14 -1
- package/src/capture.ts +36 -0
- package/src/compile.ts +112 -13
- package/src/index.ts +5 -3
- package/src/ir.ts +5 -1
- package/src/module.ts +75 -11
- package/src/nn.ts +59 -0
- package/src/runtime.ts +260 -22
- package/src/trace.ts +13 -0
- package/SPEC.md +0 -293
package/README.md
CHANGED
|
@@ -27,8 +27,8 @@ class Linear extends Module {
|
|
|
27
27
|
W: Tensor; b: Tensor
|
|
28
28
|
constructor(public inDim: number, public outDim: number) {
|
|
29
29
|
super()
|
|
30
|
-
this.W = this.param([inDim, outDim])
|
|
31
|
-
this.b = this.param([outDim])
|
|
30
|
+
this.W = this.param([inDim, outDim]) // randn, scale 0.02
|
|
31
|
+
this.b = this.param([outDim], { init: 'zeros' })
|
|
32
32
|
}
|
|
33
33
|
}
|
|
34
34
|
|
|
@@ -50,8 +50,7 @@ function loss(m: MLP, x: Tensor, y: Tensor): Tensor {
|
|
|
50
50
|
}
|
|
51
51
|
|
|
52
52
|
const B = 256
|
|
53
|
-
const
|
|
54
|
-
const compiled = await compileModule(model, loss, {
|
|
53
|
+
const compiled = await compileModule(() => new MLP(), loss, {
|
|
55
54
|
adam: { lr: 0.005 },
|
|
56
55
|
inputs: [
|
|
57
56
|
{ name: 'x', shape: [B, 1], dtype: 'f32' },
|
|
@@ -59,8 +58,8 @@ const compiled = await compileModule(model, loss, {
|
|
|
59
58
|
],
|
|
60
59
|
})
|
|
61
60
|
|
|
62
|
-
//
|
|
63
|
-
|
|
61
|
+
compiled.uploadInitialParams() // applies the per-param init declared above
|
|
62
|
+
|
|
64
63
|
for (let step = 0; step < 1000; step++) {
|
|
65
64
|
const { x, y } = generateBatch()
|
|
66
65
|
const lossVal = await compiled.step({ x, y })
|
|
@@ -94,9 +93,8 @@ If you need full JAX (vmap / pmap / dynamic shapes / multi-backend), use
|
|
|
94
93
|
|
|
95
94
|
## Scope (deliberately small)
|
|
96
95
|
|
|
97
|
-
The library only does what it does because of what it doesn't do.
|
|
98
|
-
|
|
99
|
-
"out of scope" decisions are:
|
|
96
|
+
The library only does what it does because of what it doesn't do. The
|
|
97
|
+
load-bearing "out of scope" decisions are:
|
|
100
98
|
|
|
101
99
|
- **WebGPU only** — no Wasm or WebGL fallback.
|
|
102
100
|
- **Static shapes only** — every shape is fixed at compile time. This is
|
package/dist/buffers.d.ts
CHANGED
|
@@ -30,6 +30,7 @@ export interface BufferPlan {
|
|
|
30
30
|
inputsByName: Map<string, number>;
|
|
31
31
|
paramGradsByName: Map<string, number>;
|
|
32
32
|
statesByName: Map<string, number>;
|
|
33
|
+
capturesByName: Map<string, number>;
|
|
33
34
|
outputBufferIds: number[];
|
|
34
35
|
/** End-of-step writebacks (Adam updates for params, m, v, etc.) */
|
|
35
36
|
writebacks: Writeback[];
|
package/dist/buffers.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"buffers.d.ts","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAcA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAU,MAAM,SAAS,CAAA;AAElE,MAAM,WAAW,UAAU;IACzB,yBAAyB;IACzB,EAAE,EAAE,MAAM,CAAA;IACV,QAAQ,EAAE,MAAM,CAAA;IAChB,KAAK,EAAE,KAAK,CAAA;IACZ,KAAK,EAAE,KAAK,CAAA;IACZ,IAAI,EAAE,OAAO,GAAG,YAAY,GAAG,cAAc,GAAG,OAAO,GAAG,cAAc,GAAG,QAAQ,CAAA;IACnF,sFAAsF;IACtF,IAAI,EAAE,MAAM,GAAG,IAAI,CAAA;IACnB,gFAAgF;IAChF,SAAS,CAAC,EAAE,MAAM,CAAA;CACnB;AAED;;;;GAIG;AACH,MAAM,WAAW,SAAS;IACxB,MAAM,EAAE,MAAM,CAAA;IACd,IAAI,EAAE,MAAM,CAAA;IACZ,KAAK,EAAE,MAAM,CAAA;CACd;AAED,MAAM,WAAW,UAAU;IACzB,OAAO,EAAE,UAAU,EAAE,CAAA;IACrB,iFAAiF;IACjF,cAAc,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACnC,0CAA0C;IAC1C,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,gBAAgB,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACrC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,eAAe,EAAE,MAAM,EAAE,CAAA;IACzB,mEAAmE;IACnE,UAAU,EAAE,SAAS,EAAE,CAAA;CACxB;AAUD;;;GAGG;AACH,MAAM,WAAW,aAAa;IAC5B,0EAA0E;IAC1E,MAAM,EAAE,MAAM,CAAA;IACd,gFAAgF;IAChF,QAAQ,EAAE,MAAM,CAAA;IAChB,QAAQ,EAAE,OAAO,GAAG,OAAO,CAAA;CAC5B;AAED;;;;;;GAMG;AACH,wBAAgB,WAAW,CACzB,KAAK,EAAE,KAAK,EACZ,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAClC,cAAc,GAAE,aAAa,EAAO,GACnC,UAAU,
|
|
1
|
+
{"version":3,"file":"buffers.d.ts","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAcA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAU,MAAM,SAAS,CAAA;AAElE,MAAM,WAAW,UAAU;IACzB,yBAAyB;IACzB,EAAE,EAAE,MAAM,CAAA;IACV,QAAQ,EAAE,MAAM,CAAA;IAChB,KAAK,EAAE,KAAK,CAAA;IACZ,KAAK,EAAE,KAAK,CAAA;IACZ,IAAI,EAAE,OAAO,GAAG,YAAY,GAAG,cAAc,GAAG,OAAO,GAAG,cAAc,GAAG,QAAQ,CAAA;IACnF,sFAAsF;IACtF,IAAI,EAAE,MAAM,GAAG,IAAI,CAAA;IACnB,gFAAgF;IAChF,SAAS,CAAC,EAAE,MAAM,CAAA;CACnB;AAED;;;;GAIG;AACH,MAAM,WAAW,SAAS;IACxB,MAAM,EAAE,MAAM,CAAA;IACd,IAAI,EAAE,MAAM,CAAA;IACZ,KAAK,EAAE,MAAM,CAAA;CACd;AAED,MAAM,WAAW,UAAU;IACzB,OAAO,EAAE,UAAU,EAAE,CAAA;IACrB,iFAAiF;IACjF,cAAc,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACnC,0CAA0C;IAC1C,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,gBAAgB,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACrC,YAAY,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACjC,cAAc,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IACnC,eAAe,EAAE,MAAM,EAAE,CAAA;IACzB,mEAAmE;IACnE,UAAU,EAAE,SAAS,EAAE,CAAA;CACxB;AAUD;;;GAGG;AACH,MAAM,WAAW,aAAa;IAC5B,0EAA0E;IAC1E,MAAM,EAAE,MAAM,CAAA;IACd,gFAAgF;IAChF,QAAQ,EAAE,MAAM,CAAA;IAChB,QAAQ,EAAE,OAAO,GAAG,OAAO,CAAA;CAC5B;AAED;;;;;;GAMG;AACH,wBAAgB,WAAW,CACzB,KAAK,EAAE,KAAK,EACZ,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,EAClC,cAAc,GAAE,aAAa,EAAO,GACnC,UAAU,CAmGZ"}
|
package/dist/buffers.js
CHANGED
|
@@ -109,6 +109,17 @@ export function planBuffers(graph, paramGrads, writebackDecls = []) {
|
|
|
109
109
|
}
|
|
110
110
|
return { source: sourceBufId, dest: destBufId, bytes: sourceSpec.byteSize };
|
|
111
111
|
});
|
|
112
|
-
|
|
112
|
+
// Resolve graph.captures (name -> tensor id) to (name -> buffer id).
|
|
113
|
+
// No pinning needed at the planner level: each tensor already has its own
|
|
114
|
+
// buffer (see "v1 strategy" comment at top — no pooling yet).
|
|
115
|
+
const capturesByName = new Map();
|
|
116
|
+
for (const [name, tensorId] of graph.captures) {
|
|
117
|
+
const bufId = tensorToBuffer.get(tensorId);
|
|
118
|
+
if (bufId === undefined) {
|
|
119
|
+
throw new Error(`planBuffers: capture '${name}' references unknown tensor #${tensorId}`);
|
|
120
|
+
}
|
|
121
|
+
capturesByName.set(name, bufId);
|
|
122
|
+
}
|
|
123
|
+
return { buffers, tensorToBuffer, paramsByName, inputsByName, paramGradsByName, statesByName, capturesByName, outputBufferIds, writebacks };
|
|
113
124
|
}
|
|
114
125
|
//# sourceMappingURL=buffers.js.map
|
package/dist/buffers.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"buffers.js","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,+EAA+E;AAC/E,4EAA4E;AAC5E,8EAA8E;AAC9E,4EAA4E;AAC5E,EAAE;AACF,kDAAkD;AAClD,gFAAgF;AAChF,qFAAqF;AACrF,iEAAiE;AACjE,wEAAwE;AACxE,8EAA8E;
|
|
1
|
+
{"version":3,"file":"buffers.js","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,+EAA+E;AAC/E,4EAA4E;AAC5E,8EAA8E;AAC9E,4EAA4E;AAC5E,EAAE;AACF,kDAAkD;AAClD,gFAAgF;AAChF,qFAAqF;AACrF,iEAAiE;AACjE,wEAAwE;AACxE,8EAA8E;AA2C9E,MAAM,UAAU,GAA0B,EAAE,GAAG,EAAE,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,IAAI,EAAE,CAAC,EAAE,CAAA;AAErE,SAAS,SAAS,CAAC,KAAY;IAC7B,IAAI,CAAC,GAAG,CAAC,CAAA;IACT,KAAK,MAAM,CAAC,IAAI,KAAK;QAAE,CAAC,IAAI,CAAC,CAAA;IAC7B,OAAO,CAAC,CAAA;AACV,CAAC;AAcD;;;;;;GAMG;AACH,MAAM,UAAU,WAAW,CACzB,KAAY,EACZ,UAAkC,EAClC,iBAAkC,EAAE;IAEpC,MAAM,OAAO,GAAiB,EAAE,CAAA;IAChC,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkB,CAAA;IAChD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,gBAAgB,GAAG,IAAI,GAAG,EAAkB,CAAA;IAClD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAE9C,iEAAiE;IACjE,MAAM,kBAAkB,GAAG,IAAI,GAAG,EAAkB,CAAA;IACpD,KAAK,MAAM,CAAC,IAAI,EAAE,MAAM,CAAC,IAAI,MAAM,CAAC,OAAO,CAAC,UAAU,CAAC,EAAE,CAAC;QACxD,kBAAkB,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACzC,CAAC;IACD,2EAA2E;IAC3E,MAAM,SAAS,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC3C,KAAK,MAAM,EAAE,IAAI,KAAK,CAAC,GAAG;QAAE,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,EAAE,EAAE,CAAC,CAAA;IAErD,MAAM,SAAS,GAAG,IAAI,GAAG,CAAC,KAAK,CAAC,OAAO,CAAC,CAAA;IAExC,iDAAiD;IACjD,KAAK,MAAM,CAAC,IAAI,KAAK,CAAC,OAAO,EAAE,CAAC;QAC9B,MAAM,EAAE,GAAG,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAA;QAC9B,IAAI,IAAI,GAAuB,cAAc,CAAA;QAC7C,IAAI,IAAI,GAAkB,IAAI,CAAA;QAC9B,IAAI,SAA6B,CAAA;QAEjC,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YAC/B,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,cAAc,EAAE,CAAC;YACvC,IAAI,GAAG,cAAc,CAAA;YACrB,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YACtC,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;YACd,SAAS,GAAG,EAAE,CAAC,SAAS,CAAA;QAC1B,CAAC;aAAM,IAAI,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YACxC,IAAI,GAAG,YAAY,CAAA;YACnB,IAAI,GAAG,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAE,CAAA;QACtC,CAAC;aAAM,IAAI,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YAC/B,IAAI,GAAG,QAAQ,CAAA;QACjB,CAAC;QAED,MAAM,IAAI,GAAe;YACvB,EAAE,EAAE,CAAC,CAAC,EAAE;YACR,QAAQ,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC,KAAK,CAAC,GAAG,UAAU,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;YAC/D,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,IAAI;YACJ,IAAI;YACJ,GAAG,CAAC,SAAS,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SAClD,CAAA;QACD,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;QAClB,cAAc,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA,CAAE,aAAa;QAE7C,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QACnD,IAAI,IAAI,KAAK,cAAc;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC1D,IAAI,IAAI,KAAK,YAAY;YAAE,gBAAgB,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC5D,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;IACrD,CAAC;IAED,MAAM,eAAe,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAE,CAAC,CAAA;IAExE,oEAAoE;IACpE,MAAM,UAAU,GAAgB,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE;QACxD,MAAM,WAAW,GAAG,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAC,CAAA;QACtD,IAAI,WAAW,KAAK,SAAS,EAAE,CAAC;YAC9B,MAAM,IAAI,KAAK,CAAC,yCAAyC,IAAI,CAAC,MAAM,CAAC,EAAE,eAAe,CAAC,CAAA;QACzF,CAAC;QACD,MAAM,SAAS,GAAG,IAAI,CAAC,QAAQ,KAAK,OAAO;YACzC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC;YACjC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAA;QACnC,IAAI,SAAS,KAAK,SAAS,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,+BAA+B,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,aAAa,CAAC,CAAA;QAC9F,CAAC;QACD,MAAM,UAAU,GAAG,OAAO,CAAC,WAAW,CAAE,CAAA;QACxC,MAAM,QAAQ,GAAG,OAAO,CAAC,SAAS,CAAE,CAAA;QACpC,IAAI,UAAU,CAAC,QAAQ,KAAK,QAAQ,CAAC,QAAQ,EAAE,CAAC;YAC9C,MAAM,IAAI,KAAK,CACb,4CAA4C,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,IAAI;gBAC/E,WAAW,UAAU,CAAC,QAAQ,kBAAkB,QAAQ,CAAC,QAAQ,GAAG,CACrE,CAAA;QACH,CAAC;QACD,OAAO,EAAE,MAAM,EAAE,WAAW,EAAE,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,UAAU,CAAC,QAAQ,EAAE,CAAA;IAC7E,CAAC,CAAC,CAAA;IAEF,qEAAqE;IACrE,0EAA0E;IAC1E,8DAA8D;IAC9D,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkB,CAAA;IAChD,KAAK,MAAM,CAAC,IAAI,EAAE,QAAQ,CAAC,IAAI,KAAK,CAAC,QAAQ,EAAE,CAAC;QAC9C,MAAM,KAAK,GAAG,cAAc,CAAC,GAAG,CAAC,QAAQ,CAAC,CAAA;QAC1C,IAAI,KAAK,KAAK,SAAS,EAAE,CAAC;YACxB,MAAM,IAAI,KAAK,CAAC,yBAAyB,IAAI,gCAAgC,QAAQ,EAAE,CAAC,CAAA;QAC1F,CAAC;QACD,cAAc,CAAC,GAAG,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;IACjC,CAAC;IAED,OAAO,EAAE,OAAO,EAAE,cAAc,EAAE,YAAY,EAAE,YAAY,EAAE,gBAAgB,EAAE,YAAY,EAAE,cAAc,EAAE,eAAe,EAAE,UAAU,EAAE,CAAA;AAC7I,CAAC"}
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"capture.d.ts","sourceRoot":"","sources":["../src/capture.ts"],"names":[],"mappings":"AAqBA,OAAO,KAAK,EAAE,MAAM,EAAE,MAAM,SAAS,CAAA;AAGrC,wBAAgB,OAAO,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,CAW/D"}
|
package/dist/capture.js
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
// Activation capture — opt-in readback of intermediate tensors at training step.
|
|
2
|
+
//
|
|
3
|
+
// Usage (inside the user's forward pass):
|
|
4
|
+
//
|
|
5
|
+
// import { capture } from 'tensorgrad'
|
|
6
|
+
//
|
|
7
|
+
// function attentionFwd(p, x) {
|
|
8
|
+
// const scores = mul(matmulBatched(q, kT), SCALE_QK)
|
|
9
|
+
// const attn = capture(`attn.${layerIdx}`, softmaxCausalLast(scores))
|
|
10
|
+
// return matmulBatched(attn, v)
|
|
11
|
+
// }
|
|
12
|
+
//
|
|
13
|
+
// Pass-through return type: `capture(name, t)` returns `t` unchanged so it
|
|
14
|
+
// inlines at the point of computation. Behind the scenes it registers `t.id`
|
|
15
|
+
// against `name` on the current graph; runtime exposes the registered tensors
|
|
16
|
+
// via `step(inputs, { withCaptures: true })`.
|
|
17
|
+
//
|
|
18
|
+
// Outside the user's forward trace (during `appendGrad` / `appendAdam`'s
|
|
19
|
+
// `traceInto` re-entry), `capture()` is a no-op — gradient and optimizer
|
|
20
|
+
// internals shouldn't accidentally publish themselves to the UI.
|
|
21
|
+
import { currentGraph, isCaptureEnabled } from './trace.js';
|
|
22
|
+
export function capture(name, t) {
|
|
23
|
+
if (!isCaptureEnabled())
|
|
24
|
+
return t;
|
|
25
|
+
const g = currentGraph();
|
|
26
|
+
if (g.captures.has(name)) {
|
|
27
|
+
throw new Error(`capture: name '${name}' already registered. Use unique names ` +
|
|
28
|
+
`(e.g. \`attn.\${layerIdx}\`) when capturing across a loop.`);
|
|
29
|
+
}
|
|
30
|
+
g.captures.set(name, t.id);
|
|
31
|
+
return t;
|
|
32
|
+
}
|
|
33
|
+
//# sourceMappingURL=capture.js.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"file":"capture.js","sourceRoot":"","sources":["../src/capture.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,0CAA0C;AAC1C,EAAE;AACF,yCAAyC;AACzC,EAAE;AACF,kCAAkC;AAClC,yDAAyD;AACzD,0EAA0E;AAC1E,oCAAoC;AACpC,MAAM;AACN,EAAE;AACF,2EAA2E;AAC3E,6EAA6E;AAC7E,8EAA8E;AAC9E,8CAA8C;AAC9C,EAAE;AACF,yEAAyE;AACzE,yEAAyE;AACzE,iEAAiE;AAGjE,OAAO,EAAE,YAAY,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAA;AAE3D,MAAM,UAAU,OAAO,CAAmB,IAAY,EAAE,CAAI;IAC1D,IAAI,CAAC,gBAAgB,EAAE;QAAE,OAAO,CAAC,CAAA;IACjC,MAAM,CAAC,GAAG,YAAY,EAAE,CAAA;IACxB,IAAI,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAC;QACzB,MAAM,IAAI,KAAK,CACb,kBAAkB,IAAI,yCAAyC;YAC/D,4DAA4D,CAC7D,CAAA;IACH,CAAC;IACD,CAAC,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;IAC1B,OAAO,CAAC,CAAA;AACV,CAAC"}
|
package/dist/compile.d.ts
CHANGED
|
@@ -3,7 +3,7 @@ import { type GradResult } from './grad.js';
|
|
|
3
3
|
import { type AdamConfig } from './adam.js';
|
|
4
4
|
import { type BufferPlan } from './buffers.js';
|
|
5
5
|
import { type KernelSpec } from './codegen.js';
|
|
6
|
-
import { type CompiledRuntime, type RuntimeOpts } from './runtime.js';
|
|
6
|
+
import { type CompiledRuntime, type CompiledForward, type RuntimeOpts } from './runtime.js';
|
|
7
7
|
import { Module } from './module.js';
|
|
8
8
|
/** Declares one input tensor of the model's forward function. Order matches
|
|
9
9
|
* the function's parameter list (after `model`). The `name` is used at
|
|
@@ -35,10 +35,18 @@ export interface CompileModuleOptions extends RuntimeOpts {
|
|
|
35
35
|
/** Adam hyperparameters. If omitted, no optimizer is appended (forward-only). */
|
|
36
36
|
adam?: AdamConfig;
|
|
37
37
|
}
|
|
38
|
+
export interface CompileForwardOptions extends RuntimeOpts {
|
|
39
|
+
/** Per-step data inputs to the forward function. */
|
|
40
|
+
inputs?: InputDecl[];
|
|
41
|
+
}
|
|
38
42
|
/**
|
|
39
|
-
* Compile a Module-based model.
|
|
40
|
-
* model
|
|
41
|
-
*
|
|
43
|
+
* Compile a Module-based model. Pass a *factory* `() => new Model()`, not the
|
|
44
|
+
* model instance itself: compilation mutates the tree (every `ParamSentinel`
|
|
45
|
+
* field becomes a real `Tensor`), so the instance is consumed and shouldn't be
|
|
46
|
+
* referenced afterwards. Re-call the factory if you need a fresh tree.
|
|
47
|
+
*
|
|
48
|
+
* The forward function takes the materialized model and returns the loss
|
|
49
|
+
* tensor.
|
|
42
50
|
*
|
|
43
51
|
* Walks the module tree to materialize params with auto-derived names, then
|
|
44
52
|
* runs trace → grad → adam → buffer plan → codegen → runtime.
|
|
@@ -47,7 +55,27 @@ export interface CompileModuleOptions extends RuntimeOpts {
|
|
|
47
55
|
* internal step count and injects the bias-corrected `lrt` scalar each call;
|
|
48
56
|
* users don't need to provide it themselves.
|
|
49
57
|
*/
|
|
50
|
-
export declare function compileModule<M extends Module>(
|
|
58
|
+
export declare function compileModule<M extends Module>(modelFactory: () => M, forward: (m: M, ...inputs: Tensor[]) => Tensor, opts?: CompileModuleOptions): Promise<CompiledRuntime & {
|
|
59
|
+
ir: CompiledIR;
|
|
60
|
+
uploadInitialParams: () => void;
|
|
61
|
+
}>;
|
|
62
|
+
/**
|
|
63
|
+
* Compile a Module-based model in forward-only mode (no autograd, no Adam).
|
|
64
|
+
* The forward function returns the output tensor (e.g., logits) instead of a
|
|
65
|
+
* scalar loss; runtime exposes `run(inputs)` returning the full output as a
|
|
66
|
+
* `Float32Array`.
|
|
67
|
+
*
|
|
68
|
+
* **Sharing params with a training compile.** Pass `opts.sharedParams =
|
|
69
|
+
* trainCompiled.params` to bind this graph's param buffers to an existing
|
|
70
|
+
* training runtime's GPU buffers — every train step is then immediately
|
|
71
|
+
* visible to `run()` calls here, no copies. The forward graph's
|
|
72
|
+
* `uploadInitialParams()` skips any param covered by `sharedParams`.
|
|
73
|
+
*
|
|
74
|
+
* Typical use: a B=1 inference graph alongside a B=512 training graph,
|
|
75
|
+
* built from the same `Module` factory.
|
|
76
|
+
*/
|
|
77
|
+
export declare function compileForward<M extends Module>(modelFactory: () => M, forward: (m: M, ...inputs: Tensor[]) => Tensor, opts?: CompileForwardOptions): Promise<CompiledForward & {
|
|
51
78
|
ir: CompiledIR;
|
|
79
|
+
uploadInitialParams: () => void;
|
|
52
80
|
}>;
|
|
53
81
|
//# sourceMappingURL=compile.d.ts.map
|
package/dist/compile.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"compile.d.ts","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAUA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAEnD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,
|
|
1
|
+
{"version":3,"file":"compile.d.ts","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAUA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAEnD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAc,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAe,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAuC,KAAK,eAAe,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,MAAM,cAAc,CAAA;AAChI,OAAO,EAAE,MAAM,EAAqB,MAAM,aAAa,CAAA;AAEvD;;2DAE2D;AAC3D,MAAM,WAAW,SAAS;IACxB,IAAI,EAAE,MAAM,CAAA;IACZ,KAAK,EAAE,KAAK,CAAA;IACZ,KAAK,CAAC,EAAE,KAAK,CAAA;CACd;AAED,MAAM,WAAW,UAAU;IACzB,KAAK,EAAE,UAAU,CAAC,OAAO,CAAC,CAAA;IAC1B,UAAU,EAAE,UAAU,CAAC,YAAY,CAAC,CAAA;IACpC,IAAI,EAAE,MAAM,CAAA;IACZ,IAAI,EAAE,UAAU,CAAA;IAChB,OAAO,EAAE,UAAU,EAAE,CAAA;CACtB;AAED,yEAAyE;AACzE,wBAAgB,WAAW,CAAC,OAAO,EAAE,MAAM,MAAM,GAAG,UAAU,CAM7D;AAED,0EAA0E;AAC1E,wBAAsB,OAAO,CAAC,OAAO,EAAE,MAAM,MAAM,EAAE,IAAI,GAAE,WAAgB,GAAG,OAAO,CAAC,eAAe,GAAG;IAAE,EAAE,EAAE,UAAU,CAAA;CAAE,CAAC,CAK1H;AAMD,MAAM,WAAW,oBAAqB,SAAQ,WAAW;IACvD;;;4EAGwE;IACxE,MAAM,CAAC,EAAE,SAAS,EAAE,CAAA;IACpB,iFAAiF;IACjF,IAAI,CAAC,EAAE,UAAU,CAAA;CAClB;AAED,MAAM,WAAW,qBAAsB,SAAQ,WAAW;IACxD,oDAAoD;IACpD,MAAM,CAAC,EAAE,SAAS,EAAE,CAAA;CACrB;AAED;;;;;;;;;;;;;;;GAeG;AACH,wBAAsB,aAAa,CAAC,CAAC,SAAS,MAAM,EAClD,YAAY,EAAE,MAAM,CAAC,EACrB,OAAO,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,GAAG,MAAM,EAAE,MAAM,EAAE,KAAK,MAAM,EAC9C,IAAI,GAAE,oBAAyB,GAC9B,OAAO,CAAC,eAAe,GAAG;IAAE,EAAE,EAAE,UAAU,CAAC;IAAC,mBAAmB,EAAE,MAAM,IAAI,CAAA;CAAE,CAAC,CA8DhF;AAMD;;;;;;;;;;;;;;GAcG;AACH,wBAAsB,cAAc,CAAC,CAAC,SAAS,MAAM,EACnD,YAAY,EAAE,MAAM,CAAC,EACrB,OAAO,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,GAAG,MAAM,EAAE,MAAM,EAAE,KAAK,MAAM,EAC9C,IAAI,GAAE,qBAA0B,GAC/B,OAAO,CAAC,eAAe,GAAG;IAAE,EAAE,EAAE,UAAU,CAAC;IAAC,mBAAmB,EAAE,MAAM,IAAI,CAAA;CAAE,CAAC,CAuChF"}
|
package/dist/compile.js
CHANGED
|
@@ -12,7 +12,7 @@ import { appendGrad } from './grad.js';
|
|
|
12
12
|
import { appendAdam } from './adam.js';
|
|
13
13
|
import { planBuffers } from './buffers.js';
|
|
14
14
|
import { emitKernels } from './codegen.js';
|
|
15
|
-
import { createRuntime } from './runtime.js';
|
|
15
|
+
import { createRuntime, createForwardRuntime } from './runtime.js';
|
|
16
16
|
import { Module, materializeParams } from './module.js';
|
|
17
17
|
/** Trace + autograd + buffer-plan + codegen, without touching WebGPU. */
|
|
18
18
|
export function compileToIR(traceFn) {
|
|
@@ -30,9 +30,13 @@ export async function compile(traceFn, opts = {}) {
|
|
|
30
30
|
return Object.assign(runtime, { ir });
|
|
31
31
|
}
|
|
32
32
|
/**
|
|
33
|
-
* Compile a Module-based model.
|
|
34
|
-
* model
|
|
35
|
-
*
|
|
33
|
+
* Compile a Module-based model. Pass a *factory* `() => new Model()`, not the
|
|
34
|
+
* model instance itself: compilation mutates the tree (every `ParamSentinel`
|
|
35
|
+
* field becomes a real `Tensor`), so the instance is consumed and shouldn't be
|
|
36
|
+
* referenced afterwards. Re-call the factory if you need a fresh tree.
|
|
37
|
+
*
|
|
38
|
+
* The forward function takes the materialized model and returns the loss
|
|
39
|
+
* tensor.
|
|
36
40
|
*
|
|
37
41
|
* Walks the module tree to materialize params with auto-derived names, then
|
|
38
42
|
* runs trace → grad → adam → buffer plan → codegen → runtime.
|
|
@@ -41,36 +45,117 @@ export async function compile(traceFn, opts = {}) {
|
|
|
41
45
|
* internal step count and injects the bias-corrected `lrt` scalar each call;
|
|
42
46
|
* users don't need to provide it themselves.
|
|
43
47
|
*/
|
|
44
|
-
export async function compileModule(
|
|
48
|
+
export async function compileModule(modelFactory, forward, opts = {}) {
|
|
45
49
|
const inputDecls = opts.inputs ?? [];
|
|
46
|
-
|
|
50
|
+
const model = modelFactory();
|
|
51
|
+
let materialized = { tensors: {}, initFns: {} };
|
|
47
52
|
const graph = trace(() => {
|
|
48
|
-
|
|
53
|
+
materialized = materializeParams(model);
|
|
49
54
|
const inputTensors = inputDecls.map(d => tensorInput(d.name, d.shape, d.dtype ?? 'f32'));
|
|
50
55
|
return forward(model, ...inputTensors);
|
|
51
56
|
});
|
|
52
57
|
const { paramGrads, loss } = appendGrad(graph);
|
|
53
58
|
let adamResult;
|
|
54
59
|
if (opts.adam) {
|
|
55
|
-
adamResult = appendAdam(graph, paramGrads,
|
|
60
|
+
adamResult = appendAdam(graph, paramGrads, materialized.tensors, opts.adam);
|
|
56
61
|
}
|
|
57
62
|
const plan = planBuffers(graph, paramGrads, adamResult?.writebacks ?? []);
|
|
58
63
|
const kernels = emitKernels(graph, plan);
|
|
59
64
|
const lossBufferId = plan.tensorToBuffer.get(loss.id);
|
|
60
65
|
const runtime = await createRuntime(plan, kernels, lossBufferId, opts);
|
|
61
66
|
// If Adam is enabled, wrap step() to track the step count and supply lrt.
|
|
67
|
+
// Wrap resetOptimizerState() too, so a reset zeros m/v *and* the bias-correction
|
|
68
|
+
// counter — otherwise the next step would skip Adam's warmup phase.
|
|
62
69
|
if (adamResult) {
|
|
63
70
|
const { lrtInputName, config } = adamResult;
|
|
64
71
|
let t = 0;
|
|
65
72
|
const lrtBuf = new Float32Array(1);
|
|
66
73
|
const innerStep = runtime.step.bind(runtime);
|
|
67
|
-
|
|
74
|
+
const innerReset = runtime.resetOptimizerState.bind(runtime);
|
|
75
|
+
const wrappedStep = (inputs, opts) => {
|
|
68
76
|
t++;
|
|
69
77
|
lrtBuf[0] = config.lr * Math.sqrt(1 - Math.pow(config.b2, t)) / (1 - Math.pow(config.b1, t));
|
|
70
|
-
|
|
78
|
+
const merged = { ...inputs, [lrtInputName]: lrtBuf };
|
|
79
|
+
return opts?.withCaptures ? innerStep(merged, { withCaptures: true }) : innerStep(merged);
|
|
80
|
+
};
|
|
81
|
+
runtime.step = wrappedStep;
|
|
82
|
+
runtime.resetOptimizerState = () => {
|
|
83
|
+
t = 0;
|
|
84
|
+
innerReset();
|
|
71
85
|
};
|
|
72
86
|
}
|
|
87
|
+
const { initFns } = materialized;
|
|
88
|
+
const uploadInitialParams = () => {
|
|
89
|
+
const out = {};
|
|
90
|
+
for (const [name, bufId] of plan.paramsByName) {
|
|
91
|
+
const shape = plan.buffers[bufId].shape;
|
|
92
|
+
const size = shape.reduce((a, b) => a * b, 1);
|
|
93
|
+
const initFn = initFns[name];
|
|
94
|
+
if (!initFn)
|
|
95
|
+
throw new Error(`uploadInitialParams: no init for param '${name}'`);
|
|
96
|
+
out[name] = initFn(size, shape);
|
|
97
|
+
}
|
|
98
|
+
runtime.uploadParams(out);
|
|
99
|
+
};
|
|
73
100
|
const ir = { graph, paramGrads, loss, plan, kernels };
|
|
74
|
-
return Object.assign(runtime, { ir });
|
|
101
|
+
return Object.assign(runtime, { ir, uploadInitialParams });
|
|
102
|
+
}
|
|
103
|
+
// ============================================================================
|
|
104
|
+
// Forward-only compile
|
|
105
|
+
// ============================================================================
|
|
106
|
+
/**
|
|
107
|
+
* Compile a Module-based model in forward-only mode (no autograd, no Adam).
|
|
108
|
+
* The forward function returns the output tensor (e.g., logits) instead of a
|
|
109
|
+
* scalar loss; runtime exposes `run(inputs)` returning the full output as a
|
|
110
|
+
* `Float32Array`.
|
|
111
|
+
*
|
|
112
|
+
* **Sharing params with a training compile.** Pass `opts.sharedParams =
|
|
113
|
+
* trainCompiled.params` to bind this graph's param buffers to an existing
|
|
114
|
+
* training runtime's GPU buffers — every train step is then immediately
|
|
115
|
+
* visible to `run()` calls here, no copies. The forward graph's
|
|
116
|
+
* `uploadInitialParams()` skips any param covered by `sharedParams`.
|
|
117
|
+
*
|
|
118
|
+
* Typical use: a B=1 inference graph alongside a B=512 training graph,
|
|
119
|
+
* built from the same `Module` factory.
|
|
120
|
+
*/
|
|
121
|
+
export async function compileForward(modelFactory, forward, opts = {}) {
|
|
122
|
+
const inputDecls = opts.inputs ?? [];
|
|
123
|
+
const model = modelFactory();
|
|
124
|
+
let materialized = { tensors: {}, initFns: {} };
|
|
125
|
+
const graph = trace(() => {
|
|
126
|
+
materialized = materializeParams(model);
|
|
127
|
+
const inputTensors = inputDecls.map(d => tensorInput(d.name, d.shape, d.dtype ?? 'f32'));
|
|
128
|
+
return forward(model, ...inputTensors);
|
|
129
|
+
});
|
|
130
|
+
const plan = planBuffers(graph, /* paramGrads */ {});
|
|
131
|
+
const kernels = emitKernels(graph, plan);
|
|
132
|
+
const outputTensor = graph.tensors[graph.outputs[0]];
|
|
133
|
+
const outputBufferId = plan.tensorToBuffer.get(outputTensor.id);
|
|
134
|
+
const runtime = await createForwardRuntime(plan, kernels, outputBufferId, opts);
|
|
135
|
+
const sharedParams = opts.sharedParams;
|
|
136
|
+
const { initFns } = materialized;
|
|
137
|
+
const uploadInitialParams = () => {
|
|
138
|
+
const out = {};
|
|
139
|
+
let needsUpload = false;
|
|
140
|
+
for (const [name, bufId] of plan.paramsByName) {
|
|
141
|
+
// Skip params covered by sharedParams — those are owned by the providing
|
|
142
|
+
// compile and already initialized there.
|
|
143
|
+
if (sharedParams?.has(name))
|
|
144
|
+
continue;
|
|
145
|
+
const shape = plan.buffers[bufId].shape;
|
|
146
|
+
const size = shape.reduce((a, b) => a * b, 1);
|
|
147
|
+
const initFn = initFns[name];
|
|
148
|
+
if (!initFn)
|
|
149
|
+
throw new Error(`uploadInitialParams: no init for param '${name}'`);
|
|
150
|
+
out[name] = initFn(size, shape);
|
|
151
|
+
needsUpload = true;
|
|
152
|
+
}
|
|
153
|
+
if (needsUpload)
|
|
154
|
+
runtime.uploadParams(out, { partial: !!sharedParams });
|
|
155
|
+
};
|
|
156
|
+
// CompiledIR.loss is the field name; for forward-only, it carries the user's
|
|
157
|
+
// returned tensor (e.g., logits). Same shape conceptually; just no autograd.
|
|
158
|
+
const ir = { graph, paramGrads: {}, loss: outputTensor, plan, kernels };
|
|
159
|
+
return Object.assign(runtime, { ir, uploadInitialParams });
|
|
75
160
|
}
|
|
76
161
|
//# sourceMappingURL=compile.js.map
|
package/dist/compile.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"compile.js","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,EAAE;AACF,oBAAoB;AACpB,sEAAsE;AACtE,iEAAiE;AACjE,0EAA0E;AAC1E,0EAA0E;AAC1E,2EAA2E;AAC3E,mEAAmE;AAGnE,OAAO,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,YAAY,CAAA;AAC/C,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,
|
|
1
|
+
{"version":3,"file":"compile.js","sourceRoot":"","sources":["../src/compile.ts"],"names":[],"mappings":"AAAA,2EAA2E;AAC3E,EAAE;AACF,oBAAoB;AACpB,sEAAsE;AACtE,iEAAiE;AACjE,0EAA0E;AAC1E,0EAA0E;AAC1E,2EAA2E;AAC3E,mEAAmE;AAGnE,OAAO,EAAE,KAAK,EAAE,WAAW,EAAE,MAAM,YAAY,CAAA;AAC/C,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAgE,MAAM,cAAc,CAAA;AAChI,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,MAAM,aAAa,CAAA;AAmBvD,yEAAyE;AACzE,MAAM,UAAU,WAAW,CAAC,OAAqB;IAC/C,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,CAAA;IAC5B,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAC9C,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,CAAC,CAAA;IAC3C,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;AACnD,CAAC;AAED,0EAA0E;AAC1E,MAAM,CAAC,KAAK,UAAU,OAAO,CAAC,OAAqB,EAAE,OAAoB,EAAE;IACzE,MAAM,EAAE,GAAG,WAAW,CAAC,OAAO,CAAC,CAAA;IAC/B,MAAM,YAAY,GAAG,EAAE,CAAC,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IAC5D,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,EAAE,CAAC,IAAI,EAAE,EAAE,CAAC,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAC5E,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,CAAC,CAAA;AACvC,CAAC;AAqBD;;;;;;;;;;;;;;;GAeG;AACH,MAAM,CAAC,KAAK,UAAU,aAAa,CACjC,YAAqB,EACrB,OAA8C,EAC9C,OAA6B,EAAE;IAE/B,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,IAAI,EAAE,CAAA;IACpC,MAAM,KAAK,GAAG,YAAY,EAAE,CAAA;IAC5B,IAAI,YAAY,GAAyC,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,CAAA;IACrF,MAAM,KAAK,GAAG,KAAK,CAAC,GAAG,EAAE;QACvB,YAAY,GAAG,iBAAiB,CAAC,KAAK,CAAC,CAAA;QACvC,MAAM,YAAY,GAAG,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,KAAK,CAAC,CAAC,CAAA;QACxF,OAAO,OAAO,CAAC,KAAK,EAAE,GAAG,YAAY,CAAC,CAAA;IACxC,CAAC,CAAC,CAAA;IAEF,MAAM,EAAE,UAAU,EAAE,IAAI,EAAE,GAAG,UAAU,CAAC,KAAK,CAAC,CAAA;IAE9C,IAAI,UAAqD,CAAA;IACzD,IAAI,IAAI,CAAC,IAAI,EAAE,CAAC;QACd,UAAU,GAAG,UAAU,CAAC,KAAK,EAAE,UAAU,EAAE,YAAY,CAAC,OAAO,EAAE,IAAI,CAAC,IAAI,CAAC,CAAA;IAC7E,CAAC;IAED,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,UAAU,EAAE,UAAU,EAAE,UAAU,IAAI,EAAE,CAAC,CAAA;IACzE,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,MAAM,YAAY,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE,CAAE,CAAA;IACtD,MAAM,OAAO,GAAG,MAAM,aAAa,CAAC,IAAI,EAAE,OAAO,EAAE,YAAY,EAAE,IAAI,CAAC,CAAA;IAEtE,0EAA0E;IAC1E,iFAAiF;IACjF,oEAAoE;IACpE,IAAI,UAAU,EAAE,CAAC;QACf,MAAM,EAAE,YAAY,EAAE,MAAM,EAAE,GAAG,UAAU,CAAA;QAC3C,IAAI,CAAC,GAAG,CAAC,CAAA;QACT,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAA;QAClC,MAAM,SAAS,GAAG,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAA4B,CAAA;QACvE,MAAM,UAAU,GAAG,OAAO,CAAC,mBAAmB,CAAC,IAAI,CAAC,OAAO,CAAC,CAAA;QAC5D,MAAM,WAAW,GAAG,CAClB,MAAiD,EACjD,IAAiC,EAC2C,EAAE;YAC9E,CAAC,EAAE,CAAA;YACH,MAAM,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,EAAE,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,CAAA;YAC5F,MAAM,MAAM,GAAG,EAAE,GAAG,MAAM,EAAE,CAAC,YAAY,CAAC,EAAE,MAAM,EAAE,CAAA;YACpD,OAAO,IAAI,EAAE,YAAY,CAAC,CAAC,CAAC,SAAS,CAAC,MAAM,EAAE,EAAE,YAAY,EAAE,IAAI,EAAE,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,MAAM,CAAC,CAAA;QAC3F,CAAC,CAAA;QACD,OAAO,CAAC,IAAI,GAAG,WAAsC,CAAA;QACrD,OAAO,CAAC,mBAAmB,GAAG,GAAG,EAAE;YACjC,CAAC,GAAG,CAAC,CAAA;YACL,UAAU,EAAE,CAAA;QACd,CAAC,CAAA;IACH,CAAC;IAED,MAAM,EAAE,OAAO,EAAE,GAAG,YAAY,CAAA;IAChC,MAAM,mBAAmB,GAAG,GAAG,EAAE;QAC/B,MAAM,GAAG,GAAiC,EAAE,CAAA;QAC5C,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,KAAK,CAAA;YACxC,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAA;YAC7C,MAAM,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAA;YAC5B,IAAI,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,2CAA2C,IAAI,GAAG,CAAC,CAAA;YAChF,GAAG,CAAC,IAAI,CAAC,GAAG,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;QACjC,CAAC;QACD,OAAO,CAAC,YAAY,CAAC,GAAG,CAAC,CAAA;IAC3B,CAAC,CAAA;IAED,MAAM,EAAE,GAAe,EAAE,KAAK,EAAE,UAAU,EAAE,IAAI,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;IACjE,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,mBAAmB,EAAE,CAAC,CAAA;AAC5D,CAAC;AAED,+EAA+E;AAC/E,uBAAuB;AACvB,+EAA+E;AAE/E;;;;;;;;;;;;;;GAcG;AACH,MAAM,CAAC,KAAK,UAAU,cAAc,CAClC,YAAqB,EACrB,OAA8C,EAC9C,OAA8B,EAAE;IAEhC,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,IAAI,EAAE,CAAA;IACpC,MAAM,KAAK,GAAG,YAAY,EAAE,CAAA;IAC5B,IAAI,YAAY,GAAyC,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,CAAA;IACrF,MAAM,KAAK,GAAG,KAAK,CAAC,GAAG,EAAE;QACvB,YAAY,GAAG,iBAAiB,CAAC,KAAK,CAAC,CAAA;QACvC,MAAM,YAAY,GAAG,UAAU,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,WAAW,CAAC,CAAC,CAAC,IAAI,EAAE,CAAC,CAAC,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,KAAK,CAAC,CAAC,CAAA;QACxF,OAAO,OAAO,CAAC,KAAK,EAAE,GAAG,YAAY,CAAC,CAAA;IACxC,CAAC,CAAC,CAAA;IAEF,MAAM,IAAI,GAAG,WAAW,CAAC,KAAK,EAAE,gBAAgB,CAAC,EAAE,CAAC,CAAA;IACpD,MAAM,OAAO,GAAG,WAAW,CAAC,KAAK,EAAE,IAAI,CAAC,CAAA;IACxC,MAAM,YAAY,GAAG,KAAK,CAAC,OAAO,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC,CAAE,CAAE,CAAA;IACtD,MAAM,cAAc,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,YAAY,CAAC,EAAE,CAAE,CAAA;IAChE,MAAM,OAAO,GAAG,MAAM,oBAAoB,CAAC,IAAI,EAAE,OAAO,EAAE,cAAc,EAAE,IAAI,CAAC,CAAA;IAE/E,MAAM,YAAY,GAAG,IAAI,CAAC,YAAY,CAAA;IACtC,MAAM,EAAE,OAAO,EAAE,GAAG,YAAY,CAAA;IAChC,MAAM,mBAAmB,GAAG,GAAG,EAAE;QAC/B,MAAM,GAAG,GAAiC,EAAE,CAAA;QAC5C,IAAI,WAAW,GAAG,KAAK,CAAA;QACvB,KAAK,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,IAAI,IAAI,CAAC,YAAY,EAAE,CAAC;YAC9C,yEAAyE;YACzE,yCAAyC;YACzC,IAAI,YAAY,EAAE,GAAG,CAAC,IAAI,CAAC;gBAAE,SAAQ;YACrC,MAAM,KAAK,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAE,CAAC,KAAK,CAAA;YACxC,MAAM,IAAI,GAAG,KAAK,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAA;YAC7C,MAAM,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAA;YAC5B,IAAI,CAAC,MAAM;gBAAE,MAAM,IAAI,KAAK,CAAC,2CAA2C,IAAI,GAAG,CAAC,CAAA;YAChF,GAAG,CAAC,IAAI,CAAC,GAAG,MAAM,CAAC,IAAI,EAAE,KAAK,CAAC,CAAA;YAC/B,WAAW,GAAG,IAAI,CAAA;QACpB,CAAC;QACD,IAAI,WAAW;YAAE,OAAO,CAAC,YAAY,CAAC,GAAG,EAAE,EAAE,OAAO,EAAE,CAAC,CAAC,YAAY,EAAE,CAAC,CAAA;IACzE,CAAC,CAAA;IAED,6EAA6E;IAC7E,6EAA6E;IAC7E,MAAM,EAAE,GAAe,EAAE,KAAK,EAAE,UAAU,EAAE,EAAE,EAAE,IAAI,EAAE,YAAY,EAAE,IAAI,EAAE,OAAO,EAAE,CAAA;IACnF,OAAO,MAAM,CAAC,MAAM,CAAC,OAAO,EAAE,EAAE,EAAE,EAAE,mBAAmB,EAAE,CAAC,CAAA;AAC5D,CAAC"}
|
package/dist/index.d.ts
CHANGED
|
@@ -1,12 +1,14 @@
|
|
|
1
1
|
export type { Tensor, Shape, Dtype, OpNode, Graph, CallSite } from './ir.js';
|
|
2
2
|
export { ShapeError } from './shape.js';
|
|
3
3
|
export { trace, traceInto, paramInput, tensorInput, stateInput } from './trace.js';
|
|
4
|
+
export { capture } from './capture.js';
|
|
4
5
|
export { add, sub, mul, div, sqrt, rsqrt, log, exp, relu, less, greater, where, meanLast, sumLast, reshape, transpose, matmul, matmulBatched, oneHot, arange, softmaxCausalLast, logSoftmaxLast, whereCausal, sliceLastRange, } from './ops.js';
|
|
5
6
|
export { appendGrad, type GradResult } from './grad.js';
|
|
6
7
|
export { appendAdam, type AdamConfig, type AdamResult } from './adam.js';
|
|
7
8
|
export { planBuffers, type BufferPlan, type BufferSpec, type Writeback, type WritebackDecl } from './buffers.js';
|
|
8
9
|
export { emitKernels, type KernelSpec } from './codegen.js';
|
|
9
|
-
export { createRuntime, type CompiledRuntime, type RuntimeOpts } from './runtime.js';
|
|
10
|
-
export { compile, compileToIR, compileModule, type CompiledIR, type CompileModuleOptions, type InputDecl } from './compile.js';
|
|
11
|
-
export { Module, materializeParams } from './module.js';
|
|
10
|
+
export { createRuntime, createForwardRuntime, type CompiledRuntime, type CompiledForward, type RuntimeOpts, type StepOptions, type StepWithCaptures, type RunOptions, type RunWithCaptures } from './runtime.js';
|
|
11
|
+
export { compile, compileToIR, compileModule, compileForward, type CompiledIR, type CompileModuleOptions, type CompileForwardOptions, type InputDecl } from './compile.js';
|
|
12
|
+
export { Module, materializeParams, type InitSpec, type ParamOptions, type MaterializedParams } from './module.js';
|
|
13
|
+
export * as nn from './nn.js';
|
|
12
14
|
//# sourceMappingURL=index.d.ts.map
|
package/dist/index.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAKA,YAAY,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,QAAQ,EAAE,MAAM,SAAS,CAAA;AAC5E,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAEL,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG,EAElB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI,EAE3B,IAAI,EAAE,OAAO,EAAE,KAAK,EAEpB,QAAQ,EAAE,OAAO,EAEjB,OAAO,EAAE,SAAS,EAElB,MAAM,EAAE,aAAa,EAErB,MAAM,EAAE,MAAM,EAEd,iBAAiB,EAAE,cAAc,EAAE,WAAW,EAE9C,cAAc,GACf,MAAM,UAAU,CAAA;AAMjB,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,SAAS,EAAE,KAAK,aAAa,EAAE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,MAAM,cAAc,CAAA;
|
|
1
|
+
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAKA,YAAY,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,QAAQ,EAAE,MAAM,SAAS,CAAA;AAC5E,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAAE,OAAO,EAAE,MAAM,cAAc,CAAA;AACtC,OAAO,EAEL,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG,EAElB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI,EAE3B,IAAI,EAAE,OAAO,EAAE,KAAK,EAEpB,QAAQ,EAAE,OAAO,EAEjB,OAAO,EAAE,SAAS,EAElB,MAAM,EAAE,aAAa,EAErB,MAAM,EAAE,MAAM,EAEd,iBAAiB,EAAE,cAAc,EAAE,WAAW,EAE9C,cAAc,GACf,MAAM,UAAU,CAAA;AAMjB,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,KAAK,UAAU,EAAE,KAAK,SAAS,EAAE,KAAK,aAAa,EAAE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAE,KAAK,eAAe,EAAE,KAAK,eAAe,EAAE,KAAK,WAAW,EAAE,KAAK,WAAW,EAAE,KAAK,gBAAgB,EAAE,KAAK,UAAU,EAAE,KAAK,eAAe,EAAE,MAAM,cAAc,CAAA;AAChN,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,cAAc,EAAE,KAAK,UAAU,EAAE,KAAK,oBAAoB,EAAE,KAAK,qBAAqB,EAAE,KAAK,SAAS,EAAE,MAAM,cAAc,CAAA;AAC1K,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAAE,KAAK,QAAQ,EAAE,KAAK,YAAY,EAAE,KAAK,kBAAkB,EAAE,MAAM,aAAa,CAAA;AAClH,OAAO,KAAK,EAAE,MAAM,SAAS,CAAA"}
|
package/dist/index.js
CHANGED
|
@@ -4,6 +4,7 @@
|
|
|
4
4
|
// codegen / compile() (Phase 3+) come later.
|
|
5
5
|
export { ShapeError } from './shape.js';
|
|
6
6
|
export { trace, traceInto, paramInput, tensorInput, stateInput } from './trace.js';
|
|
7
|
+
export { capture } from './capture.js';
|
|
7
8
|
export {
|
|
8
9
|
// Element-wise arithmetic. The binops accept Tensor or JS-number for the second arg.
|
|
9
10
|
add, sub, mul, div,
|
|
@@ -31,7 +32,8 @@ export { appendGrad } from './grad.js';
|
|
|
31
32
|
export { appendAdam } from './adam.js';
|
|
32
33
|
export { planBuffers } from './buffers.js';
|
|
33
34
|
export { emitKernels } from './codegen.js';
|
|
34
|
-
export { createRuntime } from './runtime.js';
|
|
35
|
-
export { compile, compileToIR, compileModule } from './compile.js';
|
|
35
|
+
export { createRuntime, createForwardRuntime } from './runtime.js';
|
|
36
|
+
export { compile, compileToIR, compileModule, compileForward } from './compile.js';
|
|
36
37
|
export { Module, materializeParams } from './module.js';
|
|
38
|
+
export * as nn from './nn.js';
|
|
37
39
|
//# sourceMappingURL=index.js.map
|
package/dist/index.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,+CAA+C;AAC/C,EAAE;AACF,8EAA8E;AAC9E,6CAA6C;AAG7C,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO;AACL,qFAAqF;AACrF,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG;AAClB,qBAAqB;AACrB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI;AAC3B,uBAAuB;AACvB,IAAI,EAAE,OAAO,EAAE,KAAK;AACpB,yEAAyE;AACzE,QAAQ,EAAE,OAAO;AACjB,YAAY;AACZ,OAAO,EAAE,SAAS;AAClB,iBAAiB;AACjB,MAAM,EAAE,aAAa;AACrB,qBAAqB;AACrB,MAAM,EAAE,MAAM;AACd,4CAA4C;AAC5C,iBAAiB,EAAE,cAAc,EAAE,WAAW;AAC9C,UAAU;AACV,cAAc,GACf,MAAM,UAAU,CAAA;AAEjB,sFAAsF;AACtF,8EAA8E;AAC9E,2EAA2E;AAC3E,qDAAqD;AACrD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAoC,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAwE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,
|
|
1
|
+
{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,+CAA+C;AAC/C,EAAE;AACF,8EAA8E;AAC9E,6CAA6C;AAG7C,OAAO,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AACvC,OAAO,EAAE,KAAK,EAAE,SAAS,EAAE,UAAU,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,YAAY,CAAA;AAClF,OAAO,EAAE,OAAO,EAAE,MAAM,cAAc,CAAA;AACtC,OAAO;AACL,qFAAqF;AACrF,GAAG,EAAE,GAAG,EAAE,GAAG,EAAE,GAAG;AAClB,qBAAqB;AACrB,IAAI,EAAE,KAAK,EAAE,GAAG,EAAE,GAAG,EAAE,IAAI;AAC3B,uBAAuB;AACvB,IAAI,EAAE,OAAO,EAAE,KAAK;AACpB,yEAAyE;AACzE,QAAQ,EAAE,OAAO;AACjB,YAAY;AACZ,OAAO,EAAE,SAAS;AAClB,iBAAiB;AACjB,MAAM,EAAE,aAAa;AACrB,qBAAqB;AACrB,MAAM,EAAE,MAAM;AACd,4CAA4C;AAC5C,iBAAiB,EAAE,cAAc,EAAE,WAAW;AAC9C,UAAU;AACV,cAAc,GACf,MAAM,UAAU,CAAA;AAEjB,sFAAsF;AACtF,8EAA8E;AAC9E,2EAA2E;AAC3E,qDAAqD;AACrD,OAAO,EAAE,UAAU,EAAmB,MAAM,WAAW,CAAA;AACvD,OAAO,EAAE,UAAU,EAAoC,MAAM,WAAW,CAAA;AACxE,OAAO,EAAE,WAAW,EAAwE,MAAM,cAAc,CAAA;AAChH,OAAO,EAAE,WAAW,EAAmB,MAAM,cAAc,CAAA;AAC3D,OAAO,EAAE,aAAa,EAAE,oBAAoB,EAAgJ,MAAM,cAAc,CAAA;AAChN,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,aAAa,EAAE,cAAc,EAA0F,MAAM,cAAc,CAAA;AAC1K,OAAO,EAAE,MAAM,EAAE,iBAAiB,EAA6D,MAAM,aAAa,CAAA;AAClH,OAAO,KAAK,EAAE,MAAM,SAAS,CAAA"}
|
package/dist/ir.d.ts
CHANGED
|
@@ -194,6 +194,7 @@ export interface Graph {
|
|
|
194
194
|
readonly ops: OpNode[];
|
|
195
195
|
readonly tensors: Tensor[];
|
|
196
196
|
readonly outputs: number[];
|
|
197
|
+
readonly captures: Map<string, number>;
|
|
197
198
|
}
|
|
198
199
|
export declare function makeGraph(): Graph;
|
|
199
200
|
export declare function addTensor(g: Graph, shape: Shape, dtype: Dtype, source: number | null, site: CallSite | null): Tensor;
|
package/dist/ir.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"ir.d.ts","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAcA,MAAM,MAAM,KAAK,GAAG,KAAK,GAAG,KAAK,GAAG,MAAM,CAAA;AAC1C,MAAM,MAAM,KAAK,GAAG,SAAS,MAAM,EAAE,CAAA;AAIrC,MAAM,WAAW,MAAM;IACrB,QAAQ,CAAC,EAAE,EAAE,MAAM,CAAA;IACnB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IACrB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IAErB,QAAQ,CAAC,MAAM,EAAE,MAAM,GAAG,IAAI,CAAA;IAG9B,QAAQ,CAAC,IAAI,EAAE,QAAQ,GAAG,IAAI,CAAA;CAC/B;AAED,MAAM,WAAW,QAAQ;IACvB,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAA;IAEvB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAA;CACvB;AAQD,MAAM,MAAM,MAAM,GAGd;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAElD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAInD;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAGrE;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAC9D;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAG9D;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACxC;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACzC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGxC;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAC7C;IAAE,IAAI,EAAE,UAAU,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG5C;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,QAAQ,EAAE,KAAK,CAAA;CAAE,GAC5D;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,SAAS,MAAM,EAAE,CAAA;CAAE,GAMtE;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAErD;IAAE,IAAI,EAAE,gBAAgB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG7D;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,OAAO,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAC9E;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAGxD;IAAE,IAAI,EAAE,qBAAqB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvD;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAIpD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAKnE;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACnD;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGtD;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAMlE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GACxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GAOxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,MAAM,CAAA;CAAE,GAM5H;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAA;CAAE,GAGhF;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAGpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAEpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAElE;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,CAAA;AAI7D,MAAM,WAAW,KAAK;IACpB,QAAQ,CAAC,GAAG,EAAE,MAAM,EAAE,CAAA;IACtB,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAG1B,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;
|
|
1
|
+
{"version":3,"file":"ir.d.ts","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAcA,MAAM,MAAM,KAAK,GAAG,KAAK,GAAG,KAAK,GAAG,MAAM,CAAA;AAC1C,MAAM,MAAM,KAAK,GAAG,SAAS,MAAM,EAAE,CAAA;AAIrC,MAAM,WAAW,MAAM;IACrB,QAAQ,CAAC,EAAE,EAAE,MAAM,CAAA;IACnB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IACrB,QAAQ,CAAC,KAAK,EAAE,KAAK,CAAA;IAErB,QAAQ,CAAC,MAAM,EAAE,MAAM,GAAG,IAAI,CAAA;IAG9B,QAAQ,CAAC,IAAI,EAAE,QAAQ,GAAG,IAAI,CAAA;CAC/B;AAED,MAAM,WAAW,QAAQ;IACvB,QAAQ,CAAC,MAAM,EAAE,MAAM,CAAA;IAEvB,QAAQ,CAAC,KAAK,EAAE,MAAM,CAAA;CACvB;AAQD,MAAM,MAAM,MAAM,GAGd;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAElD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAA;CAAE,GAInD;IAAE,IAAI,EAAE,aAAa,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAGrE;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAClD;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAC9D;IAAE,IAAI,EAAE,YAAY,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE,GAG9D;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACxC;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACzC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,KAAK,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvC;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGxC;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAC7C;IAAE,IAAI,EAAE,UAAU,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG5C;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,QAAQ,EAAE,KAAK,CAAA;CAAE,GAC5D;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,SAAS,MAAM,EAAE,CAAA;CAAE,GAMtE;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAErD;IAAE,IAAI,EAAE,gBAAgB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAG7D;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,OAAO,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAC9E;IAAE,IAAI,EAAE,QAAQ,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAGxD;IAAE,IAAI,EAAE,qBAAqB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACvD;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAIpD;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,SAAS,EAAE,MAAM,CAAA;CAAE,GAKnE;IAAE,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GACnD;IAAE,IAAI,EAAE,SAAS,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAGtD;IAAE,IAAI,EAAE,OAAO,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAA;CAAE,GAMlE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GACxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,GAOxE;IAAE,IAAI,EAAE,eAAe,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,IAAI,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,MAAM,CAAA;CAAE,GAM5H;IAAE,IAAI,EAAE,kBAAkB,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,GAAG,EAAE,MAAM,CAAA;CAAE,GAGhF;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAGpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,WAAW,EAAE,KAAK,CAAA;CAAE,GAEpE;IAAE,IAAI,EAAE,cAAc,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,KAAK,EAAE,KAAK,CAAA;CAAE,GAElE;IAAE,IAAI,EAAE,WAAW,CAAC;IAAC,GAAG,EAAE,MAAM,CAAC;IAAC,CAAC,EAAE,MAAM,CAAC;IAAC,EAAE,EAAE,MAAM,CAAA;CAAE,CAAA;AAI7D,MAAM,WAAW,KAAK;IACpB,QAAQ,CAAC,GAAG,EAAE,MAAM,EAAE,CAAA;IACtB,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAG1B,QAAQ,CAAC,OAAO,EAAE,MAAM,EAAE,CAAA;IAI1B,QAAQ,CAAC,QAAQ,EAAE,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CACvC;AAED,wBAAgB,SAAS,IAAI,KAAK,CAEjC;AAGD,wBAAgB,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,GAAG,IAAI,EAAE,IAAI,EAAE,QAAQ,GAAG,IAAI,GAAG,MAAM,CAKpH;AAMD,wBAAgB,KAAK,CAAC,CAAC,SAAS,MAAM,CAAC,MAAM,CAAC,EAC5C,CAAC,EAAE,KAAK,EACR,IAAI,EAAE,CAAC,EACP,KAAK,EAAE,KAAK,EACZ,KAAK,EAAE,KAAK,EACZ,IAAI,EAAE,QAAQ,GAAG,IAAI,EACrB,MAAM,EAAE,IAAI,CAAC,OAAO,CAAC,MAAM,EAAE;IAAE,IAAI,EAAE,CAAC,CAAA;CAAE,CAAC,EAAE,MAAM,GAAG,KAAK,CAAC,GACzD,MAAM,CAMR;AAID,wBAAgB,WAAW,CAAC,MAAM,EAAE,MAAM,GAAG,QAAQ,CAIpD;AAID,wBAAgB,UAAU,CAAC,IAAI,EAAE,QAAQ,GAAG,MAAM,CAYjD"}
|
package/dist/ir.js
CHANGED
|
@@ -12,7 +12,7 @@
|
|
|
12
12
|
// Design intent: keep this file boring. No tracing logic, no shape inference,
|
|
13
13
|
// no codegen — those live in their own modules and consume `Graph` / `OpNode`.
|
|
14
14
|
export function makeGraph() {
|
|
15
|
-
return { ops: [], tensors: [], outputs: [] };
|
|
15
|
+
return { ops: [], tensors: [], outputs: [], captures: new Map() };
|
|
16
16
|
}
|
|
17
17
|
// Internal: register a fresh tensor in the graph and return its id.
|
|
18
18
|
export function addTensor(g, shape, dtype, source, site) {
|
package/dist/ir.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"ir.js","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAAA,uDAAuD;AACvD,EAAE;AACF,gFAAgF;AAChF,+EAA+E;AAC/E,8EAA8E;AAC9E,EAAE;AACF,0DAA0D;AAC1D,uCAAuC;AACvC,8EAA8E;AAC9E,wFAAwF;AACxF,EAAE;AACF,8EAA8E;AAC9E,+EAA+E;
|
|
1
|
+
{"version":3,"file":"ir.js","sourceRoot":"","sources":["../src/ir.ts"],"names":[],"mappings":"AAAA,uDAAuD;AACvD,EAAE;AACF,gFAAgF;AAChF,+EAA+E;AAC/E,8EAA8E;AAC9E,EAAE;AACF,0DAA0D;AAC1D,uCAAuC;AACvC,8EAA8E;AAC9E,wFAAwF;AACxF,EAAE;AACF,8EAA8E;AAC9E,+EAA+E;AAyI/E,MAAM,UAAU,SAAS;IACvB,OAAO,EAAE,GAAG,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,OAAO,EAAE,EAAE,EAAE,QAAQ,EAAE,IAAI,GAAG,EAAE,EAAE,CAAA;AACnE,CAAC;AAED,oEAAoE;AACpE,MAAM,UAAU,SAAS,CAAC,CAAQ,EAAE,KAAY,EAAE,KAAY,EAAE,MAAqB,EAAE,IAAqB;IAC1G,MAAM,EAAE,GAAG,CAAC,CAAC,OAAO,CAAC,MAAM,CAAA;IAC3B,MAAM,CAAC,GAAW,EAAE,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,CAAA;IACpD,CAAC,CAAC,OAAO,CAAC,IAAI,CAAC,CAAC,CAAC,CAAA;IACjB,OAAO,CAAC,CAAA;AACV,CAAC;AAED,kFAAkF;AAClF,0EAA0E;AAC1E,+EAA+E;AAC/E,kFAAkF;AAClF,MAAM,UAAU,KAAK,CACnB,CAAQ,EACR,IAAO,EACP,KAAY,EACZ,KAAY,EACZ,IAAqB,EACrB,MAA0D;IAE1D,MAAM,OAAO,GAAG,CAAC,CAAC,GAAG,CAAC,MAAM,CAAA;IAC5B,MAAM,GAAG,GAAG,SAAS,CAAC,CAAC,EAAE,KAAK,EAAE,KAAK,EAAE,OAAO,EAAE,IAAI,CAAC,CAAA;IACrD,MAAM,IAAI,GAAG,EAAE,IAAI,EAAE,GAAG,EAAE,GAAG,CAAC,EAAE,EAAE,GAAG,MAAM,EAAkC,CAAA;IAC7E,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;IAChB,OAAO,GAAG,CAAA;AACZ,CAAC;AAED,0EAA0E;AAC1E,iFAAiF;AACjF,MAAM,UAAU,WAAW,CAAC,MAAc;IACxC,+EAA+E;IAC/E,MAAM,KAAK,GAAG,CAAC,IAAI,KAAK,EAAE,CAAC,CAAC,KAAK,IAAI,EAAE,CAAA;IACvC,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,CAAA;AAC1B,CAAC;AAED,8EAA8E;AAC9E,2DAA2D;AAC3D,MAAM,UAAU,UAAU,CAAC,IAAc;IACvC,MAAM,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,CAAC,CAAA;IACpC,2EAA2E;IAC3E,iEAAiE;IACjE,MAAM,UAAU,GAAa,EAAE,CAAA;IAC/B,KAAK,MAAM,IAAI,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC;QAClC,IAAI,IAAI,CAAC,QAAQ,CAAC,kBAAkB,CAAC,IAAI,IAAI,CAAC,QAAQ,CAAC,qBAAqB,CAAC;YAAE,SAAQ;QACvF,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,EAAE,CAAC,CAAA;QAC5B,IAAI,UAAU,CAAC,MAAM,IAAI,CAAC;YAAE,MAAK;IACnC,CAAC;IACD,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC;QAAE,OAAO,IAAI,IAAI,CAAC,MAAM,yBAAyB,CAAA;IAC5E,OAAO,IAAI,IAAI,CAAC,MAAM,QAAQ,UAAU,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAA;AACzD,CAAC"}
|
package/dist/module.d.ts
CHANGED
|
@@ -1,4 +1,21 @@
|
|
|
1
1
|
import type { Tensor, Shape, Dtype } from './ir.js';
|
|
2
|
+
/** How a parameter's initial values are produced.
|
|
3
|
+
* - `'randn'` — Gaussian, with `scale` (default 0.02). The common case for
|
|
4
|
+
* weight matrices and embeddings.
|
|
5
|
+
* - `'zeros'` — fill with 0. Common for biases and LayerNorm beta.
|
|
6
|
+
* - `'ones'` — fill with 1. Common for LayerNorm gain.
|
|
7
|
+
* - Custom function — receives total element count and shape, returns the
|
|
8
|
+
* Float32Array. Use for fan-in scaling or any non-standard scheme.
|
|
9
|
+
*/
|
|
10
|
+
export type InitSpec = 'randn' | 'zeros' | 'ones' | ((size: number, shape: readonly number[]) => Float32Array);
|
|
11
|
+
export interface ParamOptions {
|
|
12
|
+
dtype?: Dtype;
|
|
13
|
+
/** Init kind. Default: `'randn'`. */
|
|
14
|
+
init?: InitSpec;
|
|
15
|
+
/** Std dev for `'randn'`. Default 0.02. Ignored for non-randn init. */
|
|
16
|
+
scale?: number;
|
|
17
|
+
}
|
|
18
|
+
type InitFn = (size: number, shape: readonly number[]) => Float32Array;
|
|
2
19
|
export declare abstract class Module {
|
|
3
20
|
/**
|
|
4
21
|
* Declare a learnable parameter at this module. Must be called from inside
|
|
@@ -6,16 +23,25 @@ export declare abstract class Module {
|
|
|
6
23
|
* that gets replaced with a real Tensor at compile time.
|
|
7
24
|
*
|
|
8
25
|
* The parameter's name is auto-derived from its property path in the model
|
|
9
|
-
* tree (e.g. `layers.0.attn.W_q`).
|
|
26
|
+
* tree (e.g. `layers.0.attn.W_q`). Init metadata travels with the param;
|
|
27
|
+
* call `compiled.uploadInitialParams()` to apply it after compile.
|
|
10
28
|
*/
|
|
11
|
-
protected param(shape: Shape,
|
|
29
|
+
protected param(shape: Shape, opts?: ParamOptions): Tensor;
|
|
30
|
+
}
|
|
31
|
+
export interface MaterializedParams {
|
|
32
|
+
/** Map from auto-derived path (e.g. `layers.0.attn.W_q`) to its Tensor. */
|
|
33
|
+
tensors: Record<string, Tensor>;
|
|
34
|
+
/** Init function per param path. Used by `uploadInitialParams`. */
|
|
35
|
+
initFns: Record<string, InitFn>;
|
|
12
36
|
}
|
|
13
37
|
/**
|
|
14
38
|
* Walk the module tree and replace every ParamSentinel with a real Tensor
|
|
15
39
|
* created via `paramInput(autoName, ...)`. Must be called inside an active
|
|
16
40
|
* trace context (paramInput appends to the current graph).
|
|
17
41
|
*
|
|
18
|
-
* Returns
|
|
42
|
+
* Returns the param tensors keyed by path, plus init functions for use by
|
|
43
|
+
* `uploadInitialParams`.
|
|
19
44
|
*/
|
|
20
|
-
export declare function materializeParams(root: Module):
|
|
45
|
+
export declare function materializeParams(root: Module): MaterializedParams;
|
|
46
|
+
export {};
|
|
21
47
|
//# sourceMappingURL=module.d.ts.map
|
package/dist/module.d.ts.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"file":"module.d.ts","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AA2BA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;
|
|
1
|
+
{"version":3,"file":"module.d.ts","sourceRoot":"","sources":["../src/module.ts"],"names":[],"mappings":"AA2BA,OAAO,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,KAAK,EAAE,MAAM,SAAS,CAAA;AAOnD;;;;;;;GAOG;AACH,MAAM,MAAM,QAAQ,GAChB,OAAO,GACP,OAAO,GACP,MAAM,GACN,CAAC,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,MAAM,EAAE,KAAK,YAAY,CAAC,CAAA;AAE9D,MAAM,WAAW,YAAY;IAC3B,KAAK,CAAC,EAAE,KAAK,CAAA;IACb,qCAAqC;IACrC,IAAI,CAAC,EAAE,QAAQ,CAAA;IACf,uEAAuE;IACvE,KAAK,CAAC,EAAE,MAAM,CAAA;CACf;AAED,KAAK,MAAM,GAAG,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,SAAS,MAAM,EAAE,KAAK,YAAY,CAAA;AA2CtE,8BAAsB,MAAM;IAC1B;;;;;;;;OAQG;IACH,SAAS,CAAC,KAAK,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,CAAC,EAAE,YAAY,GAAG,MAAM;CAK3D;AAMD,MAAM,WAAW,kBAAkB;IACjC,2EAA2E;IAC3E,OAAO,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;IAC/B,mEAAmE;IACnE,OAAO,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CAChC;AAED;;;;;;;GAOG;AACH,wBAAgB,iBAAiB,CAAC,IAAI,EAAE,MAAM,GAAG,kBAAkB,CAYlE"}
|