tapewasm 0.1.1 → 0.2.0

package/index.d.ts

@@ -1,6 +1,7 @@
 // Types for `index.js`; the shapes come from `pkg/`, which `make wasm` writes.
 export {
   AotSampler,
+  AdviResult,
   CompiledTape,
   compileTape,
   tapewasmVersion,

package/index.js

@@ -5,6 +5,7 @@
 import init from "./pkg/tapewasm.js";
 export {
   AotSampler,
+  AdviResult,
   CompiledTape,
   compileTape,
   tapewasmVersion,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tapewasm",
-  "version": "0.1.1",
+  "version": "0.2.0",
   "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.1.1",
+  "version": "0.2.0",
   "license": "MIT OR Apache-2.0",
   "repository": {
     "type": "git",

package/pkg/tapewasm.d.ts

@@ -1,6 +1,31 @@
 /* tslint:disable */
 /* eslint-disable */
 
+/**
+ * The fitted mean-field variational distribution and its ELBO trajectory.
+ */
+export class AdviResult {
+    private constructor();
+    free(): void;
+    [Symbol.dispose](): void;
+    /**
+     * The ELBO estimate after each iteration, in order.
+     */
+    readonly elboTrace: Float64Array;
+    /**
+     * The raw (pre-averaging) iterate `μ` every `snapshot_every` iterations,
+     * one snapshot's `n_params` values after another. Empty unless `advi`
+     * was called with `snapshot_every > 0`.
+     */
+    readonly muSnapshots: Float64Array;
+    readonly mu: Float64Array;
+    readonly sigma: Float64Array;
+    /**
+     * The iteration number each entry of `muSnapshots` was taken at.
+     */
+    readonly snapshotIters: Float64Array;
+}
+
 /**
  * A sampler over one compiled module.
  *
@@ -11,6 +36,32 @@
 export class AotSampler {
     free(): void;
     [Symbol.dispose](): void;
+    /**
+     * Mean-field ADVI: fits `q(θ) = N(μ, diag(σ²))` to the log density by
+     * Adam-ascending the ELBO under the reparameterization trick
+     * `θ = μ + σ·η`, `η ~ N(0, I)`. `init` seeds `μ`; `ω = log σ` starts at
+     * `0` (`σ = 1`).
+     *
+     * The ELBO gradient combines the tape's `∂logp/∂θ` (from `log_prob_grad`)
+     * with the reparameterization's own Jacobian and the mean-field Gaussian
+     * entropy's gradient, both exact rather than estimated:
+     * `∂ELBO/∂μ_i = ∂logp/∂θ_i`, `∂ELBO/∂ω_i = ∂logp/∂θ_i · σ_i · η_i + 1`.
+     * Each is averaged over `mc_samples` reparameterized draws per iteration.
+     *
+     * `AdviResult::mu`/`sigma` average `(μ, ω)` over the second half of
+     * `num_iters` (Polyak averaging): a constant Adam step size never settles
+     * on one point, so the plain last iterate keeps wandering a "noise ball"
+     * around the optimum. `elbo_trace` still reports the raw iterate's ELBO
+     * per step, for diagnosing convergence rather than for the fitted result.
+     *
+     * `snapshot_every` (0 disables it) additionally copies the raw iterate
+     * `μ` every that-many iterations into `AdviResult::muSnapshots` — a film
+     * strip of the one training run, rather than several shorter ones
+     * 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.
+     */
+    advi(init: Float64Array, num_iters: number, mc_samples: number, learning_rate: number, seed: bigint, snapshot_every: number): AdviResult;
     /**
      * `[log_prob, d/dparam...]`.
      */
@@ -98,8 +149,15 @@ export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembl
 
 export interface InitOutput {
     readonly memory: WebAssembly.Memory;
+    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 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_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];

package/pkg/tapewasm.js

@@ -1,11 +1,84 @@
 /* @ts-self-types="./tapewasm.d.ts" */
-import { aot_logp } from './snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js';
-import * as import1 from "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js"
-import * as import2 from "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js"
-import * as import3 from "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js"
-import * as import4 from "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js"
+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"
 
 
+/**
+ * The fitted mean-field variational distribution and its ELBO trajectory.
+ */
+export class AdviResult {
+    static __wrap(ptr) {
+        const obj = Object.create(AdviResult.prototype);
+        obj.__wbg_ptr = ptr;
+        AdviResultFinalization.register(obj, obj.__wbg_ptr, obj);
+        return obj;
+    }
+    __destroy_into_raw() {
+        const ptr = this.__wbg_ptr;
+        this.__wbg_ptr = 0;
+        AdviResultFinalization.unregister(this);
+        return ptr;
+    }
+    free() {
+        const ptr = this.__destroy_into_raw();
+        wasm.__wbg_adviresult_free(ptr, 0);
+    }
+    /**
+     * The ELBO estimate after each iteration, in order.
+     * @returns {Float64Array}
+     */
+    get elboTrace() {
+        const ret = wasm.adviresult_elboTrace(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * The raw (pre-averaging) iterate `μ` every `snapshot_every` iterations,
+     * one snapshot's `n_params` values after another. Empty unless `advi`
+     * was called with `snapshot_every > 0`.
+     * @returns {Float64Array}
+     */
+    get muSnapshots() {
+        const ret = wasm.adviresult_muSnapshots(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * @returns {Float64Array}
+     */
+    get mu() {
+        const ret = wasm.adviresult_mu(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * @returns {Float64Array}
+     */
+    get sigma() {
+        const ret = wasm.adviresult_sigma(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+    /**
+     * The iteration number each entry of `muSnapshots` was taken at.
+     * @returns {Float64Array}
+     */
+    get snapshotIters() {
+        const ret = wasm.adviresult_snapshotIters(this.__wbg_ptr);
+        var v1 = getArrayF64FromWasm0(ret[0], ret[1]).slice();
+        wasm.__wbindgen_free(ret[0], ret[1] * 8, 8);
+        return v1;
+    }
+}
+if (Symbol.dispose) AdviResult.prototype[Symbol.dispose] = AdviResult.prototype.free;
+
 /**
  * A sampler over one compiled module.
  *
@@ -24,6 +97,47 @@ export class AotSampler {
         const ptr = this.__destroy_into_raw();
         wasm.__wbg_aotsampler_free(ptr, 0);
     }
+    /**
+     * Mean-field ADVI: fits `q(θ) = N(μ, diag(σ²))` to the log density by
+     * Adam-ascending the ELBO under the reparameterization trick
+     * `θ = μ + σ·η`, `η ~ N(0, I)`. `init` seeds `μ`; `ω = log σ` starts at
+     * `0` (`σ = 1`).
+     *
+     * The ELBO gradient combines the tape's `∂logp/∂θ` (from `log_prob_grad`)
+     * with the reparameterization's own Jacobian and the mean-field Gaussian
+     * entropy's gradient, both exact rather than estimated:
+     * `∂ELBO/∂μ_i = ∂logp/∂θ_i`, `∂ELBO/∂ω_i = ∂logp/∂θ_i · σ_i · η_i + 1`.
+     * Each is averaged over `mc_samples` reparameterized draws per iteration.
+     *
+     * `AdviResult::mu`/`sigma` average `(μ, ω)` over the second half of
+     * `num_iters` (Polyak averaging): a constant Adam step size never settles
+     * on one point, so the plain last iterate keeps wandering a "noise ball"
+     * around the optimum. `elbo_trace` still reports the raw iterate's ELBO
+     * per step, for diagnosing convergence rather than for the fitted result.
+     *
+     * `snapshot_every` (0 disables it) additionally copies the raw iterate
+     * `μ` every that-many iterations into `AdviResult::muSnapshots` — a film
+     * strip of the one training run, rather than several shorter ones
+     * 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.
+     * @param {Float64Array} init
+     * @param {number} num_iters
+     * @param {number} mc_samples
+     * @param {number} learning_rate
+     * @param {bigint} seed
+     * @param {number} snapshot_every
+     * @returns {AdviResult}
+     */
+    advi(init, num_iters, mc_samples, learning_rate, seed, snapshot_every) {
+        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);
+        if (ret[2]) {
+            throw takeFromExternrefTable0(ret[1]);
+        }
+        return AdviResult.__wrap(ret[0]);
+    }
     /**
      * `[log_prob, d/dparam...]`.
      * @param {Float64Array} params
@@ -253,7 +367,7 @@ function __wbg_get_imports() {
         __wbg___wbindgen_throw_5d9e815e6fdf150f: function(arg0, arg1) {
             throw new Error(getStringFromWasm0(arg0, arg1));
         },
-        __wbg_aot_logp_f9a6d0f4683f5890: function(arg0, arg1, arg2, arg3) {
+        __wbg_aot_logp_897c94d09e163a67: function(arg0, arg1, arg2, arg3) {
             const ret = aot_logp(arg0 >>> 0, arg1 >>> 0, arg2 >>> 0, arg3 >>> 0);
             return ret;
         },
@@ -292,13 +406,16 @@ function __wbg_get_imports() {
     return {
         __proto__: null,
         "./tapewasm_bg.js": import0,
-        "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js": import1,
-        "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js": import2,
-        "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js": import3,
-        "./snippets/tapewasm-efcc61b460c384cc/js/aot_bridge.js": import4,
+        "./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,
     };
 }
 
+const AdviResultFinalization = (typeof FinalizationRegistry === 'undefined')
+    ? { register: () => {}, unregister: () => {} }
+    : new FinalizationRegistry(ptr => wasm.__wbg_adviresult_free(ptr, 1));
 const AotSamplerFinalization = (typeof FinalizationRegistry === 'undefined')
     ? { register: () => {}, unregister: () => {} }
     : new FinalizationRegistry(ptr => wasm.__wbg_aotsampler_free(ptr, 1));

package/pkg/tapewasm_bg.wasm.d.ts

@@ -1,8 +1,15 @@
 /* tslint:disable */
 /* eslint-disable */
 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 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_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];