pymc-devs / pymc-devs/nutpie

Make the PyMC/Numba frontend usable without _lib for WASM integration

Open
#347 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
207
Forks
28
Avg merge
7d 21h
Merged PRs (30d)
1

Description

Problem

In nuts-rs-wasm, we run PyMC/Numba model evaluation in a Xeus/Emscripten Python runtime and call nuts-rs through a separate WASM adapter. The current integration does not require Nutpie's native Python extension, but independently implements model compilation, expansion, metadata and mutable-data handling. This overlaps substantially with Nutpie and risks diverging behavior as both implementations evolve.

We would like to reuse Nutpie's Python model frontend without having to build or load nutpie._lib in the browser.

Current import blockers

At commit 98c0e879cdd110c14e6bc82a001251b53a990f58:

The graph preparation in _make_functions and the Numba callback builders appear separable from native sampler construction. Moving a single import would not resolve the complete dependency chain.

Proposed direction

Expose an independently importable Python/Numba frontend, with its exact API and module layout agreed here before implementation:

  1. Separate graph compilation, callback creation, metadata and buffer ownership from native model/sampler construction.
  2. Return a compiled model description containing density/gradient and expansion callbacks, dimensions/shapes/coordinates, initial-point generation, and the owned shared-data buffers needed to keep pointers valid.
  3. Let the existing native backend consume that description. Let a WASM adapter consume the same description through its own callback bridge.
  4. Make package initialization and shared types usable without _lib; load the extension when native functionality is requested, with an actionable error if unavailable.
  5. Provide a supported way to install/distribute the frontend without requiring a native-extension build. The packaging choice is open; this need not imply creating a separate package.

The initial scope would be the PyMC/Numba path. This is not a request to port the full native sampler, Stan, or every backend to WASM.

Browser compatibility still needs validation

An extension-independent import is necessary but does not by itself prove browser compatibility:

  • Nutpie's callbacks have status returns, separate log-density output storage, 64-bit dimension arguments and a user-data pointer. Our current bridge uses different signatures and must adapt, including WASM i64 handling and error propagation.
  • Shared arrays and their dtype/pointer/shape metadata can remain in the Python runtime. The separate Rust memory does not inherently require our custom float64 packing, but Nutpie's Numba pointer extraction must be tested in the actual Xeus build.
  • Expansion ordering/selection and initial-point semantics differ from our current adapter. We should explicitly map outputs and use PyMC's initial-point generation rather than assume identical layouts or historical RNG sequences.
  • Callback and buffer lifetime must remain valid across repeated sampling, data updates and release.

The browser worker, cross-memory bridge and synchronous nuts-rs execution adapter would remain in nuts-rs-wasm initially. Result conversion could be shared in a later change; it should not block the compiler extraction.

Acceptance criteria

  • The supported frontend import and compilation of a simple continuous PyMC model work when _lib is absent, without a fake extension or copying private functions downstream.
  • Existing native compilation and sampling behavior remains covered by regression checks.
  • A bounded Xeus probe validates density, gradients and expansion against PyMC for Gaussian, transformed, simplex and deterministic models, followed by the MMM example.
  • The probe covers shared-data dtypes and updates, callback status handling, pointer widths, memory growth and ownership/lifetime.
  • The public boundary, installation route and remaining browser limitations are documented.

Would this frontend/backend split be an acceptable upstream direction? In particular, should the reusable result be a standalone compiled-model description or a refactoring of CompiledPyMCModel with native operations separated out?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading python/nutpie/init.py, compile_pymc.py, sample.py, and compiled_pyfunc.py at the referenced commit to trace the _lib dependency chain. Review _make_functions, the Numba callback builders, and existing native regression checks before proposing the frontend boundary. Done means an agreed API and packaging route, extension-independent compilation, preserved native behavior, and the specified Xeus validation and documentation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, wasm
Domain
backend, compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.