A Stan runtime that runs on wasm

Lately I have been hooked on data analysis and WebAssembly.
There is something about it that feels like the mixed martial arts of AI and computing, and I like that.
On the not-very-considered grounds that Bayesian inference in the browser would be fun, I built stanwasm, which runs Stan on WebAssembly.
This post introduces it, and tries to put the thinking behind it into words.
stanwasm takes a subset of the Stan language and does everything — parsing, compilation, sampling — inside the browser.
The parser, the evaluator, the autodiff tape and the code generator were written from scratch in Rust for this project; only the sampler is borrowed, by compiling nuts-rs by Adrian Seyboldt and the PyMC developers to wasm.

• Run the above in your browser, no install required
• From Python on JupyterLite — the same thing runs through Pyodide

habakan/stanwasmA Stan subset that runs entirely in the browser. Apache-2.0, on npm as stanwasmgithub.com

It is still alpha and the API may change, but npm install [email protected] will get you it.

import init, { StanModel } from "stanwasm";
await init();

const model = new StanModel(stanCode, JSON.stringify(data));
const draws = model.sample(model.randomInit(42n), 1000, 1000, 42n);
//                          init, warmup, draws, seed

It is not a replacement for CmdStan or Stan Playground; it aims at the "embed it in a browser" case those do not fit.
That said, after looking at several designs, keeping everything inside the browser and simple ended up meaning that almost none of Stan's existing assets are used.
So what this really is, is a runtime of my own that samples with NUTS from a model written in the Stan language.

What to watch out for when putting Stan on wasm

Let me start with a quick tour of Stan's toolchain.
Stan is implemented on top of C++.
You describe a probabilistic graphical model in Stan's syntax, and from that, the sampling code that computes the posterior with MCMC and friends is compiled.
The distributions written in Stan are parsed and then turned into derivatives so that HMC can sample them. What handles the differentiation is Stan Math Library, an autodiff library optimised for speed.
It is this long-maintained C++ code that lets Stan sample posteriors fast.

If it is implemented in C++, surely emscripten can turn it into wasm right away? Some of you will think that, but there are a few obstacles in the way.
They come down to the fact that a compiler is part of the toolchain that gets you from Stan source to a posterior.

Laid out, the structure looks like this.

model.stan is first turned into C++ source (model.hpp) by stanc3.
That is then built by a C++ compiler on your machine, together with Stan Math for autodiff and stan::services for NUTS / HMC.
What comes out is an executable specific to that model; run it and the draws come out as CSV.
Both the derivatives and the sampler are pulled into the executable at this point.
Whichever interface you use, this build runs, so you need a C++ compiler on hand.
Only the data is passed at run time, so the same executable can be reused with different data.

If the interface is a web page where someone writes Stan, then of course the runtime cannot know in advance which distributions or functions will show up.
That narrows things to two options.

One is to reproduce the existing path in the browser.
On top of stanc3 and Stan Math, that means bringing the C++ compiler itself into the browser.

The other is to compile every distribution and function Stan offers ahead of time, and ship the whole set as a runtime.
Since you cannot know which model is coming, you end up carrying things that are never used.
No compiler is needed, but the download grows accordingly.

So the question for running Stan on wasm becomes: how do you represent Stan's toolchain compactly in the world of wasm, and for that, what representation belongs between the Stan source and the sampler? The answer I settled on for now was to make the autodiff graph an explicit intermediate representation, and emit wasm instructions from it.

How stanwasm is put together

Before that, let me set out why autodiff and a tape are needed at all.

NUTS and HMC look at the gradient of the posterior to choose the next point.
To run sampling, you need to know how each distribution in the model contributes to the log density — its gradient.

When distributions and operations are chained inside a model, gradients compose through the chain rule.
Doing that automatically is autodiff.

The composition runs backwards, from the output towards the inputs, so it cannot be computed unless you know which operations built the log likelihood, and in what order.
So you record the process of the computation itself.
That record is what is called a tape.

Stan, in the course of compiling a model into an executable, pulls in the code that computes that model's derivatives by autodiff.
Because it builds per model, the types of the nodes pushed onto the tape are fixed to that model too.
This is what makes its posterior computation fast.

Where stanwasm diverges is when the tape is built, and what is done with it afterwards.

Laid out in the same frame, the stanwasm side looks like this.

stanwasm builds the tape once, at load time. Inside new StanModel(src, data) it parses, binds the data, and records while evaluating — all of that is done by then.

On the path from the previous section, the operations that build and run the derivatives land inside the executable at the moment the C++ is compiled.

stanwasm turns the forward and backward passes into a wasm module after it has seen a specific model in the browser. What makes that possible is that stanwasm keeps the derivatives explicitly as data. Stan Math builds a tape at run time too. var_stack_ is a sequence of vari_base*, and grad() reads it in reverse. The difference is not whether there is a tape, but what the nodes hold.

Read the tape every time, or emit instructions once

stanwasm's Tape is a struct that just holds op, arg1, val and grad as arrays of the same length. op[i] is Op::Mul — that is, the value 3 — and the backward pass reads it and branches on it itself.

What it means for the tape to be a sequence of values is that you can change how you read it. Read it front to back and you get the log density; read it back to front and you get the gradient; read it forwards while emitting instructions and you get a wasm module for that model.

Let me first say why you would do that.

NUTS calls for a gradient tens of thousands of times, and over all of those the structure of the computation never changes once.
What changes is only the values of the parameters.
The path that reads the tape is re-reading that settled structure on every gradient.
If the structure is fixed from the start, it should be enough to build code with it baked in once, and then just call that.
Call it the difference between an interpreter and compiled code, if you like.

What it takes is simple.
The backward pass reads op[i] and branches on "the derivative of this operation is this"; you swap that branch for "the wasm instruction for this operation is this" (stanwasm-codegen).
The shape of the tape being read is the same; the only difference is whether, on seeing Op::Mul, you multiply there and then, or lay down one multiply instruction.
In the implementation the load-time tape is not reused — it reads one recorded afresh with dummy parameters.

This is done for both directions, and the instructions that compute the log density and the instructions that accumulate the gradient are collected into a single function, log_prob_grad.
It is not that the log density and the gradient are switched between; both sets of instructions sit side by side in the same function. NUTS needs both at once anyway.

Stan Math cannot do the same thing.
Calling the chain method attached to a node advances the differentiation, but you cannot pull out "what operation is this" as a value. Which chain gets called was fixed when the C++ was compiled.
Adding a mechanism to emit instructions would mean rebuilding Stan Math, and we are back to the original problem of needing a compiler.

The instructions being f64x2 SIMD is also something only a per-model module can do. When you read a tape generically you do not know what operation comes next, so you cannot decide to pair up two doubles.

What the implementation showed

A rough implementation seems to get this far.

Correctness is checked against posteriordb, a collection of Stan models with reference results.

The posterior mean check is means only; R-hat, ESS and divergences have not been compared yet.
The 1.5–12x is also an internal comparison between my own two paths, not a benchmark against Stan.

That said, because so much of this was written from scratch, plenty has been given up.

• The language is a subset of Stan, not all of it
• The sampler is separate code, so the same seed does not reproduce Stan's draws
• Fixes to Stan Math do not flow through. The special functions and the distribution edge cases were written and tested here

And above all, things whose computation changes shape with the parameters sit badly with recording once.
Adaptive ODE solvers are the clear example: how they take steps depends on the parameters, so a recorded graph freezes into the shape it had at the tracing point.
The record-and-replay path refuses these, and sampleFresh() re-records the tape on every gradient instead. integrate_ode_rk45 works on this path, at 5–8x the cost of replay.
With ode_rk4_fixed, where the caller fixes the step count, the graph does not move, so it runs on replay and AOT alike. integrate_ode_bdf is not there yet.

There is no single right answer to what belongs between the Stan source and the sampler.
Putting a static graph on top of precompiled kernels is one, and I think the difference is just in the trade between compatibility, artifact size and how far you specialise.

Indeed, stanli goes down that road.
I only learned it existed after building stanwasm.
It makes good use of existing assets such as Stan Math while putting the per-model autodiff on wasm as an interpreter, which I think is wonderful.

Closing

To run Stan models with nothing but a browser, I ended up holding the autodiff graph as an explicit intermediate representation and emitting wasm instructions from it.
It is still alpha and the language subset is narrow, but it is something you can try without installing anything.

I thought I was thinking about Bayes, and somewhere along the way it turned into a study of assembly and compilers.