Request for code review: stan2tfp python wrapper

Hi,

My name is Adam Haber and I’m a neuroscience PhD student. A couple of months ago I’ve joined the stanc3 team, mostly working on the Tensorflow Probability backend. I’ve started working on a small python library that wraps the stan2tfp compiler. The goal (the way I see it, at least - would be happy to hear other thoughts!) is to enable python users to compile Stan programs to TFP models “without leaving the notebook”. There’s a repo with an initial version, but it can probably be much better. Here’s a simple example.

If anyone’s interested in having a look, it’d be much appreciated. Thanks!

Adam

6 Likes

This looks like a posteriordb sorta thing: Beta-release Bayesian Posterior Database

Maybe you can hook in there and get models to test for correctness.

The advantage of this over default Stan would be taking advantage of Tensorflow’s scalability, right? So big data stuff?

1 Like

This looks like a posteriordb sorta thing: Beta-release Bayesian Posterior Database
Maybe you can hook in there and get models to test for correctness

Thanks! I’ll look into it.

The advantage of this over default Stan would be taking advantage of Tensorflow’s scalability, right? So big data stuff?

Sure. Additional reasons might be:

  1. Users who are used to TFP and/or like the TF ecosystem as a whole.
  2. Time-to-first-draw - the TFP code generation is nearly instantaneous.

There are probably other/better reasons, for me these seem like the main motivations…

1 Like

Looks very useful!

Ben mentioned correctness, but I would expect also that posteriordb and stan2tfp it would be very useful for testing efficiency and accuracy of alternative inference algorithms available for tfp.

It’d be compelling if there were models that could be fit in TFP but not in Stan. Or even ones that would be a lot faster in TFP.

The code gen’s always fast. It’s the compile time that kills us for the C++. It’s about 7s right now. How fast is TFP code ready to run from when it’s generated?

The code gen’s always fast. It’s the compile time that kills us for the C++. It’s about 7s right now. How fast is TFP code ready to run from when it’s generated?

Haven’t thoroughly profiled it but certainly less than a second for everything I’ve played with…

Eager or non-eager or both? How much does the size of autodiff tree affect non-eager?

Eager or non-eager or both? How much does the size of autodiff tree affect non-eager?

By “ready to run” I meant time until an instantiated TFP model object, on which .sample() can be called; basically just the time it takes the stan2tfp compiler to generate the TFP code + overhead. But you’re right, this doesn’t take the XLA compilation time into account, since the sampling function is compiled only when .sample() is called. I’m not sure how to profile the XLA compilation time separately from the actual sampling. I can ask in the TFP google group if this is of interest.

Also, AFAIK, TF2 is eager by default, and sampling without @tf.function(experimental_compile=True) is much, much slower.

1 Like

Hi!

This sounds super-interesting! In posteriordb our thought is to add TFP code (as well as other PPF) in the long run. If it is possible convert directly . between Stan and TFP that would simplify this alot!

I know to little about TFP to know if it is possible to convert every stan model or is there only some part of the Stan syntax that is handled. We could probably see if we can use posteriordb to also improve stan2tfp.

With kind regards
Måns

It’ll depend on how it’s implemented.

I imagine there are a lot of functions and gradients that only one of the two libraries (TFP or Stan) have implemented.

At the moment, about 40% of the distributions in Stan’s functions reference are supported by stan2tfp. This covers most of the “basic” distributions, which are also implemented by TFP (possibly with different parameterisations).

Roughly 20% more could be implemented once we’ll add support for simplex/ordered. The remaining 40% include distributions not available in TFP, different _glm and GPs .

1 Like