Add experimental WASM adapter for browser PyMC sampling - #345
Closed
twiecki wants to merge 1 commit into
Closed
Conversation
Member
Author
|
Will implement as a separate project. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 makepip install nutpieornutpie.samplework in the browser.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).