Skip to content

Add experimental WASM adapter for browser PyMC sampling - #345

Closed
twiecki wants to merge 1 commit into
pymc-devs:mainfrom
twiecki:codex/browser-wasm-adapter
Closed

Add experimental WASM adapter for browser PyMC sampling#345
twiecki wants to merge 1 commit into
pymc-devs:mainfrom
twiecki:codex/browser-wasm-adapter

Conversation

@twiecki

@twiecki twiecki commented Sep 6, 2026

Copy link
Copy Markdown
Member

PyMC models can already compile logp/gradients with Numba in a compatible browser WASM runtime, but Python still executes the NUTS control flow. This adds an experimental, standalone WASM adapter around Nutpie's sequential nuts-rs core so the trajectory loop runs in Rust and invokes the compiled model without a Python call per evaluation.

The proposed home is Nutpie because this is the PyMC-to-sampler integration layer. The existing native package and public API are unchanged. This draft starts under examples/browser/ to review the ABI and distribution approach before promoting it to a supported backend; it does not make pip install nutpie or nutpie.sample work in the browser.

  • Independent Rust cdylib with a pinned nuts-rs dependency and no Python extension/storage dependencies.
  • General continuous-model PyMC/Numba compiler that freezes shared data and exports a logp/gradient C callback.
  • Worker-side JS bridge between the separate Rust and Emscripten memories/function tables.
  • Reproducible WASM build, smoke tests, and CI artifact containing the static adapter files. The compatible Xeus Python runtime is supplied separately; no public release is automatically published.

Validation: release WASM build; real Rust WASM Gaussian posterior moments, memory growth and failure/progress handling through the JS bridge; native PyMC callback agreement for transformed parameters and frozen shared data. The original adapter also completed a real 179-week, 15-parameter PyMC-Marketing MMM in Xeus with Numba (2 × 750 warmup + 500 retained draws). Native/browser sampling was 0.80/8.41 seconds, excluding compilation. R-hat was 1.018/1.023, so this is feasibility evidence, not a speedup claim at matched convergence.

Remaining before a supported backend: supported Xeus memory export instead of a local loader patch; runtime/package distribution and compatibility testing; constrained-variable/deterministic expansion and InferenceData integration; initialization, configurable adaptation, cancellation and stronger error recovery. The README documents these limits and the exact timer boundaries.

The generalized compiler and bridge in this draft were additionally tested with that MMM in the actual browser: 8.24 s sampling, 14.85 s preparation/compilation, and identical posterior diagnostics to the prototype (80,911 evaluations, zero divergences, max R-hat 1.0233, min bulk ESS 183.4).

@twiecki

twiecki commented Sep 6, 2026

Copy link
Copy Markdown
Member Author

Will implement as a separate project.

@twiecki twiecki closed this Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant