Why does Pyro use messengers in a trace instead of a computational graph

Hi all,

I have been using Pyro for some time and recently looked at the code base to understand how it’s working. I signed up on this forum just to ask this question. I noticed that pyro computes everything sequentially maintaining a trace and then back/forward propogates through this trace. Is there a reason for this choice? I am not sure but intuitively I feel like building a computational graph might be more efficient since we can also run minor optimizations on the graph and remove unnecessary computations (for e.g. sampling a variable that is not used anywhere in the graph).

What does the community think of this design choice and are there concrete reasons behind why it is implemented the way it is now?

Like PyTorch, we try to keep the experience of writing Pyro models as Pythonic as possible, and introducing a user-visible computational graph as in e.g. TensorFlow would break that experience. Many Pyro inference algorithms perform statistical optimizations like Rao-Blackwellization, but since the core Pyro team is fairly small, we try wherever possible to defer standard program transformations like dead code elimination that aren’t specific to probabilistic programs to the PyTorch autograd engine and JIT compiler (or to XLA in NumPyro).

We’re also working on replacing much of Pyro’s internals with Funsor, a first-order intermediate language for integration that’s probably closer to what you have in mind but doesn’t violate the two design goals above.

Hi eb8680_2! I am so sorry for the late reply. Thank you so much for letting me know of the funsor api. On a high level look, it does seem to be what I had in mind. I will take a deeper look at the repo. Since a graphical model is best represented as a graph, I was expecting to see a model graph in pyro and was just surprised to see that everything is done using traces. I hope I can contribute to the development, although I am not at all used to open source contribution. Nonetheless, these things are for github. Thanks again for sharing that information. Have a great week.

FYI if you’re interested in understanding Funsor, you should probably start with our paper “Functional Tensors for Probabilistic Programming”.

I’m glad to hear you’d like to contribute to Pyro development. We welcome all sorts of open source contributors - if you’re new to open source, I recommend having a look at GitHub’s nice intro to open source contribution. The best way to get started contributing to Pyro specifically is to look at open GitHub issues labeled “good first issue” or “help wanted.”

Funsor itself is still a bit of a science project, but is in dire need of more documentation and examples (the most valuable things open source contributors can do for almost any project). Note that it may be difficult to do more than that for Funsor right now unless you have some expertise in Bayesian inference or programming languages and compilers, in which case we can give you more specific guidance.

1 Like