tapewasm 0.3.0 → 0.3.2

package/index.d.ts

@@ -11,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

@@ -14,4 +14,11 @@ export {
   clearAotExports,
   sharedMemory,
 } from "./pkg/tapewasm.js";
+export {
+  calibrateReroll,
+  compileTapeCalibrated,
+  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.3.0",
+  "version": "0.3.2",
   "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.3.0",
+  "version": "0.3.2",
   "license": "MIT OR Apache-2.0",
   "repository": {
     "type": "git",

package/pkg/snippets/{tapewasm-6c5813ce65cccbcf → tapewasm-c400eec12145fd93}/js/aot_bridge.js

@@ -1,11 +1,16 @@
 // Bridge between tapewasm.wasm and a per-model compiled wasm, which imports
-// tapewasm's memory and exports `log_prob_grad`, `tapewasm_layout_id` and
-// `tapewasm_abi_version`.
+// tapewasm's memory and exports `log_prob_grad`, `tapewasm_layout_id`,
+// `tapewasm_abi_version` and `tapewasm_n_outputs` — plus `evaluate` when the
+// tape named any outputs.
 //
 // The binding is per page while the scratch buffer belongs to one model, so
 // the sampler reads the id back to refuse a mismatched pair.
 
 let aotLogProbGrad = null;
+let aotEvaluate = null;
+// How many values `evaluate` writes. 0 when the tape named none, and for a
+// module built before outputs existed, which exports neither.
+let aotNOutputs = 0;
 // NaN means nothing is bound, or no id is exported; no u32 id collides with it.
 let aotLayoutId = NaN;
 // Likewise: a module from before the global existed reads as unknown, not as 0.
@@ -13,6 +18,9 @@ let aotAbiVersion = NaN;
 
 export function set_aot_exports(exports) {
   aotLogProbGrad = exports.log_prob_grad;
+  aotEvaluate = exports.evaluate ?? null;
+  const n = exports.tapewasm_n_outputs;
+  aotNOutputs = n ? n.value >>> 0 : 0;
   const g = exports.tapewasm_layout_id;
   aotLayoutId = g ? g.value >>> 0 : NaN;
   const v = exports.tapewasm_abi_version;
@@ -21,6 +29,8 @@ export function set_aot_exports(exports) {
 
 export function clear_aot_exports() {
   aotLogProbGrad = null;
+  aotEvaluate = null;
+  aotNOutputs = 0;
   aotLayoutId = NaN;
   aotAbiVersion = NaN;
 }
@@ -39,3 +49,17 @@ export function aot_logp(paramsPtr, gradsPtr, nParams, scratchPtr) {
   }
   return aotLogProbGrad(paramsPtr, gradsPtr, nParams, scratchPtr);
 }
+
+export function aot_n_outputs() {
+  return aotNOutputs;
+}
+
+export function aot_evaluate(paramsPtr, outPtr, nParams, scratchPtr) {
+  if (!aotEvaluate) {
+    throw new Error(
+      "the bound module exports no evaluate — compile the tape with an " +
+        "`outputs` line naming what it should report",
+    );
+  }
+  return aotEvaluate(paramsPtr, outPtr, nParams, scratchPtr);
+}

package/pkg/tapewasm.d.ts

@@ -67,8 +67,25 @@ export class AotSampler {
      * once it returns.
      */
     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;
+    /**
+     * What the module's `outputs` named, at one point in the parameter space.
+     *
+     * The forward pass alone, so this is the way to a pointwise
+     * log-likelihood — one call per draw, a term per observation — or to a
+     * deterministic quantity the density does not return. It shares the
+     * scratch buffer with `logProbGrad`, whose own call recomputes what it
+     * needs, so the two interleave freely.
+     *
+     * Throws when the bound module was compiled from a tape that named no
+     * outputs.
+     */
+    evaluate(params: Float64Array): Float64Array;
     /**
      * `[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;
     /**
@@ -116,6 +133,11 @@ export class CompiledTape {
     free(): void;
     [Symbol.dispose](): void;
     readonly layoutId: number;
+    /**
+     * How many values `AotSampler.evaluate` returns — what the tape's
+     * `outputs` line named, and 0 when it had none.
+     */
+    readonly nOutputs: number;
     readonly nParams: number;
     readonly scratchInit: Float64Array;
     /**
@@ -172,12 +194,22 @@ export function clearAotExports(): void;
  * The format is not an artifact and carries no compatibility promise: a tape
  * is written and consumed inside one call.
  *
+ * An `outputs` line names nodes the module reports through
+ * `AotSampler.evaluate` — a pointwise log-likelihood's per-observation terms,
+ * or a deterministic quantity. Without one the module is what it always was.
+ *
  * `reroll` says when a vectorised statement becomes a wasm loop: `"auto"`
- * (the default, straight-line below a size threshold), `"always"` or
- * `"never"`. Which is faster is an engine's preference, not the model's —
- * on one real model straight-line was faster on V8 and slower on
- * SpiderMonkey and JavaScriptCore — and `"always"` is also the smallest
- * module, often by an order of magnitude.
+ * (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, reroll?: string | null): CompiledTape;
 
@@ -220,6 +252,7 @@ export interface InitOutput {
     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, i: number) => [number, number, number];
+    readonly aotsampler_evaluate: (a: number, b: number, c: number) => [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];
@@ -230,6 +263,7 @@ export interface InitOutput {
     readonly clearAotExports: () => void;
     readonly compileTape: (a: number, b: number, c: number, d: number) => [number, number, number];
     readonly compiledtape_layoutId: (a: number) => number;
+    readonly compiledtape_nOutputs: (a: number) => number;
     readonly compiledtape_nParams: (a: number) => number;
     readonly compiledtape_scratchInit: (a: number) => [number, number];
     readonly compiledtape_wasm: (a: number) => [number, number];
@@ -245,10 +279,10 @@ export interface InitOutput {
     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 __wbindgen_free: (a: number, b: number, c: number) => void;
     readonly __externref_table_dealloc: (a: number) => void;
     readonly __wbindgen_start: () => void;
 }

package/pkg/tapewasm.js

@@ -1,9 +1,10 @@
 /* @ts-self-types="./tapewasm.d.ts" */
-import { aot_logp } from './snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js';
-import * as import1 from "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js"
-import * as import2 from "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js"
-import * as import3 from "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js"
-import * as import4 from "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js"
+import { aot_evaluate, aot_logp } from './snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js';
+import * as import1 from "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js"
+import * as import2 from "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js"
+import * as import3 from "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js"
+import * as import4 from "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js"
+import * as import5 from "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js"
 
 
 /**
@@ -144,8 +145,37 @@ export class AotSampler {
         }
         return AdviResult.__wrap(ret[0]);
     }
+    /**
+     * What the module's `outputs` named, at one point in the parameter space.
+     *
+     * The forward pass alone, so this is the way to a pointwise
+     * log-likelihood — one call per draw, a term per observation — or to a
+     * deterministic quantity the density does not return. It shares the
+     * scratch buffer with `logProbGrad`, whose own call recomputes what it
+     * needs, so the two interleave freely.
+     *
+     * Throws when the bound module was compiled from a tape that named no
+     * outputs.
+     * @param {Float64Array} params
+     * @returns {Float64Array}
+     */
+    evaluate(params) {
+        const ptr0 = passArrayF64ToWasm0(params, wasm.__wbindgen_malloc);
+        const len0 = WASM_VECTOR_LEN;
+        const ret = wasm.aotsampler_evaluate(this.__wbg_ptr, ptr0, len0);
+        if (ret[3]) {
+            throw takeFromExternrefTable0(ret[2]);
+        }
+        var v2 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v2;
+    }
     /**
      * `[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}
      */
@@ -285,6 +315,15 @@ export class CompiledTape {
         const ret = wasm.compiledtape_layoutId(this.__wbg_ptr);
         return ret >>> 0;
     }
+    /**
+     * How many values `AotSampler.evaluate` returns — what the tape's
+     * `outputs` line named, and 0 when it had none.
+     * @returns {number}
+     */
+    get nOutputs() {
+        const ret = wasm.compiledtape_nOutputs(this.__wbg_ptr);
+        return ret >>> 0;
+    }
     /**
      * @returns {number}
      */
@@ -416,12 +455,22 @@ export function clearAotExports() {
  * The format is not an artifact and carries no compatibility promise: a tape
  * is written and consumed inside one call.
  *
+ * An `outputs` line names nodes the module reports through
+ * `AotSampler.evaluate` — a pointwise log-likelihood's per-observation terms,
+ * or a deterministic quantity. Without one the module is what it always was.
+ *
  * `reroll` says when a vectorised statement becomes a wasm loop: `"auto"`
- * (the default, straight-line below a size threshold), `"always"` or
- * `"never"`. Which is faster is an engine's preference, not the model's —
- * on one real model straight-line was faster on V8 and slower on
- * SpiderMonkey and JavaScriptCore — and `"always"` is also the smallest
- * module, often by an order of magnitude.
+ * (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}
@@ -514,11 +563,15 @@ function __wbg_get_imports() {
         __wbg___wbindgen_throw_5d9e815e6fdf150f: function(arg0, arg1) {
             throw new Error(getStringFromWasm0(arg0, arg1));
         },
-        __wbg_aot_logp_73e1e4a75800334a: function(arg0, arg1, arg2, arg3) {
+        __wbg_aot_evaluate_3748b6e27554efe5: function() { return handleError(function (arg0, arg1, arg2, arg3) {
+            const ret = aot_evaluate(arg0 >>> 0, arg1 >>> 0, arg2 >>> 0, arg3 >>> 0);
+            return ret;
+        }, arguments); },
+        __wbg_aot_logp_d2c28ba4c292a1e0: function(arg0, arg1, arg2, arg3) {
             const ret = aot_logp(arg0 >>> 0, arg1 >>> 0, arg2 >>> 0, arg3 >>> 0);
             return ret;
         },
-        __wbg_call_5c65f3296b077120: function() { return handleError(function (arg0, arg1, arg2, arg3, arg4, arg5, arg6) {
+        __wbg_call_ccf57e66805df431: 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();
@@ -561,10 +614,11 @@ function __wbg_get_imports() {
     return {
         __proto__: null,
         "./tapewasm_bg.js": import0,
-        "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js": import1,
-        "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js": import2,
-        "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js": import3,
-        "./snippets/tapewasm-6c5813ce65cccbcf/js/aot_bridge.js": import4,
+        "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js": import1,
+        "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js": import2,
+        "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js": import3,
+        "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js": import4,
+        "./snippets/tapewasm-c400eec12145fd93/js/aot_bridge.js": import5,
     };
 }
 

package/pkg/tapewasm_bg.wasm.d.ts

@@ -11,6 +11,7 @@ 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, i: number) => [number, number, number];
+export const aotsampler_evaluate: (a: number, b: number, c: number) => [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];
@@ -21,6 +22,7 @@ export const aotsampler_setTargetAccept: (a: number, b: number) => [number, numb
 export const clearAotExports: () => void;
 export const compileTape: (a: number, b: number, c: number, d: number) => [number, number, number];
 export const compiledtape_layoutId: (a: number) => number;
+export const compiledtape_nOutputs: (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];
@@ -36,9 +38,9 @@ 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 __wbindgen_free: (a: number, b: number, c: number) => void;
 export const __externref_table_dealloc: (a: number) => void;
 export const __wbindgen_start: () => void;