stanwasm 0.1.0 → 0.1.1
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 +18 -0
- package/package.json +1 -1
- package/pkg/README.md +4 -0
- package/pkg/stanwasm.d.ts +20 -62
- package/pkg/stanwasm.js +20 -62
- package/pkg/stanwasm_bg.wasm +0 -0
package/README.md
CHANGED
|
@@ -1,5 +1,10 @@
|
|
|
1
1
|
# stanwasm
|
|
2
2
|
|
|
3
|
+
[](https://www.npmjs.com/package/stanwasm)
|
|
4
|
+
[](https://crates.io/crates/stanwasm)
|
|
5
|
+
[](https://github.com/habakan/stanwasm/blob/main/docs/en/BENCHMARKS.md)
|
|
6
|
+
[](https://github.com/habakan/stanwasm/blob/main/LICENSE)
|
|
7
|
+
|
|
3
8
|
Stan probabilistic models compiled and sampled **entirely in the browser**.
|
|
4
9
|
Pure Rust compiled to WebAssembly, a single ~482 KB bundle (~180 KB gzipped),
|
|
5
10
|
with [nuts-rs](https://github.com/pymc-devs/nuts-rs) embedded as the sampler.
|
|
@@ -11,6 +16,19 @@ No server, no cmdstan, no round trip.
|
|
|
11
16
|
> [Stan Playground](https://github.com/flatironinstitute/stan-playground) —
|
|
12
17
|
> this is for browser-embedded use cases where those do not fit.
|
|
13
18
|
|
|
19
|
+
## Browser support
|
|
20
|
+
|
|
21
|
+
Chrome, Edge, Firefox, Safari and Node.js, plus every browser on iOS and
|
|
22
|
+
iPadOS. Verified under Playwright's three engines: Chromium 151, Firefox 153
|
|
23
|
+
and WebKit 26.5 all instantiate the module and sample.
|
|
24
|
+
|
|
25
|
+
Safari used to fail outright — WebKit rejects a module containing
|
|
26
|
+
[relaxed SIMD](https://github.com/WebAssembly/relaxed-simd) opcodes at
|
|
27
|
+
validation time, and `pulp` emitted them. This package ships the prebuilt wasm
|
|
28
|
+
with that resolved, so nothing is required of you; see
|
|
29
|
+
[the README](https://github.com/habakan/stanwasm#browser-support) for the
|
|
30
|
+
details and the one caveat that applies to the Rust crate.
|
|
31
|
+
|
|
14
32
|
## Install
|
|
15
33
|
|
|
16
34
|
```bash
|
package/package.json
CHANGED
package/pkg/README.md
CHANGED
|
@@ -1,5 +1,9 @@
|
|
|
1
1
|
# stanwasm
|
|
2
2
|
|
|
3
|
+
[](https://crates.io/crates/stanwasm)
|
|
4
|
+
[](https://www.npmjs.com/package/stanwasm)
|
|
5
|
+
[](https://github.com/habakan/stanwasm/blob/main/LICENSE)
|
|
6
|
+
|
|
3
7
|
The wasm-bindgen API behind stanwasm: parse, compile and sample a Stan model
|
|
4
8
|
entirely inside the browser.
|
|
5
9
|
|
package/pkg/stanwasm.d.ts
CHANGED
|
@@ -2,31 +2,25 @@
|
|
|
2
2
|
/* eslint-disable */
|
|
3
3
|
|
|
4
4
|
/**
|
|
5
|
-
* One compiled Stan model
|
|
6
|
-
*
|
|
7
|
-
* the `Compiled` and re-builds it from the retained AST afterwards, so
|
|
8
|
-
* the same `StanModel` instance can be sampled repeatedly.
|
|
5
|
+
* One compiled Stan model: the parsed AST plus a pre-traced `Compiled`.
|
|
6
|
+
* Sampling consumes the `Compiled` and rebuilds it from the AST after.
|
|
9
7
|
*/
|
|
10
8
|
export class StanModel {
|
|
11
9
|
free(): void;
|
|
12
10
|
[Symbol.dispose](): void;
|
|
13
11
|
/**
|
|
14
|
-
* AOT-compile this model to a self-contained wasm module.
|
|
15
|
-
*
|
|
16
|
-
* obtain an independent log_prob_grad runtime — useful for Web Workers
|
|
17
|
-
* or for inspection).
|
|
12
|
+
* AOT-compile this model to a self-contained wasm module. Pass the bytes to
|
|
13
|
+
* `WebAssembly.instantiate` for an independent log_prob_grad runtime.
|
|
18
14
|
*/
|
|
19
15
|
compileToWasm(): Uint8Array;
|
|
20
16
|
/**
|
|
21
|
-
* Constrained
|
|
22
|
-
* unconstrained draw
|
|
23
|
-
* in `paramNames()` order.
|
|
17
|
+
* Constrained `parameters` + `transformed parameters` for one
|
|
18
|
+
* unconstrained draw, flattened in `paramNames()` order.
|
|
24
19
|
*/
|
|
25
20
|
constrainDraw(unconstrained: Float64Array): Float64Array;
|
|
26
21
|
/**
|
|
27
|
-
* Stop step-sampling
|
|
28
|
-
*
|
|
29
|
-
* re-tracing, same as `sample()` does at the end of a run.
|
|
22
|
+
* Stop step-sampling (safe after it ended naturally too) and restore
|
|
23
|
+
* `logProbGrad`/`sample` by re-tracing.
|
|
30
24
|
*/
|
|
31
25
|
finishStepSampling(): void;
|
|
32
26
|
/**
|
|
@@ -35,18 +29,8 @@ export class StanModel {
|
|
|
35
29
|
*/
|
|
36
30
|
genQuantityNames(): string[];
|
|
37
31
|
/**
|
|
38
|
-
*
|
|
39
|
-
*
|
|
40
|
-
* shape `(n_draws, n_params)`; the result is `(n_draws, n_gen_quantities)`,
|
|
41
|
-
* row-major, in `genQuantityNames()` order. A single RNG stream (seeded
|
|
42
|
-
* by `seed`) is shared across all draws so repeated `_rng` calls don't
|
|
43
|
-
* repeat the same values draw-to-draw.
|
|
44
|
-
*
|
|
45
|
-
* Note: unlike `sampleViaAot`, there is no AOT-compiled counterpart of
|
|
46
|
-
* this method — `compileToWasm` only exports `log_prob_grad`. Generated
|
|
47
|
-
* quantities involve RNG and branching that the flat-tape AOT codegen
|
|
48
|
-
* doesn't model, and (running once per draw rather than once per NUTS
|
|
49
|
-
* leapfrog step) don't need it for performance.
|
|
32
|
+
* `generated quantities` over row-major `(n_draws, n_params)` draws.
|
|
33
|
+
* Result is row-major in `genQuantityNames()` order; one seeded RNG stream.
|
|
50
34
|
*/
|
|
51
35
|
generatedQuantities(draws: Float64Array, num_draws: number, seed: bigint): Float64Array;
|
|
52
36
|
/**
|
|
@@ -69,38 +53,18 @@ export class StanModel {
|
|
|
69
53
|
*/
|
|
70
54
|
sample(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): Float64Array;
|
|
71
55
|
/**
|
|
72
|
-
*
|
|
73
|
-
*
|
|
74
|
-
* V8 JITs the unrolled forward+backward pass in the AOT module, which
|
|
75
|
-
* can be substantially faster than the in-process tape replay used by
|
|
76
|
-
* `sample`. Both produce identical samples for a given seed.
|
|
56
|
+
* `sample` through a `setAotExports`-bound AOT wasm instead of tape replay;
|
|
57
|
+
* V8 JITs the unrolled pass. Identical samples for a given seed.
|
|
77
58
|
*/
|
|
78
59
|
sampleViaAot(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): Float64Array;
|
|
79
60
|
/**
|
|
80
|
-
* Start a
|
|
81
|
-
*
|
|
82
|
-
* the sampler's state alive in the `StanModel` instance so `stepDraw()`
|
|
83
|
-
* can advance it one draw at a time — genuinely watching the sampler
|
|
84
|
-
* work, not replaying an already-finished chain. Consumes the internal
|
|
85
|
-
* `Compiled` the same way `sample()` does; call `finishStepSampling()`
|
|
86
|
-
* (or exhaust `stepDraw()` up to `num_warmup + num_draws` calls, which
|
|
87
|
-
* does it automatically) before using `logProbGrad`/`sample` again.
|
|
61
|
+
* Start a NUTS run that `stepDraw()` advances one draw at a time. Consumes
|
|
62
|
+
* the `Compiled`: call `finishStepSampling()` before `logProbGrad`/`sample`.
|
|
88
63
|
*/
|
|
89
64
|
startStepSampling(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): void;
|
|
90
65
|
/**
|
|
91
|
-
* Advance
|
|
92
|
-
*
|
|
93
|
-
* then `1.0`/`0.0` for whether this draw was still in the warmup
|
|
94
|
-
* (tuning) phase, then `1.0`/`0.0` for whether it diverged, then the
|
|
95
|
-
* leapfrog `step_size` and `num_steps` nuts-rs actually used for this
|
|
96
|
-
* draw — these come straight out of nuts-rs's own dual-averaging
|
|
97
|
-
* adaptation and trajectory-length search, not anything this crate
|
|
98
|
-
* computes, so they're a way to show the real sampler internals at
|
|
99
|
-
* work rather than just the resulting draw. Once the requested
|
|
100
|
-
* `num_warmup + num_draws` draws have all been returned, this
|
|
101
|
-
* automatically restores `logProbGrad`/`sample` (by re-tracing, same as
|
|
102
|
-
* `sample()` does) and further calls fail until `startStepSampling` runs
|
|
103
|
-
* again.
|
|
66
|
+
* Advance one draw: `n_params` positions, tuning and diverging as `1.0`/`0.0`,
|
|
67
|
+
* then nuts-rs's own `step_size` and `num_steps`. Restores `sample` when done.
|
|
104
68
|
*/
|
|
105
69
|
stepDraw(): Float64Array;
|
|
106
70
|
/**
|
|
@@ -115,13 +79,8 @@ export class StanModel {
|
|
|
115
79
|
export function clearAotExports(): void;
|
|
116
80
|
|
|
117
81
|
/**
|
|
118
|
-
*
|
|
119
|
-
*
|
|
120
|
-
* instead, but a handful of internal-invariant panics remain, e.g. index
|
|
121
|
-
* out of bounds on a malformed AST) to `console.error` with a real message
|
|
122
|
-
* and backtrace, instead of an opaque `RuntimeError: unreachable`. The
|
|
123
|
-
* panicking call still traps the instance — this is diagnostics, not
|
|
124
|
-
* recovery — but it means a bug report can include what actually broke.
|
|
82
|
+
* Forwards Rust panics to `console.error` with a message and backtrace rather
|
|
83
|
+
* than an opaque `RuntimeError: unreachable`. Diagnostics: the instance still traps.
|
|
125
84
|
*/
|
|
126
85
|
export function init_panic_hook(): void;
|
|
127
86
|
|
|
@@ -132,9 +91,8 @@ export function init_panic_hook(): void;
|
|
|
132
91
|
export function setAotExports(exports: any): void;
|
|
133
92
|
|
|
134
93
|
/**
|
|
135
|
-
*
|
|
136
|
-
*
|
|
137
|
-
* AOT model so the two modules share buffers (zero-copy bridge).
|
|
94
|
+
* The linear memory backing this module. Pass as the `stan.memory` import when
|
|
95
|
+
* instantiating an AOT model so the two share buffers.
|
|
138
96
|
*/
|
|
139
97
|
export function sharedMemory(): any;
|
|
140
98
|
|
package/pkg/stanwasm.js
CHANGED
|
@@ -5,10 +5,8 @@ import * as import2 from "./snippets/stanwasm-2319b34998707c5e/js/aot_bridge.js"
|
|
|
5
5
|
|
|
6
6
|
|
|
7
7
|
/**
|
|
8
|
-
* One compiled Stan model
|
|
9
|
-
*
|
|
10
|
-
* the `Compiled` and re-builds it from the retained AST afterwards, so
|
|
11
|
-
* the same `StanModel` instance can be sampled repeatedly.
|
|
8
|
+
* One compiled Stan model: the parsed AST plus a pre-traced `Compiled`.
|
|
9
|
+
* Sampling consumes the `Compiled` and rebuilds it from the AST after.
|
|
12
10
|
*/
|
|
13
11
|
export class StanModel {
|
|
14
12
|
__destroy_into_raw() {
|
|
@@ -22,10 +20,8 @@ export class StanModel {
|
|
|
22
20
|
wasm.__wbg_stanmodel_free(ptr, 0);
|
|
23
21
|
}
|
|
24
22
|
/**
|
|
25
|
-
* AOT-compile this model to a self-contained wasm module.
|
|
26
|
-
*
|
|
27
|
-
* obtain an independent log_prob_grad runtime — useful for Web Workers
|
|
28
|
-
* or for inspection).
|
|
23
|
+
* AOT-compile this model to a self-contained wasm module. Pass the bytes to
|
|
24
|
+
* `WebAssembly.instantiate` for an independent log_prob_grad runtime.
|
|
29
25
|
* @returns {Uint8Array}
|
|
30
26
|
*/
|
|
31
27
|
compileToWasm() {
|
|
@@ -38,9 +34,8 @@ export class StanModel {
|
|
|
38
34
|
return v1;
|
|
39
35
|
}
|
|
40
36
|
/**
|
|
41
|
-
* Constrained
|
|
42
|
-
* unconstrained draw
|
|
43
|
-
* in `paramNames()` order.
|
|
37
|
+
* Constrained `parameters` + `transformed parameters` for one
|
|
38
|
+
* unconstrained draw, flattened in `paramNames()` order.
|
|
44
39
|
* @param {Float64Array} unconstrained
|
|
45
40
|
* @returns {Float64Array}
|
|
46
41
|
*/
|
|
@@ -56,9 +51,8 @@ export class StanModel {
|
|
|
56
51
|
return v2;
|
|
57
52
|
}
|
|
58
53
|
/**
|
|
59
|
-
* Stop step-sampling
|
|
60
|
-
*
|
|
61
|
-
* re-tracing, same as `sample()` does at the end of a run.
|
|
54
|
+
* Stop step-sampling (safe after it ended naturally too) and restore
|
|
55
|
+
* `logProbGrad`/`sample` by re-tracing.
|
|
62
56
|
*/
|
|
63
57
|
finishStepSampling() {
|
|
64
58
|
wasm.stanmodel_finishStepSampling(this.__wbg_ptr);
|
|
@@ -75,18 +69,8 @@ export class StanModel {
|
|
|
75
69
|
return v1;
|
|
76
70
|
}
|
|
77
71
|
/**
|
|
78
|
-
*
|
|
79
|
-
*
|
|
80
|
-
* shape `(n_draws, n_params)`; the result is `(n_draws, n_gen_quantities)`,
|
|
81
|
-
* row-major, in `genQuantityNames()` order. A single RNG stream (seeded
|
|
82
|
-
* by `seed`) is shared across all draws so repeated `_rng` calls don't
|
|
83
|
-
* repeat the same values draw-to-draw.
|
|
84
|
-
*
|
|
85
|
-
* Note: unlike `sampleViaAot`, there is no AOT-compiled counterpart of
|
|
86
|
-
* this method — `compileToWasm` only exports `log_prob_grad`. Generated
|
|
87
|
-
* quantities involve RNG and branching that the flat-tape AOT codegen
|
|
88
|
-
* doesn't model, and (running once per draw rather than once per NUTS
|
|
89
|
-
* leapfrog step) don't need it for performance.
|
|
72
|
+
* `generated quantities` over row-major `(n_draws, n_params)` draws.
|
|
73
|
+
* Result is row-major in `genQuantityNames()` order; one seeded RNG stream.
|
|
90
74
|
* @param {Float64Array} draws
|
|
91
75
|
* @param {number} num_draws
|
|
92
76
|
* @param {bigint} seed
|
|
@@ -178,11 +162,8 @@ export class StanModel {
|
|
|
178
162
|
return v2;
|
|
179
163
|
}
|
|
180
164
|
/**
|
|
181
|
-
*
|
|
182
|
-
*
|
|
183
|
-
* V8 JITs the unrolled forward+backward pass in the AOT module, which
|
|
184
|
-
* can be substantially faster than the in-process tape replay used by
|
|
185
|
-
* `sample`. Both produce identical samples for a given seed.
|
|
165
|
+
* `sample` through a `setAotExports`-bound AOT wasm instead of tape replay;
|
|
166
|
+
* V8 JITs the unrolled pass. Identical samples for a given seed.
|
|
186
167
|
* @param {Float64Array} init
|
|
187
168
|
* @param {number} num_warmup
|
|
188
169
|
* @param {number} num_draws
|
|
@@ -201,14 +182,8 @@ export class StanModel {
|
|
|
201
182
|
return v2;
|
|
202
183
|
}
|
|
203
184
|
/**
|
|
204
|
-
* Start a
|
|
205
|
-
*
|
|
206
|
-
* the sampler's state alive in the `StanModel` instance so `stepDraw()`
|
|
207
|
-
* can advance it one draw at a time — genuinely watching the sampler
|
|
208
|
-
* work, not replaying an already-finished chain. Consumes the internal
|
|
209
|
-
* `Compiled` the same way `sample()` does; call `finishStepSampling()`
|
|
210
|
-
* (or exhaust `stepDraw()` up to `num_warmup + num_draws` calls, which
|
|
211
|
-
* does it automatically) before using `logProbGrad`/`sample` again.
|
|
185
|
+
* Start a NUTS run that `stepDraw()` advances one draw at a time. Consumes
|
|
186
|
+
* the `Compiled`: call `finishStepSampling()` before `logProbGrad`/`sample`.
|
|
212
187
|
* @param {Float64Array} init
|
|
213
188
|
* @param {number} num_warmup
|
|
214
189
|
* @param {number} num_draws
|
|
@@ -223,19 +198,8 @@ export class StanModel {
|
|
|
223
198
|
}
|
|
224
199
|
}
|
|
225
200
|
/**
|
|
226
|
-
* Advance
|
|
227
|
-
*
|
|
228
|
-
* then `1.0`/`0.0` for whether this draw was still in the warmup
|
|
229
|
-
* (tuning) phase, then `1.0`/`0.0` for whether it diverged, then the
|
|
230
|
-
* leapfrog `step_size` and `num_steps` nuts-rs actually used for this
|
|
231
|
-
* draw — these come straight out of nuts-rs's own dual-averaging
|
|
232
|
-
* adaptation and trajectory-length search, not anything this crate
|
|
233
|
-
* computes, so they're a way to show the real sampler internals at
|
|
234
|
-
* work rather than just the resulting draw. Once the requested
|
|
235
|
-
* `num_warmup + num_draws` draws have all been returned, this
|
|
236
|
-
* automatically restores `logProbGrad`/`sample` (by re-tracing, same as
|
|
237
|
-
* `sample()` does) and further calls fail until `startStepSampling` runs
|
|
238
|
-
* again.
|
|
201
|
+
* Advance one draw: `n_params` positions, tuning and diverging as `1.0`/`0.0`,
|
|
202
|
+
* then nuts-rs's own `step_size` and `num_steps`. Restores `sample` when done.
|
|
239
203
|
* @returns {Float64Array}
|
|
240
204
|
*/
|
|
241
205
|
stepDraw() {
|
|
@@ -258,13 +222,8 @@ export function clearAotExports() {
|
|
|
258
222
|
}
|
|
259
223
|
|
|
260
224
|
/**
|
|
261
|
-
*
|
|
262
|
-
*
|
|
263
|
-
* instead, but a handful of internal-invariant panics remain, e.g. index
|
|
264
|
-
* out of bounds on a malformed AST) to `console.error` with a real message
|
|
265
|
-
* and backtrace, instead of an opaque `RuntimeError: unreachable`. The
|
|
266
|
-
* panicking call still traps the instance — this is diagnostics, not
|
|
267
|
-
* recovery — but it means a bug report can include what actually broke.
|
|
225
|
+
* Forwards Rust panics to `console.error` with a message and backtrace rather
|
|
226
|
+
* than an opaque `RuntimeError: unreachable`. Diagnostics: the instance still traps.
|
|
268
227
|
*/
|
|
269
228
|
export function init_panic_hook() {
|
|
270
229
|
wasm.init_panic_hook();
|
|
@@ -280,9 +239,8 @@ export function setAotExports(exports) {
|
|
|
280
239
|
}
|
|
281
240
|
|
|
282
241
|
/**
|
|
283
|
-
*
|
|
284
|
-
*
|
|
285
|
-
* AOT model so the two modules share buffers (zero-copy bridge).
|
|
242
|
+
* The linear memory backing this module. Pass as the `stan.memory` import when
|
|
243
|
+
* instantiating an AOT model so the two share buffers.
|
|
286
244
|
* @returns {any}
|
|
287
245
|
*/
|
|
288
246
|
export function sharedMemory() {
|
package/pkg/stanwasm_bg.wasm
CHANGED
|
Binary file
|