Nikolas Siccha writes:
I got inspired by Bob’s latest blog post to bring your attention to this code generation tool. It’s currently supremely buggy and only really usable by myself – but it may still be interesting for you to see, as it IMO solves many of Stan’s “problems.”
“Supremely buggy and only really usable by myself” . . . that describes lots of the software that I write!
Follow the link and here’s what you see:
Brings Julia syntax to Stan models by implementing a (limited) Julia to Stan transpilation with many caveats. See
test/slic.jlfor implementations of a few simpleposteriordbmodels and seesrc/slic_stan/builtin.jlfor a list of built-in functions and examples of user defined functions.Current features include
- activity analysis (automatically determines what is
data,transformed_data,parameters,transformed parameters,model, orgenerated quantities),- automatically inferred types, shapes and constraints – including for user defined functions (including the function arguments, function body, and function return type),
- automatic posterior pointwise likelihood and predictive generation,
- (variadic) user defined functions,
- higher order (user defined) functions (such as
map,broadcasted,sumand more),- sub models,
- post-hoc model adjustment,
- named tuples,
- (approximate) automatic code formatting à la Blue,
- and more.
Upcoming features include, in order of priority and estimated arrival,
- easy runtime assertions (like Julia’s
@assert) – and support for other macros,- model docstrings,
- custom types (for method dispatch – this would help with more “Julia-style” broadcasting, e.g. via
Ref),- closures via Julia’s
Do-Block Syntax(to make within chain parallelization viareduce_sumless painful),- lower transpilation runtimes (currently, transpilation can sometimes take longer than compilation – there is currently at least one algorithmic inefficiency on top of the systemic implementation inefficiency),
- a much better user experience,
- more and better tests,
- keyword arguments,
- default arguments,
- inlining (to reduce potential runtime overhead),
- easier custom parameter transformations (going from sampler parametrization to user parametrization – aka as constraining),
- array comprehensions,
- a more complete (and more correct) coverage of built-in Stan functions,
- better name resolution (currently user defined functions or sub models have to be defined in
Main),- functions that mutate their arguments (solved via inlining),
- and more.
Almost anything that’s possible in Julia should be possible to be transpiled to Stan. Of course, unless Stan is much faster than Julia (+Mooncake or Enzyme) for the model in question, just sticking to Julia comes with many advantages.
Features which I am on the fence about, but currently not planning to implement:
- a Julia backend,
target +=statements,- top level control flow,
- top level mutability,
- getting rid of superfluous parentheses.
Features which are NOT planned:
- (automatically) transpiling Julia functions which have not been defined via
@deffun.The
earn_height.stanmodel below becomes
using StanBlocks
import PosteriorDB, StanLogDensityProblems, JSON
# Get data from PosteriorDB
pdb = PosteriorDB.database()
post = PosteriorDB.posterior(pdb, "earnings-earn_height")
(;earn, height) = (;Dict([Symbol(k)=>v for (k, v) in pairs(PosteriorDB.load(PosteriorDB.dataset(post)))])...)
# Model definition
earn_height_model = @slic begin
beta ~ flat(;n=2)
sigma ~ flat(;lower=0.)
earn ~ normal(beta[1]+beta[2]*to_vector(height), sigma)
end
# Not compiled yet
earn_height_posterior = earn_height_model(; earn, height)
# Prints the Stan model code
println(stan_code(earn_height_posterior))
# Compiled (requires StanLogDensityProblems and JSON)
earn_height_problem = stan_instantiate(earn_height_posterior)
I don’t quite get that code at the end, but then again I don’t speak Julia.
One reason I think this is cool is that Julia and Stan started at the same time! Both projects were originally supported by a Department of Energy grant that funded Alan Edelman and me to hire postdocs and programmers. Alan’s group built Julia; we built Stan. Thanks, Bob, Matt, Daniel, and everybody else who did this!
Stepping back, it’s good to be able to go between languages and not be stuck in just one framework. We wrote Stan for two reasons:
1. To be able to easily and flexibly express Bayesian models and fit them to data.
2. To be able to perform statistical workflow, which includes fitting multiple models, postprocessing inference, plotting data and fitted models, etc.
For step 2, it was absolutely necessary for Stan to link smoothly to general-purpose computing environments such as R and Python, which is why we built rstan and pystan (which have since mostly been replaced by the lighter interfaces cmdstanr and cmdstanpy; thanks, Jonah, Mitzi, and everybody else who did this!). You can also run Stan from Julia. What Niko is doing here is going beyond that in allowing some sort of improved version of Stan that can be written in Julia, which seems like it could be really useful.
Beyond all this, as a scholar I’ve always tried to avoid any territorial attitude. I’ve just seen this too often, with researchers who have some method or another and are out there dissing the competition. OK, don’t get me wrong here, I’m happy to diss, when it’s appropriate, but not just cos it’s competition. Science isn’t a competition; we’re all in this together.
So, yeah, Stan is great. I love Stan, and I use it. I eat my own dogfood. But the purpose of Stan is Bayesian workflow. And if Stan can be helpful to Bayesian workflow in other ways, that’s great.
As I understand this, it literally takes julia code, and creates a Stan file under the hood that implements the same computation and then calls Stan for sampling. If you call certain user-written Julia functions, it maybe transpiles the code for those functions into your Stan file.
Stan is extremely well optimized to do its special task, so this can be faster than Julia.
I have found however that staying in Julia and writing models in Turing.jl with Mooncake or Enzyme as differentiation backends can be sufficiently fast as well. It also gives you access to a relatively easy to use profiler, debugging tools, and the full range of Julia code. So if there’s a package that does something special you need in your model, it pretty much just works to include that package and use it in the model. That includes things like Orthogonal Polynomials, the massive SciML differential equation solving toolkit, geospatial stuff, graph algorithms, or whatever.
I’m with Andrew when he says not to diss the competition just because it’s competition. Stan is great stuff that runs in a very optimized way. Julia as a language is also great stuff that runs in an optimized way, anything that lets the two work better together is welcome news. I do encourage R and Python users to look into Julia. It’s a great language that has made incredible strides since its 1.0 release in 2018.
> I’m with Andrew when he says not to diss the competition just because it’s competition
Yeah agreed. And similarly for the “Supremely buggy and only really usable by myself”. I think I’m always a little skeptical of things getting reinvented, but practically there’s so many reinvented things that I end up loving!
I’ve been so happy the last few years switching from nothing to poetry to uv for Python project management.
And then on the Stan side, the pystan -> cmdstanpy -> nutpie + bridgestan — all really good reinventions in my mind.
Maybe reinvention is more of an eye-of-the-beholder-thing — like I’m sure the people doing these things expect the thing to be new and different, but in my mind I don’t really expect much difference, and then I am thrilled and delighted with what pops out the other side (so it seems like a reinvent to me).
Python seems to be reinventing itself quite rapidly these days. It still seems like a pretty awkward host for stats problems — my friend that uses R is much more efficient than me in many ways. Maybe one of these days one of these Julia things will finally pull me over to that ship (I know you’ve been an advocate here for it for a while), or maybe Python will keep catching up to these other languages.
Is someone aware of not-too-old benchmarks of stan, Turing.jl and jax? I could not find anything. I would guess that the speedup from Turing to stan is probably not too great. Maybe factor 2? For some models maybe more? No idea really. So I guess I am not sure in which cases this tool will be useful. Using an external tool to create stan code form Julia code, running the model in stan and then again analyzing output in julia seems like a lot of overhead for possible relatively little value.
https://sethaxen.github.io/PosteriorDB.jl/stable/
It’d be interesting to choose 3-4 PosteriorDB models that take a nontrivial wall time and compare them in Stan vs Julia vs Jax.
I don’t have the time to do this entire project, nor do I currently have a working Stan (though I’m sure I could install one) but I’d be willing to translate Stan models to Julia and provide suggestions or bug fixes as part of this effort if someone was interested in setting up the benchmark project.
Perhaps someone reading the blog, like a grad student or whatever would like to set up such a comparison, maybe make it available as a Codeberg or Github repo and post a link here.
I do have a bunch of Julia implementations of PosteriorDB models lying around (not using Turing.jl, but the “original” StanBlocks.jl, which was a simple but relatively efficient Julia implementation of Stan-like syntax), which should be as close to being “as efficiently implemented as the Stan models” as can reasonably be expected.
The implementations are here: https://github.com/nsiccha/StanBlocks.jl/blob/main/ext/PosteriorDBExt.jl
I also at some point had some benchmarks online, but thought I had taken them offline for various reasons, the main reason being that I had forgotten to turn BridgeStan’s `propto` argument to `false` for the non-AD run, making Stan look slower than it can be, see the Bridgestan manual (https://roualdes.us/bridgestan/latest/internals/details.html) for an explanation. However, the benchmarks are actually still online below, though they should be taken with a big helping of salt. The main caveats are a) the one mentioned above, and b) the Stan PosteriorDB implementations not trying to be as efficient as possible. Link to the outdated benchmarks: https://nsiccha.github.io/StanBlocks.jl/performance.html
There’s also this set of Turing.jl reimplementations of PosteriorDB models, but I don’t know how much care has been taken to make them efficient: https://github.com/JasonPekos/TuringPosteriorDB.jl
Maybe someone wants to take the above as a starting point to do some benchmarks – be warned though that it’s difficult to do right.
> So I guess I am not sure in which cases this tool will be useful.
I do think that there are very few people who like Julia’s syntax as much as I do but still have/want to produce Stan models which can’t be generated by brms. That’s also one of the reasons why I’ve never bothered to make it less buggy or more user friendly – I don’t anticipate many people benefiting from that effort.
That being said, I do actually always use StanBlocks.jl, and never use Turing.jl or write Stan models manually…
Niko recently used StanBlocks.jl on a couple of our projects, and we were able to build quite complicated Stan models fairly quickly, so that’s a win. None of us is brave enough to use StanBlocks on our own, but we are thinking a lot about model composability, which is relatively hard in vanilla Stan. It seems our industry (biotech) has less and less patience with long projects, so anything we can do to increase modeling efficiency, rather than say computational efficiency, is important.
On the Julia front, I have been avoiding all J languages for a long time, but Niko very subtly convinced me to try it, and coming from R, I was pleasantly surprised. I wrote a short post about it (https://www.generable.com/post/porting-sir-ode-model-to-julia-1), in which I ported a simple SIR model from R/Stan to Julia/Turing. I still like R, probably because I am used to it, but if I were starting out today, I would go with Julia or Python.
Thanks for sharing, Andrew!