Minijyro - a toy Pyro-like PPL

Hi everyone,

I implemented a small PPL, Minijyro.jl, that is based on the ideas of effect handlers in Pyro and thought it might be of interest to people. This project is mostly me trying to learn more about Julia and probabilistic programming so it’s not meant to be a language that is seriously used by anyone. Here’s a quick linear regression example (sampling is done with AdvancedHMC.jl):

@jyro function model(xs)
    D = size(xs)[2]
    w ~ MvNormal(zeros(D), I)
    y ~ MvNormal(xs * w, 0.1*I)
end

cond_model = condition(model, Dict(:y => y_obs))
samples, stats = nuts(cond_model, (X,), 1000)

Effect handlers offer a really handy abstraction to implement many common operations in probabilistic programming such as conditioning, computing the density function, model reparametrization and many more!

They also seem similar to the context system that is used in Turing but maybe someone from the Turing team can give a description of the context system used there and how it differs from effect handlers?

Some high-level implementation details can be found in the README of the repo. I based my implementation a lot on the Mini-Pyro implementation so the code might look a bit Pythonic. I’m very happy for any suggestions to make the code more “Julian”.

Please let me know if you have any questions!

7 Likes

This is very cool!

Sounds great, looking forward to checking it out!

Turing contexts and Pyro effect handlers seem almost identical. Both target modifying the behavior of a probabilistic program without manually modifying it. Turing contexts have zero run-time overhead though thanks to the beautiful Julia compiler.

2 Likes

That sounds great! I was wondering how to achieve zero run-time overhead so I will have a closer look at the Turing internals.

The zero run-time overhead can be achieved because Julia compiles a different method for each context type. There is nothing Turing-specific here. We just have different types for different contexts and use multiple dispatch to define the specialized behavior of each context. Then Julia does the rest.

1 Like