tapewasm 0.2.0 → 0.3.1

package/index.d.ts

@@ -3,6 +3,7 @@ export {
   AotSampler,
   AdviResult,
   CompiledTape,
+  SampleResult,
   compileTape,
   tapewasmVersion,
   setAotExports,
@@ -10,3 +11,30 @@ export {
   sharedMemory,
 } from "./pkg/tapewasm.js";
 export { default } from "./pkg/tapewasm.js";
+
+/** What `"auto"` uses: right for SpiderMonkey and JavaScriptCore. */
+export const RE_ROLL_ABOVE: number;
+/** For an engine that prefers straight-line, as V8 does. */
+export const V8_RE_ROLL_ABOVE: number;
+
+export interface Calibration {
+  /** Threshold to pass to `compileTape`. */
+  above: number;
+  /** False when the measurement could not run; `above` is the built-in value. */
+  measured: boolean;
+  prefersStraight?: boolean;
+  straightMs?: number;
+  loopedMs?: number;
+  error?: string;
+}
+
+/**
+ * Measure which shape this engine prefers and return a threshold for
+ * `compileTape`'s `reroll` argument. Call after `init()`. Cached; pass
+ * `{ force: true }` to measure again. Never throws — on failure it returns
+ * the built-in threshold.
+ */
+export function calibrateReroll(options?: { force?: boolean; rounds?: number }): Promise<number>;
+
+/** What the last `calibrateReroll` measured, or `null` before the first call. */
+export function lastCalibration(): Calibration | null;

package/index.js

@@ -7,10 +7,12 @@ export {
   AotSampler,
   AdviResult,
   CompiledTape,
+  SampleResult,
   compileTape,
   tapewasmVersion,
   setAotExports,
   clearAotExports,
   sharedMemory,
 } from "./pkg/tapewasm.js";
+export { calibrateReroll, lastCalibration, RE_ROLL_ABOVE, V8_RE_ROLL_ABOVE } from "./calibrate.js";
 export default init;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tapewasm",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "description": "Compile an autodiff tape to a WebAssembly module and sample it with nuts-rs, in the browser.",
   "type": "module",
   "main": "./index.js",

package/pkg/package.json

@@ -2,7 +2,7 @@
   "name": "tapewasm",
   "type": "module",
   "description": "Compile an autodiff tape to a wasm module and sample it with nuts-rs, in the browser. The wasm-bindgen API over tapewasm-codegen.",
-  "version": "0.2.0",
+  "version": "0.3.1",
   "license": "MIT OR Apache-2.0",
   "repository": {
     "type": "git",

package/pkg/tapewasm.d.ts

@@ -60,10 +60,19 @@ export class AotSampler {
      * spliced together: restarting Adam's moment estimates partway through
      * measurably converges to a worse optimum, so this is the only way to
      * watch a run progress without paying for that.
+     *
+     * `on_snapshot`, if given, is also called at each snapshot with the
+     * iteration, a copy of that `μ`, and the ELBO trace since the previous
+     * call — so a run inside a Worker can report as it goes rather than only
+     * once it returns.
      */
-    advi(init: Float64Array, num_iters: number, mc_samples: number, learning_rate: number, seed: bigint, snapshot_every: number): AdviResult;
+    advi(init: Float64Array, num_iters: number, mc_samples: number, learning_rate: number, seed: bigint, snapshot_every: number, on_snapshot?: ((iter: number, mu: Float64Array, elbo: Float64Array) => void) | null): AdviResult;
     /**
      * `[log_prob, d/dparam...]`.
+     *
+     * The evaluator behind it is built on the first call and reused, so
+     * calling this in a row costs one `scratch_init` copy rather than one per
+     * call — the same evaluator `sample` and `advi` keep for a whole run.
      */
     logProbGrad(params: Float64Array): Float64Array;
     /**
@@ -78,10 +87,28 @@ export class AotSampler {
      * an empty array to go without.
      */
     constructor(n_params: number, scratch_init: Float64Array, layout_id: number, param_names: string[]);
+    /**
+     * [`sample`](Self::sample) with each draw's sampler statistics beside it —
+     * what ArviZ keeps as `sample_stats`. `chain` only labels the run: a
+     * different seed is what keeps two chains apart.
+     */
+    sampleWithStats(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint, chain: number): SampleResult;
     /**
      * `num_warmup + num_draws` draws, row-major, `n_params` wide.
      */
     sample(init: Float64Array, num_warmup: number, num_draws: number, seed: bigint): Float64Array;
+    /**
+     * Estimate the diagonal metric from the gradients as well as the draws, as
+     * nuts-rs does by default. Off unless set, to match the reference
+     * posteriors; see [`nuts_settings`].
+     */
+    setGradBasedEstimate(on: boolean): void;
+    /**
+     * Aim warmup's step-size adaptation at this acceptance rate instead of
+     * nuts-rs's 0.8. Higher adapts a smaller step: fewer divergences on a hard
+     * geometry, more gradients per draw.
+     */
+    setTargetAccept(target: number): void;
     readonly nParams: number;
 }
 
@@ -102,6 +129,37 @@ export class CompiledTape {
     readonly wasm: Uint8Array;
 }
 
+/**
+ * Draws and, beside each, the sampler's statistics — warmup first, as
+ * [`AotSampler::sample`] returns them.
+ */
+export class SampleResult {
+    private constructor();
+    free(): void;
+    [Symbol.dispose](): void;
+    /**
+     * 1 where the draw's trajectory diverged.
+     */
+    readonly diverging: Uint8Array;
+    /**
+     * `num_warmup + num_draws` draws, row-major, `n_params` wide.
+     */
+    readonly draws: Float64Array;
+    /**
+     * The log density at each draw.
+     */
+    readonly lp: Float64Array;
+    /**
+     * Leapfrog steps the draw's trajectory took.
+     */
+    readonly numSteps: Uint32Array;
+    readonly stepSize: Float64Array;
+    /**
+     * 1 for the warmup draws.
+     */
+    readonly tuning: Uint8Array;
+}
+
 /**
  * Release the bound module's exports. The next draw will throw.
  */
@@ -117,8 +175,21 @@ export function clearAotExports(): void;
  *
  * The format is not an artifact and carries no compatibility promise: a tape
  * is written and consumed inside one call.
+ *
+ * `reroll` says when a vectorised statement becomes a wasm loop: `"auto"`
+ * (the default, straight-line below a size threshold), `"always"`, `"never"`,
+ * or a node count to re-roll past, written as a number.
+ *
+ * **Which is faster is an engine's preference, not the model's.** Measured
+ * across three engines, straight-line and re-rolled cross over between 8,026
+ * and 24,564 nodes in V8 and around 2,000 in SpiderMonkey and JavaScriptCore
+ * — an order of magnitude apart, so no single threshold serves all three. `"auto"` takes the
+ * lower one: near-optimal for two of the three, and up to 7.6x off on V8 for
+ * a trace between them. A caller that knows its engine passes the number.
+ *
+ * `"always"` is also the smallest module, often by an order of magnitude.
  */
-export function compileTape(tape: string): CompiledTape;
+export function compileTape(tape: string, reroll?: string | null): CompiledTape;
 
 /**
  * Forwards Rust panics to `console.error` with a message and backtrace rather
@@ -152,32 +223,43 @@ export interface InitOutput {
     readonly __wbg_adviresult_free: (a: number, b: number) => void;
     readonly __wbg_aotsampler_free: (a: number, b: number) => void;
     readonly __wbg_compiledtape_free: (a: number, b: number) => void;
+    readonly __wbg_sampleresult_free: (a: number, b: number) => void;
     readonly adviresult_elboTrace: (a: number) => [number, number];
     readonly adviresult_mu: (a: number) => [number, number];
     readonly adviresult_muSnapshots: (a: number) => [number, number];
     readonly adviresult_sigma: (a: number) => [number, number];
     readonly adviresult_snapshotIters: (a: number) => [number, number];
-    readonly aotsampler_advi: (a: number, b: number, c: number, d: number, e: number, f: number, g: bigint, h: number) => [number, number, number];
+    readonly aotsampler_advi: (a: number, b: number, c: number, d: number, e: number, f: number, g: bigint, h: number, i: number) => [number, number, number];
     readonly aotsampler_logProbGrad: (a: number, b: number, c: number) => [number, number, number, number];
     readonly aotsampler_nParams: (a: number) => number;
     readonly aotsampler_new: (a: number, b: number, c: number, d: number, e: number, f: number) => [number, number, number];
     readonly aotsampler_sample: (a: number, b: number, c: number, d: number, e: number, f: bigint) => [number, number, number, number];
+    readonly aotsampler_sampleWithStats: (a: number, b: number, c: number, d: number, e: number, f: bigint, g: number) => [number, number, number];
+    readonly aotsampler_setGradBasedEstimate: (a: number, b: number) => void;
+    readonly aotsampler_setTargetAccept: (a: number, b: number) => [number, number];
     readonly clearAotExports: () => void;
-    readonly compileTape: (a: number, b: number) => [number, number, number];
+    readonly compileTape: (a: number, b: number, c: number, d: number) => [number, number, number];
     readonly compiledtape_layoutId: (a: number) => number;
     readonly compiledtape_nParams: (a: number) => number;
     readonly compiledtape_scratchInit: (a: number) => [number, number];
     readonly compiledtape_wasm: (a: number) => [number, number];
     readonly init_panic_hook: () => void;
+    readonly sampleresult_diverging: (a: number) => [number, number];
+    readonly sampleresult_draws: (a: number) => [number, number];
+    readonly sampleresult_lp: (a: number) => [number, number];
+    readonly sampleresult_numSteps: (a: number) => [number, number];
+    readonly sampleresult_stepSize: (a: number) => [number, number];
+    readonly sampleresult_tuning: (a: number) => [number, number];
     readonly setAotExports: (a: any) => void;
     readonly sharedMemory: () => any;
     readonly tapewasmVersion: () => [number, number];
     readonly __wbindgen_malloc: (a: number, b: number) => number;
     readonly __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
     readonly __wbindgen_free: (a: number, b: number, c: number) => void;
+    readonly __wbindgen_exn_store: (a: number) => void;
+    readonly __externref_table_alloc: () => number;
     readonly __wbindgen_externrefs: WebAssembly.Table;
     readonly __externref_table_dealloc: (a: number) => void;
-    readonly __externref_table_alloc: () => number;
     readonly __wbindgen_start: () => void;
 }
 

package/pkg/tapewasm.js

@@ -1,9 +1,9 @@
 /* @ts-self-types="./tapewasm.d.ts" */
-import { aot_logp } from './snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js';
-import * as import1 from "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js"
-import * as import2 from "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js"
-import * as import3 from "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js"
-import * as import4 from "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js"
+import { aot_logp } from './snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js';
+import * as import1 from "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js"
+import * as import2 from "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js"
+import * as import3 from "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js"
+import * as import4 from "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js"
 
 
 /**
@@ -121,18 +121,24 @@ export class AotSampler {
      * spliced together: restarting Adam's moment estimates partway through
      * measurably converges to a worse optimum, so this is the only way to
      * watch a run progress without paying for that.
+     *
+     * `on_snapshot`, if given, is also called at each snapshot with the
+     * iteration, a copy of that `μ`, and the ELBO trace since the previous
+     * call — so a run inside a Worker can report as it goes rather than only
+     * once it returns.
      * @param {Float64Array} init
      * @param {number} num_iters
      * @param {number} mc_samples
      * @param {number} learning_rate
      * @param {bigint} seed
      * @param {number} snapshot_every
+     * @param {((iter: number, mu: Float64Array, elbo: Float64Array) => void) | null} [on_snapshot]
      * @returns {AdviResult}
      */
-    advi(init, num_iters, mc_samples, learning_rate, seed, snapshot_every) {
+    advi(init, num_iters, mc_samples, learning_rate, seed, snapshot_every, on_snapshot) {
         const ptr0 = passArrayF64ToWasm0(init, wasm.__wbindgen_malloc);
         const len0 = WASM_VECTOR_LEN;
-        const ret = wasm.aotsampler_advi(this.__wbg_ptr, ptr0, len0, num_iters, mc_samples, learning_rate, seed, snapshot_every);
+        const ret = wasm.aotsampler_advi(this.__wbg_ptr, ptr0, len0, num_iters, mc_samples, learning_rate, seed, snapshot_every, isLikeNone(on_snapshot) ? 0 : addToExternrefTable0(on_snapshot));
         if (ret[2]) {
             throw takeFromExternrefTable0(ret[1]);
         }
@@ -140,6 +146,10 @@ export class AotSampler {
     }
     /**
      * `[log_prob, d/dparam...]`.
+     *
+     * The evaluator behind it is built on the first call and reused, so
+     * calling this in a row costs one `scratch_init` copy rather than one per
+     * call — the same evaluator `sample` and `advi` keep for a whole run.
      * @param {Float64Array} params
      * @returns {Float64Array}
      */
@@ -189,6 +199,26 @@ export class AotSampler {
         AotSamplerFinalization.register(this, this.__wbg_ptr, this);
         return this;
     }
+    /**
+     * [`sample`](Self::sample) with each draw's sampler statistics beside it —
+     * what ArviZ keeps as `sample_stats`. `chain` only labels the run: a
+     * different seed is what keeps two chains apart.
+     * @param {Float64Array} init
+     * @param {number} num_warmup
+     * @param {number} num_draws
+     * @param {bigint} seed
+     * @param {number} chain
+     * @returns {SampleResult}
+     */
+    sampleWithStats(init, num_warmup, num_draws, seed, chain) {
+        const ptr0 = passArrayF64ToWasm0(init, wasm.__wbindgen_malloc);
+        const len0 = WASM_VECTOR_LEN;
+        const ret = wasm.aotsampler_sampleWithStats(this.__wbg_ptr, ptr0, len0, num_warmup, num_draws, seed, chain);
+        if (ret[2]) {
+            throw takeFromExternrefTable0(ret[1]);
+        }
+        return SampleResult.__wrap(ret[0]);
+    }
     /**
      * `num_warmup + num_draws` draws, row-major, `n_params` wide.
      * @param {Float64Array} init
@@ -208,6 +238,27 @@ export class AotSampler {
         wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
         return v2;
     }
+    /**
+     * Estimate the diagonal metric from the gradients as well as the draws, as
+     * nuts-rs does by default. Off unless set, to match the reference
+     * posteriors; see [`nuts_settings`].
+     * @param {boolean} on
+     */
+    setGradBasedEstimate(on) {
+        wasm.aotsampler_setGradBasedEstimate(this.__wbg_ptr, on);
+    }
+    /**
+     * Aim warmup's step-size adaptation at this acceptance rate instead of
+     * nuts-rs's 0.8. Higher adapts a smaller step: fewer divergences on a hard
+     * geometry, more gradients per draw.
+     * @param {number} target
+     */
+    setTargetAccept(target) {
+        const ret = wasm.aotsampler_setTargetAccept(this.__wbg_ptr, target);
+        if (ret[1]) {
+            throw takeFromExternrefTable0(ret[0]);
+        }
+    }
 }
 if (Symbol.dispose) AotSampler.prototype[Symbol.dispose] = AotSampler.prototype.free;
 
@@ -268,6 +319,89 @@ export class CompiledTape {
 }
 if (Symbol.dispose) CompiledTape.prototype[Symbol.dispose] = CompiledTape.prototype.free;
 
+/**
+ * Draws and, beside each, the sampler's statistics — warmup first, as
+ * [`AotSampler::sample`] returns them.
+ */
+export class SampleResult {
+    static __wrap(ptr) {
+        const obj = Object.create(SampleResult.prototype);
+        obj.__wbg_ptr = ptr;
+        SampleResultFinalization.register(obj, obj.__wbg_ptr, obj);
+        return obj;
+    }
+    __destroy_into_raw() {
+        const ptr = this.__wbg_ptr;
+        this.__wbg_ptr = 0;
+        SampleResultFinalization.unregister(this);
+        return ptr;
+    }
+    free() {
+        const ptr = this.__destroy_into_raw();
+        wasm.__wbg_sampleresult_free(ptr, 0);
+    }
+    /**
+     * 1 where the draw's trajectory diverged.
+     * @returns {Uint8Array}
+     */
+    get diverging() {
+        const ret = wasm.sampleresult_diverging(this.__wbg_ptr);
+        var v1 = getArrayU8FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 1, 1);
+        return v1;
+    }
+    /**
+     * `num_warmup + num_draws` draws, row-major, `n_params` wide.
+     * @returns {Float64Array}
+     */
+    get draws() {
+        const ret = wasm.sampleresult_draws(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * The log density at each draw.
+     * @returns {Float64Array}
+     */
+    get lp() {
+        const ret = wasm.sampleresult_lp(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * Leapfrog steps the draw's trajectory took.
+     * @returns {Uint32Array}
+     */
+    get numSteps() {
+        const ret = wasm.sampleresult_numSteps(this.__wbg_ptr);
+        var v1 = getArrayU32FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 4, 4);
+        return v1;
+    }
+    /**
+     * @returns {Float64Array}
+     */
+    get stepSize() {
+        const ret = wasm.sampleresult_stepSize(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * 1 for the warmup draws.
+     * @returns {Uint8Array}
+     */
+    get tuning() {
+        const ret = wasm.sampleresult_tuning(this.__wbg_ptr);
+        var v1 = getArrayU8FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 1, 1);
+        return v1;
+    }
+}
+if (Symbol.dispose) SampleResult.prototype[Symbol.dispose] = SampleResult.prototype.free;
+
 /**
  * Release the bound module's exports. The next draw will throw.
  */
@@ -285,13 +419,29 @@ export function clearAotExports() {
  *
  * The format is not an artifact and carries no compatibility promise: a tape
  * is written and consumed inside one call.
+ *
+ * `reroll` says when a vectorised statement becomes a wasm loop: `"auto"`
+ * (the default, straight-line below a size threshold), `"always"`, `"never"`,
+ * or a node count to re-roll past, written as a number.
+ *
+ * **Which is faster is an engine's preference, not the model's.** Measured
+ * across three engines, straight-line and re-rolled cross over between 8,026
+ * and 24,564 nodes in V8 and around 2,000 in SpiderMonkey and JavaScriptCore
+ * — an order of magnitude apart, so no single threshold serves all three. `"auto"` takes the
+ * lower one: near-optimal for two of the three, and up to 7.6x off on V8 for
+ * a trace between them. A caller that knows its engine passes the number.
+ *
+ * `"always"` is also the smallest module, often by an order of magnitude.
  * @param {string} tape
+ * @param {string | null} [reroll]
  * @returns {CompiledTape}
  */
-export function compileTape(tape) {
+export function compileTape(tape, reroll) {
     const ptr0 = passStringToWasm0(tape, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
     const len0 = WASM_VECTOR_LEN;
-    const ret = wasm.compileTape(ptr0, len0);
+    var ptr1 = isLikeNone(reroll) ? 0 : passStringToWasm0(reroll, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+    var len1 = WASM_VECTOR_LEN;
+    const ret = wasm.compileTape(ptr0, len0, ptr1, len1);
     if (ret[2]) {
         throw takeFromExternrefTable0(ret[1]);
     }
@@ -352,6 +502,13 @@ function __wbg_get_imports() {
             const ret = Error(getStringFromWasm0(arg0, arg1));
             return ret;
         },
+        __wbg___wbindgen_debug_string_0e68cf47c9cbd9b0: function(arg0, arg1) {
+            const ret = debugString(arg1);
+            const ptr1 = passStringToWasm0(ret, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+            const len1 = WASM_VECTOR_LEN;
+            getDataViewMemory0().setInt32(arg0 + 4 * 1, len1, true);
+            getDataViewMemory0().setInt32(arg0 + 4 * 0, ptr1, true);
+        },
         __wbg___wbindgen_memory_3f8442e22540244f: function() {
             const ret = wasm.memory;
             return ret;
@@ -367,10 +524,18 @@ function __wbg_get_imports() {
         __wbg___wbindgen_throw_5d9e815e6fdf150f: function(arg0, arg1) {
             throw new Error(getStringFromWasm0(arg0, arg1));
         },
-        __wbg_aot_logp_897c94d09e163a67: function(arg0, arg1, arg2, arg3) {
+        __wbg_aot_logp_88e25d7a78bf7680: function(arg0, arg1, arg2, arg3) {
             const ret = aot_logp(arg0 >>> 0, arg1 >>> 0, arg2 >>> 0, arg3 >>> 0);
             return ret;
         },
+        __wbg_call_d9fe056b59ca67cb: function() { return handleError(function (arg0, arg1, arg2, arg3, arg4, arg5, arg6) {
+            var v0 = getArrayF64FromWasm0(arg3, arg4).slice();
+            wasm.__wbindgen_free(arg3, arg4 * 8, 8);
+            var v1 = getArrayF64FromWasm0(arg5, arg6).slice();
+            wasm.__wbindgen_free(arg5, arg6 * 8, 8);
+            const ret = arg0.call(arg1, arg2, v0, v1);
+            return ret;
+        }, arguments); },
         __wbg_error_757e9472f8410341: function(arg0, arg1) {
             let deferred0_0;
             let deferred0_1;
@@ -406,10 +571,10 @@ function __wbg_get_imports() {
     return {
         __proto__: null,
         "./tapewasm_bg.js": import0,
-        "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js": import1,
-        "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js": import2,
-        "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js": import3,
-        "./snippets/tapewasm-3fbd96d1430b56d5/js/aot_bridge.js": import4,
+        "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js": import1,
+        "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js": import2,
+        "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js": import3,
+        "./snippets/tapewasm-cbf82a3d7bc31b1b/js/aot_bridge.js": import4,
     };
 }
 
@@ -422,6 +587,9 @@ const AotSamplerFinalization = (typeof FinalizationRegistry === 'undefined')
 const CompiledTapeFinalization = (typeof FinalizationRegistry === 'undefined')
     ? { register: () => {}, unregister: () => {} }
     : new FinalizationRegistry(ptr => wasm.__wbg_compiledtape_free(ptr, 1));
+const SampleResultFinalization = (typeof FinalizationRegistry === 'undefined')
+    ? { register: () => {}, unregister: () => {} }
+    : new FinalizationRegistry(ptr => wasm.__wbg_sampleresult_free(ptr, 1));
 
 function addToExternrefTable0(obj) {
     const idx = wasm.__externref_table_alloc();
@@ -429,11 +597,81 @@ function addToExternrefTable0(obj) {
     return idx;
 }
 
+function debugString(val) {
+    // primitive types
+    const type = typeof val;
+    if (type == 'number' || type == 'boolean' || val == null) {
+        return  `${val}`;
+    }
+    if (type == 'string') {
+        return `"${val}"`;
+    }
+    if (type == 'symbol') {
+        const description = val.description;
+        if (description == null) {
+            return 'Symbol';
+        } else {
+            return `Symbol(${description})`;
+        }
+    }
+    if (type == 'function') {
+        const name = val.name;
+        if (typeof name == 'string' && name.length > 0) {
+            return `Function(${name})`;
+        } else {
+            return 'Function';
+        }
+    }
+    // objects
+    if (Array.isArray(val)) {
+        const length = val.length;
+        let debug = '[';
+        if (length > 0) {
+            debug += debugString(val[0]);
+        }
+        for(let i = 1; i < length; i++) {
+            debug += ', ' + debugString(val[i]);
+        }
+        debug += ']';
+        return debug;
+    }
+    // Test for built-in
+    const builtInMatches = /\[object ([^\]]+)\]/.exec(toString.call(val));
+    let className;
+    if (builtInMatches && builtInMatches.length > 1) {
+        className = builtInMatches[1];
+    } else {
+        // Failed to match the standard '[object ClassName]'
+        return toString.call(val);
+    }
+    if (className == 'Object') {
+        // we're a user defined class or Object
+        // JSON.stringify avoids problems with cycles, and is generally much
+        // easier than looping through ownProperties of `val`.
+        try {
+            return 'Object(' + JSON.stringify(val) + ')';
+        } catch (_) {
+            return 'Object';
+        }
+    }
+    // errors
+    if (val instanceof Error) {
+        return `${val.name}: ${val.message}\n${val.stack}`;
+    }
+    // TODO we could test for more things here, like `Set`s and `Map`s.
+    return className;
+}
+
 function getArrayF64FromWasm0(ptr, len) {
     ptr = ptr >>> 0;
     return getFloat64ArrayMemory0().subarray(ptr / 8, ptr / 8 + len);
 }
 
+function getArrayU32FromWasm0(ptr, len) {
+    ptr = ptr >>> 0;
+    return getUint32ArrayMemory0().subarray(ptr / 4, ptr / 4 + len);
+}
+
 function getArrayU8FromWasm0(ptr, len) {
     ptr = ptr >>> 0;
     return getUint8ArrayMemory0().subarray(ptr / 1, ptr / 1 + len);
@@ -459,6 +697,14 @@ function getStringFromWasm0(ptr, len) {
     return decodeText(ptr >>> 0, len);
 }
 
+let cachedUint32ArrayMemory0 = null;
+function getUint32ArrayMemory0() {
+    if (cachedUint32ArrayMemory0 === null || cachedUint32ArrayMemory0.byteLength === 0) {
+        cachedUint32ArrayMemory0 = new Uint32Array(wasm.memory.buffer);
+    }
+    return cachedUint32ArrayMemory0;
+}
+
 let cachedUint8ArrayMemory0 = null;
 function getUint8ArrayMemory0() {
     if (cachedUint8ArrayMemory0 === null || cachedUint8ArrayMemory0.byteLength === 0) {
@@ -467,6 +713,15 @@ function getUint8ArrayMemory0() {
     return cachedUint8ArrayMemory0;
 }
 
+function handleError(f, args) {
+    try {
+        return f.apply(this, args);
+    } catch (e) {
+        const idx = addToExternrefTable0(e);
+        wasm.__wbindgen_exn_store(idx);
+    }
+}
+
 function isLikeNone(x) {
     return x === undefined || x === null;
 }
@@ -567,6 +822,7 @@ function __wbg_finalize_init(instance, module) {
     wasmModule = module;
     cachedDataViewMemory0 = null;
     cachedFloat64ArrayMemory0 = null;
+    cachedUint32ArrayMemory0 = null;
     cachedUint8ArrayMemory0 = null;
     wasm.__wbindgen_start();
     return wasm;

package/pkg/tapewasm_bg.wasm.d.ts

@@ -4,30 +4,41 @@ export const memory: WebAssembly.Memory;
 export const __wbg_adviresult_free: (a: number, b: number) => void;
 export const __wbg_aotsampler_free: (a: number, b: number) => void;
 export const __wbg_compiledtape_free: (a: number, b: number) => void;
+export const __wbg_sampleresult_free: (a: number, b: number) => void;
 export const adviresult_elboTrace: (a: number) => [number, number];
 export const adviresult_mu: (a: number) => [number, number];
 export const adviresult_muSnapshots: (a: number) => [number, number];
 export const adviresult_sigma: (a: number) => [number, number];
 export const adviresult_snapshotIters: (a: number) => [number, number];
-export const aotsampler_advi: (a: number, b: number, c: number, d: number, e: number, f: number, g: bigint, h: number) => [number, number, number];
+export const aotsampler_advi: (a: number, b: number, c: number, d: number, e: number, f: number, g: bigint, h: number, i: number) => [number, number, number];
 export const aotsampler_logProbGrad: (a: number, b: number, c: number) => [number, number, number, number];
 export const aotsampler_nParams: (a: number) => number;
 export const aotsampler_new: (a: number, b: number, c: number, d: number, e: number, f: number) => [number, number, number];
 export const aotsampler_sample: (a: number, b: number, c: number, d: number, e: number, f: bigint) => [number, number, number, number];
+export const aotsampler_sampleWithStats: (a: number, b: number, c: number, d: number, e: number, f: bigint, g: number) => [number, number, number];
+export const aotsampler_setGradBasedEstimate: (a: number, b: number) => void;
+export const aotsampler_setTargetAccept: (a: number, b: number) => [number, number];
 export const clearAotExports: () => void;
-export const compileTape: (a: number, b: number) => [number, number, number];
+export const compileTape: (a: number, b: number, c: number, d: number) => [number, number, number];
 export const compiledtape_layoutId: (a: number) => number;
 export const compiledtape_nParams: (a: number) => number;
 export const compiledtape_scratchInit: (a: number) => [number, number];
 export const compiledtape_wasm: (a: number) => [number, number];
 export const init_panic_hook: () => void;
+export const sampleresult_diverging: (a: number) => [number, number];
+export const sampleresult_draws: (a: number) => [number, number];
+export const sampleresult_lp: (a: number) => [number, number];
+export const sampleresult_numSteps: (a: number) => [number, number];
+export const sampleresult_stepSize: (a: number) => [number, number];
+export const sampleresult_tuning: (a: number) => [number, number];
 export const setAotExports: (a: any) => void;
 export const sharedMemory: () => any;
 export const tapewasmVersion: () => [number, number];
 export const __wbindgen_malloc: (a: number, b: number) => number;
 export const __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
 export const __wbindgen_free: (a: number, b: number, c: number) => void;
+export const __wbindgen_exn_store: (a: number) => void;
+export const __externref_table_alloc: () => number;
 export const __wbindgen_externrefs: WebAssembly.Table;
 export const __externref_table_dealloc: (a: number) => void;
-export const __externref_table_alloc: () => number;
 export const __wbindgen_start: () => void;