User experience: Python vs R, PyMC vs Stan, PyTensor vs JAX

[bob-carpenter]

@ricardoV94

nutpie for example is pushing some normalizing flow-related automatic reparametrizations that would be harder in a Stan-like language where the core object is a sum of densities and not a graphical model.

Adrian’s draft paper hints that he couldn’t get the black-box version working past 1K dimensions but hopes that using model structure somehow will help. Depending on what the input is around model structure, it may or may not be possible with Stan. It’ll certainly be easier to use model structure automatically in PyMC.

@ricardoV94

PyMC and numpyro are specific PPLs built on top of general python computational packages: PyTensor and JAX, respectively.

I get the part about (embedded) PPLs, but I thought PyMC was more flexible about output format. Thomas Wiecki sent me a notebook producing JAX and some examples under pymc_experimental using Blackjax, but the latter are now 404. Is that just embedding into JAX rather than translating into JAX? We can embed Stan models in JAX that way, but it doesn’t let us jit them or run them on GPUs.

@ricardoV94

If you wanted to implement Stan in Python you would probably go for something like that and not build on top of an existing PPL.

If JAX had existed when we started (2010), we probably would have used that. The big issue is R integration—the stats community is still largely weddd to R. I, on the other hand, have moved to Python.

My point is really that with something like JAX, you don’t need a PPL to write models the way we write them in Stan. Stan’s really just a differentiable programming language with some convenient stats abstractions and a simple syntax. ll Stan’s doing anyway other than giving you a way to structure code and compute posterior predictive quantities in the same languag eas the model. But with JAX, we can just vmap the output, so we don’t even need that block. Similarly, we can just define functions in JAX, so we don’t need that block. And pytree makes it very easy to structure arguments. Oryx is now a pretty extensive transform library and JAX and many additional packages have a pretty big special function library. One thing I haven’t looked into is all the solvers we use in Stan, like algebraic equation solvers, root finders, 1D integrators, ODE integrators, etc.

@ricardoV94

things like graph manipulation, algebraic rewrites, numerical optimizations that compilers don’t dare to touch.

We do this for Stan, too. In using JAX straight up, that’s all on the coder and the JAX developers, which is going to be rough on the statisticians. I’m not sure we ever got a lot out of optimizations in Stan, but then optimizations should be easier if there’s a graphical model base.

@ricardoV94

@aseyboldt would probably suggest you use JAX to get model densities and gradients, but write your own sampler in something that gives you more control over CUDA streams.

That’s exactly what we’re doing.