tapewasm 0.3.1 → 0.3.3

package/calibrate.js

@@ -0,0 +1,141 @@
+// Which shape this engine prefers, measured once.
+//
+// `compileTape`'s `"auto"` threshold cannot be right for every engine:
+// straight-line and re-rolled cross over between 8,026 and 24,564 nodes in V8
+// and around 2,000 in SpiderMonkey and JavaScriptCore. The built-in value serves the
+// latter two and costs V8 up to 7.6x on a trace between them. This measures
+// which side the engine is on and returns a threshold to pass to `compileTape`.
+//
+//   await init(...);
+//   const above = await calibrateReroll();
+//   const built = compileTape(text, String(above));
+//
+// The answer is cached: it is a property of the engine, not of the model.
+
+import { compileTape } from "./pkg/tapewasm.js";
+
+/// What `"auto"` uses. Right for SpiderMonkey and JavaScriptCore.
+export const RE_ROLL_ABOVE = 2_000;
+
+/// For an engine that prefers straight-line, as V8 does.
+///
+/// Bracketed by measurement rather than picked: on eleven posteriordb models,
+/// straight-line still wins at 8,026 nodes and has lost by 24,564. A threshold
+/// of 24,000 already costs `low_dim_gauss_mix` 1.21x, so this sits clear of
+/// that edge. At 20,000 the five models whose shape changes get 2.54x in the
+/// geometric mean and none is slower.
+export const V8_RE_ROLL_ABOVE = 20_000;
+
+/// Where the two engine families disagree most clearly. At this size V8 prefers
+/// straight-line on every model measured, and the other two prefer loops on all
+/// but one.
+const PROBE_NODES = 4_000;
+
+let cached = null;
+
+/// A tape of the commonest shape — one accumulated term per observation —
+/// using nothing but arithmetic, so the module it compiles to imports only
+/// memory and the probe needs no maths from the caller.
+function probeTape(terms) {
+  // Instruction 0, 1, 2 are the parameters; `at` is the index the next line takes.
+  const out = ["n_params 3", "new_var 0.4", "new_var 1.1", "new_var 0.3"];
+  let at = 3;
+  const push = (line) => {
+    out.push(line);
+    return at++;
+  };
+  let acc = null;
+  for (let i = 0; i < terms; i++) {
+    const x = (0.001 * (i % 997)).toFixed(6);
+    const y = (0.5 + 0.001 * (i % 991)).toFixed(6);
+    const bx = push(`mul_c 1 ${x}`);     // beta * x
+    const mu = push(`add ${bx} 0`);      // + alpha
+    const r = push(`rsub_c ${mu} ${y}`); // y - mu
+    const z = push(`div ${r} 2`);        // / sigma
+    const sq = push(`mul ${z} ${z}`);
+    acc = acc === null ? sq : push(`add ${acc} ${sq}`);
+  }
+  const half = push(`mul_c ${acc} -0.5`);
+  out.push(`root ${half}`);
+  return out.join("\n") + "\n";
+}
+
+async function timeOne(wasm, nParams, scratchInit, rounds) {
+  const need = nParams * 16 + scratchInit.length * 8;
+  const memory = new WebAssembly.Memory({ initial: Math.ceil(need / 65536) + 2 });
+  const { instance } = await WebAssembly.instantiate(wasm, { tapewasm: { memory } });
+  const view = new Float64Array(memory.buffer);
+  view.set(scratchInit, nParams * 2);
+  for (let i = 0; i < nParams; i++) view[i] = 0.1 * (i + 1);
+  const lpg = instance.exports.log_prob_grad;
+  const call = () => lpg(0, nParams * 8, nParams, nParams * 16);
+
+  for (let i = 0; i < 50; i++) call();          // past the first tier
+  const t0 = performance.now();
+  let n = 0;
+  while (performance.now() - t0 < 4) { call(); n++; }
+  const iters = Math.max(1, n);
+  let best = Infinity;
+  for (let r = 0; r < rounds; r++) {
+    const t = performance.now();
+    for (let j = 0; j < iters; j++) call();
+    best = Math.min(best, (performance.now() - t) / iters);
+  }
+  return best;
+}
+
+/**
+ * Measure which shape this engine prefers and return a threshold for
+ * `compileTape`'s `reroll` argument. Call after `init()`.
+ *
+ * Cached after the first call — pass `{ force: true }` to measure again.
+ * On any failure it returns the built-in threshold, so a caller can use the
+ * result without guarding it.
+ */
+export async function calibrateReroll({ force = false, rounds = 7 } = {}) {
+  if (cached !== null && !force) return cached.above;
+  let result;
+  try {
+    const text = probeTape(Math.max(1, Math.round(PROBE_NODES / 5)));
+    const straight = compileTape(text, "never");
+    const looped = compileTape(text, "always");
+    // Interleaved, so a drift in machine state lands on both.
+    let sBest = Infinity, lBest = Infinity;
+    for (let r = 0; r < 2; r++) {
+      sBest = Math.min(sBest, await timeOne(straight.wasm, straight.nParams, straight.scratchInit, rounds));
+      lBest = Math.min(lBest, await timeOne(looped.wasm, looped.nParams, looped.scratchInit, rounds));
+    }
+    const prefersStraight = sBest < lBest;
+    result = {
+      above: prefersStraight ? V8_RE_ROLL_ABOVE : RE_ROLL_ABOVE,
+      prefersStraight,
+      straightMs: sBest,
+      loopedMs: lBest,
+      measured: true,
+    };
+  } catch (e) {
+    // A threshold is an optimisation; failing to measure one is not an error.
+    result = { above: RE_ROLL_ABOVE, measured: false, error: String(e?.message ?? e) };
+  }
+  cached = result;
+  return result.above;
+}
+
+/** What the last `calibrateReroll` measured, or `null` before the first call. */
+export function lastCalibration() {
+  return cached;
+}
+
+/**
+ * `compileTape` at the threshold this engine prefers, measured once.
+ *
+ * The plain `compileTape` is synchronous and cannot wait for a measurement, so
+ * its `"auto"` is the built-in 2,000 — right for SpiderMonkey and
+ * JavaScriptCore, and low for V8. On posteriordb's 23 models in Node, that
+ * default loses on thirteen of them, by up to 12x. This is the same call with
+ * the measured number, and the measurement happens once per process.
+ */
+export async function compileTapeCalibrated(text, opts) {
+  const above = await calibrateReroll(opts);
+  return compileTape(text, String(above));
+}

package/index.js

@@ -14,5 +14,11 @@ export {
   clearAotExports,
   sharedMemory,
 } from "./pkg/tapewasm.js";
-export { calibrateReroll, lastCalibration, RE_ROLL_ABOVE, V8_RE_ROLL_ABOVE } from "./calibrate.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.1",
+  "version": "0.3.3",
   "description": "Compile an autodiff tape to a WebAssembly module and sample it with nuts-rs, in the browser.",
   "type": "module",
   "main": "./index.js",
@@ -13,7 +13,7 @@
     },
     "./pkg/*": "./pkg/*"
   },
-  "files": ["index.js", "index.d.ts", "pkg/", "LICENSE-APACHE", "LICENSE-MIT", "README.md"],
+  "files": ["index.js", "index.d.ts", "calibrate.js", "pkg/", "LICENSE-APACHE", "LICENSE-MIT", "README.md"],
   "scripts": {
     "build:wasm": "make -C .. wasm"
   },

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.1",
+  "version": "0.3.3",
   "license": "MIT OR Apache-2.0",
   "repository": {
     "type": "git",

package/pkg/snippets/{tapewasm-cbf82a3d7bc31b1b → tapewasm-a5eaf252e13afb7a}/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,6 +67,19 @@ 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...]`.
      *
@@ -120,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;
     /**
@@ -176,6 +194,10 @@ 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"`, `"never"`,
  * or a node count to re-roll past, written as a number.
@@ -230,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];
@@ -240,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];
@@ -255,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-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"
+import { aot_evaluate, aot_logp } from './snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js';
+import * as import1 from "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js"
+import * as import2 from "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js"
+import * as import3 from "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js"
+import * as import4 from "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js"
+import * as import5 from "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js"
 
 
 /**
@@ -144,6 +145,31 @@ 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...]`.
      *
@@ -289,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}
      */
@@ -420,6 +455,10 @@ 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"`, `"never"`,
  * or a node count to re-roll past, written as a number.
@@ -524,11 +563,15 @@ function __wbg_get_imports() {
         __wbg___wbindgen_throw_5d9e815e6fdf150f: function(arg0, arg1) {
             throw new Error(getStringFromWasm0(arg0, arg1));
         },
-        __wbg_aot_logp_88e25d7a78bf7680: function(arg0, arg1, arg2, arg3) {
+        __wbg_aot_evaluate_90ed2a4aec33a889: 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_212144b205cad1c9: 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) {
+        __wbg_call_ae8ca521a7cc221c: 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();
@@ -571,10 +614,11 @@ function __wbg_get_imports() {
     return {
         __proto__: null,
         "./tapewasm_bg.js": import0,
-        "./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,
+        "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js": import1,
+        "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js": import2,
+        "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js": import3,
+        "./snippets/tapewasm-a5eaf252e13afb7a/js/aot_bridge.js": import4,
+        "./snippets/tapewasm-a5eaf252e13afb7a/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;