stanwasm 0.1.0

package/pkg/stanwasm.d.ts

@@ -0,0 +1,195 @@
+/* tslint:disable */
+/* eslint-disable */
+
+/**
+ * One compiled Stan model. Holds both the parsed AST (`Model`) and a
+ * pre-traced `Compiled` for fast log-prob evaluation. Sampling consumes
+ * the `Compiled` and re-builds it from the retained AST afterwards, so
+ * the same `StanModel` instance can be sampled repeatedly.
+ */
+export class StanModel {
+    free(): void;
+    [Symbol.dispose](): void;
+    /**
+     * AOT-compile this model to a self-contained wasm module. Returns the
+     * wasm bytes (callers can pass these to `WebAssembly.instantiate` to
+     * obtain an independent log_prob_grad runtime — useful for Web Workers
+     * or for inspection).
+     */
+    compileToWasm(): Uint8Array;
+    /**
+     * Constrained values of `parameters` + `transformed parameters` for one
+     * unconstrained draw (e.g. one row out of `sample()`'s output), flattened
+     * in `paramNames()` order.
+     */
+    constrainDraw(unconstrained: Float64Array): Float64Array;
+    /**
+     * Stop step-sampling early (or clean up after it finished naturally —
+     * safe to call either way) and restore `logProbGrad`/`sample` by
+     * re-tracing, same as `sample()` does at the end of a run.
+     */
+    finishStepSampling(): void;
+    /**
+     * Names of the top-level `generated quantities` declarations, flattened
+     * the same way as `paramNames()`.
+     */
+    genQuantityNames(): string[];
+    /**
+     * Evaluate `generated quantities` for a batch of unconstrained draws
+     * (e.g. `sample()`'s output). `draws` is a flat row-major buffer of
+     * shape `(n_draws, n_params)`; the result is `(n_draws, n_gen_quantities)`,
+     * row-major, in `genQuantityNames()` order. A single RNG stream (seeded
+     * by `seed`) is shared across all draws so repeated `_rng` calls don't
+     * repeat the same values draw-to-draw.
+     *
+     * Note: unlike `sampleViaAot`, there is no AOT-compiled counterpart of
+     * this method — `compileToWasm` only exports `log_prob_grad`. Generated
+     * quantities involve RNG and branching that the flat-tape AOT codegen
+     * doesn't model, and (running once per draw rather than once per NUTS
+     * leapfrog step) don't need it for performance.
+     */
+    generatedQuantities(draws: Float64Array, num_draws: number, seed: bigint): Float64Array;
+    /**
+     * Evaluate log_prob and gradient at `params`. Returns a flat array of
+     * length `n_params + 1`: the log-prob is at index 0, gradients follow.
+     */
+    logProbGrad(params: Float64Array): Float64Array;
+    /**
+     * Parse `stan_src`, bind `data_json`, trace the model on the autodiff
+     * tape, and return a handle ready for sampling.
+     */
+    constructor(stan_src: string, data_json: string);
+    /**
+     * Constrained parameter names (parameters then transformed parameters).
+     */
+    paramNames(): string[];
+    /**
+     * Run NUTS sampling. Returns a flat row-major buffer of shape
+     * `(num_warmup + num_draws) × n_params`. Tuning draws come first.
+     */
+    sample(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): Float64Array;
+    /**
+     * Same as `sample`, but evaluates `log_prob_grad` through a
+     * pre-instantiated AOT-compiled model wasm bound via `setAotExports`.
+     * V8 JITs the unrolled forward+backward pass in the AOT module, which
+     * can be substantially faster than the in-process tape replay used by
+     * `sample`. Both produce identical samples for a given seed.
+     */
+    sampleViaAot(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): Float64Array;
+    /**
+     * Start a step-by-step NUTS run: unlike `sample()`, which runs the whole
+     * chain inside one wasm call and returns only at the end, this leaves
+     * the sampler's state alive in the `StanModel` instance so `stepDraw()`
+     * can advance it one draw at a time — genuinely watching the sampler
+     * work, not replaying an already-finished chain. Consumes the internal
+     * `Compiled` the same way `sample()` does; call `finishStepSampling()`
+     * (or exhaust `stepDraw()` up to `num_warmup + num_draws` calls, which
+     * does it automatically) before using `logProbGrad`/`sample` again.
+     */
+    startStepSampling(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): void;
+    /**
+     * Advance the step-sampling chain started by `startStepSampling` by
+     * exactly one draw. Returns a flat array: `n_params` position values,
+     * then `1.0`/`0.0` for whether this draw was still in the warmup
+     * (tuning) phase, then `1.0`/`0.0` for whether it diverged, then the
+     * leapfrog `step_size` and `num_steps` nuts-rs actually used for this
+     * draw — these come straight out of nuts-rs's own dual-averaging
+     * adaptation and trajectory-length search, not anything this crate
+     * computes, so they're a way to show the real sampler internals at
+     * work rather than just the resulting draw. Once the requested
+     * `num_warmup + num_draws` draws have all been returned, this
+     * automatically restores `logProbGrad`/`sample` (by re-tracing, same as
+     * `sample()` does) and further calls fail until `startStepSampling` runs
+     * again.
+     */
+    stepDraw(): Float64Array;
+    /**
+     * Number of unconstrained parameters.
+     */
+    readonly n_params: number;
+}
+
+/**
+ * Release the bound AOT exports. The next `sampleViaAot` call will throw.
+ */
+export function clearAotExports(): void;
+
+/**
+ * Runs once when the wasm module is instantiated. Forwards Rust panics
+ * (Stan-typo'd names and invalid RNG parameters are now clean `JsError`s
+ * instead, but a handful of internal-invariant panics remain, e.g. index
+ * out of bounds on a malformed AST) to `console.error` with a real message
+ * and backtrace, instead of an opaque `RuntimeError: unreachable`. The
+ * panicking call still traps the instance — this is diagnostics, not
+ * recovery — but it means a bug report can include what actually broke.
+ */
+export function init_panic_hook(): void;
+
+/**
+ * Bind a freshly-instantiated AOT model wasm's exports so subsequent
+ * `sampleViaAot` calls dispatch through it. Pass `instance.exports`.
+ */
+export function setAotExports(exports: any): void;
+
+/**
+ * Returns the linear memory backing this wasm module. Pass to
+ * `WebAssembly.instantiate` as the `stan.memory` import when bringing up an
+ * AOT model so the two modules share buffers (zero-copy bridge).
+ */
+export function sharedMemory(): any;
+
+export function version(): string;
+
+export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
+
+export interface InitOutput {
+    readonly memory: WebAssembly.Memory;
+    readonly __wbg_stanmodel_free: (a: number, b: number) => void;
+    readonly setAotExports: (a: any) => void;
+    readonly stanmodel_compileToWasm: (a: number) => [number, number, number, number];
+    readonly stanmodel_constrainDraw: (a: number, b: number, c: number) => [number, number, number, number];
+    readonly stanmodel_finishStepSampling: (a: number) => void;
+    readonly stanmodel_genQuantityNames: (a: number) => [number, number];
+    readonly stanmodel_generatedQuantities: (a: number, b: number, c: number, d: number, e: bigint) => [number, number, number, number];
+    readonly stanmodel_logProbGrad: (a: number, b: number, c: number) => [number, number, number, number];
+    readonly stanmodel_n_params: (a: number) => number;
+    readonly stanmodel_new: (a: number, b: number, c: number, d: number) => [number, number, number];
+    readonly stanmodel_paramNames: (a: number) => [number, number];
+    readonly stanmodel_sample: (a: number, b: number, c: number, d: number, e: number, f: bigint) => [number, number, number, number];
+    readonly stanmodel_sampleViaAot: (a: number, b: number, c: number, d: number, e: number, f: bigint) => [number, number, number, number];
+    readonly stanmodel_startStepSampling: (a: number, b: number, c: number, d: number, e: number, f: bigint) => [number, number];
+    readonly stanmodel_stepDraw: (a: number) => [number, number, number, number];
+    readonly version: () => [number, number];
+    readonly clearAotExports: () => void;
+    readonly init_panic_hook: () => void;
+    readonly sharedMemory: () => any;
+    readonly __wbindgen_free: (a: number, b: number, c: number) => void;
+    readonly __wbindgen_malloc: (a: number, b: number) => number;
+    readonly __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
+    readonly __wbindgen_externrefs: WebAssembly.Table;
+    readonly __externref_table_dealloc: (a: number) => void;
+    readonly __externref_drop_slice: (a: number, b: number) => void;
+    readonly __wbindgen_start: () => void;
+}
+
+export type SyncInitInput = BufferSource | WebAssembly.Module;
+
+/**
+ * Instantiates the given `module`, which can either be bytes or
+ * a precompiled `WebAssembly.Module`.
+ *
+ * @param {{ module: SyncInitInput }} module - Passing `SyncInitInput` directly is deprecated.
+ *
+ * @returns {InitOutput}
+ */
+export function initSync(module: { module: SyncInitInput } | SyncInitInput): InitOutput;
+
+/**
+ * If `module_or_path` is {RequestInfo} or {URL}, makes a request and
+ * for everything else, calls `WebAssembly.instantiate` directly.
+ *
+ * @param {{ module_or_path: InitInput | Promise<InitInput> }} module_or_path - Passing `InitInput` directly is deprecated.
+ *
+ * @returns {Promise<InitOutput>}
+ */
+export default function __wbg_init (module_or_path?: { module_or_path: InitInput | Promise<InitInput> } | InitInput | Promise<InitInput>): Promise<InitOutput>;