I’d like to introduce the stanc3 project, a complete rewrite of the Stan 2 compiler in OCaml.

Join us!

With this rewrite and migration to OCaml, there’s a great opportunity to join us on the ground floor of a new era. Your enthusiasm for or expertise in programming language theory and compiler development can help bring Stan into the modern world of language design, implementation, and optimization. If this sounds interesting, we could really use your help! We’re meeting twice a week for a quick standup on Mondays and Wednesdays at 10am EST, and I’m always happy to help people get started via email, hangout, or coffee. If you’re an existing Stan user, get ready for friendly new language features and performance upgrades to existing code! It might be a little bumpy along the way, but we have a really great bunch of people working on it who all care most about making Stan a great platform for practicing scientists with bespoke modeling needs.

The opportunity

Stan is a successful and mature modeling language with core abstractions that have struck a chord, but our current C++ compiler inhibits some next-gen features that we think our community is ready for. Our users and contributors have poured a huge amount of statistical expertise into the Stan project, and we now have the opportunity to put similar amounts of programming language theory and compiler craftsmanship into practice. The rewrite will also aim at a more modular architecture, which will enable tooling to be built on top of the Stan compiler enabling features like IDE auto-completion and error highlighting, as well as programming and statistical code linters that can help users with common sources of modeling issues. OCaml’s powerful and elegant pattern matching and seasoned parsing library make it a natural fit for the kinds of symbolic computation required of compilers. This makes it much more pleasant and productive for the task at hand, and is reflected by its frequent use by programming language researchers and compiler implementers. OCaml’s flagship parsing library Menhir enabled Matthijs Vákár to rewrite the Stan parsing phase in about a week, adding hundreds of new custom error messages in another week. Matthijs is obviously a beast, but I think he would agree that OCaml & Menhir definitely helped. Come join us and see for yourself :)

New language features

After we replicate Stan’s current compilers functionality, we will be targeting new language features. The to-do list includes, but is not necessarily limited to:

• tuples
• tools for representing and working with ragged arrays
• higher order functions (functions that take in other functions)
• variadic functions
• annotations
  • to bring methods like Posterior Predictive Checking and Simulation-Based Calibration into Stan itself
  • to label variables as “silent” (not output), or as living on a GPU or other separate hardware
  • to assist those who would like to use Stan as an algorithms workbench
• user-defined gradients
• representations for missing data and sparse matrices
• discrete parameter marginalization

Next-gen optimization

But back to the next-gen features. Here is just some of the low-hanging fruit:

• peephole optimizations: we might notice when a user types log(1- x) and replace it with log1m(x) automatically
• finding redundant computations and sharing the results
• moving computation up outside of loops (including the sampling loop!)
• using the data sizes to ahead-of-time compile a specialized version of the Stan program in which we can easily unroll loops, inline functions, and pre-allocate memory
• pulling parts of the Math library into the Stan compiler to e.g. avoid checking input matrices for positive-definiteness on every iteration of HMC

There is a wealth of information at the Stan language level we can take advantage to produce more efficient code than the more mature C++ compilers we rely on, and we can use the new compiler to pass some of that information along to the C++ code we generate. Maria Gorinova showed us with SlicStan how to move code to its most efficient (Stan) block automatically as well as a nice composition-friendly syntax. We can use similar static analysis tools in a probabilistic setting to e.g. allow for discrete parameters via automated Rao-Blackwellization (i.e. integrating them out) or discover conjugacy relationships and use analytic solutions where applicable. We can go a step further and integrate with a symbolic differentiation library to get symbolic derivatives for Stan code as a fast substitution for automatic differentiation.

Advanced techniques

Once we’ve created a platform for expressing Stan language concepts and optimizing them, we’ll naturally want to bring as much of our computation onto that platform as possible so we can optimize holistically. This will mean either using techniques like Lightweight Modular Staging to parse our existing C++ library into our Stan compiler representation, or beginning a project to rewrite and simplify the Stan Math library into Stan language itself. We hope that with some of the extensions above, we’ll be able to express the vast majority of the Math library in the Stan language, and lean heavily on a symbolic differentiation library and the stanc3 code generator to generate optimized C++ code. This should shrink the size of our Math library by something like 12x, and takes the code generation techniques used in PyTorch to the next level.

Alternative backend targets (TensorFlow, Pytorch, etc.)

At that point, targeting multiple backends will become fairly trivial. We can put compilation times squarely in the cross-hairs and provide an interpreted Stan that immediately gives feedback and has minimal time-to-first-draw. We can also target other backends like TensorFlow Probability and PyTorch that do not possess the wealth of specialty statistical functions and distributions that we do, but may make better use of the 100,000 spare GPUs you have sitting in your garage.