Turing models as distributions

[mohamed82008]

I am opening this issue to propose a way to treat Turing models as distributions. After TuringLang/Turing.jl#997, we will have a nice syntax to query the logpdf of a NamedTuple of symbols given another NamedTuple of symbols, e.g. logprob"x = [1, 2] | y = [1, 1], model = mymodel". This takes care of the logpdf function in the Distributions API. So if we can also have a way to sample a NamedTuple from the probability distribution P(x | y= [1, 1], model = mymodel) and the likes, we can define a simple syntax that generates such distribution structs from Turing models, something like prob"x | y = [1, 1], model = mymodel". Defining the posterior predictive distribution from there is also somewhat easy. The question is how to sample NamedTuples from Turing models.

Initial proposal

Let the inner function called when calling an instance of the Turing.Model be:

function inner_function(vi, sampler, ctx, model)
     #function body
end

Let's assume we have random variables s and m in the model. The idea I initially had was to replace the above expression with:

function inner_function(vi, sampler, ctx, model)
    if ctx isa DistributionContext
        s = missing
        m = missing
    end
    @inline function f()
        #function body
    end
    model_output = f()
    if ctx isa DistributionContext
        return (s = s, m = m)
    else
        return model_output
    end
end

Semantics

Calling the model in the DistributionContext will therefore return us the NamedTuple of random variables sampled. The DistributionContext can wrap another context and forward the tilde and dot_tilde functions to the context wrapped. This will give us a way to return the sampled random variables from the model. If a random variable wasn't sampled, missing will be returned. If a vector or matrix random variable can be partially sampled, i.e. the number of random variables is itself random, we can tell the users to initialize such variables using:

x = Vector{Union{Missing, T}}(undef, 10)`

for example inside the model body. The unsampled random variables will therefore be missing. These semantics for "sampling from Turing models" seem reasonable to me even in the case when the number of random variables is itself random.

Performance and alternative implementation

Given that we inline f, I thought the Julia compiler should be able to optimize out s = missing and m = missing if s and m were defined inside #function body, because the compiler does optimize these away when I don't use a closure. Unfortunately, the following minimal example doesn't infer properly:

function f(a)
    # initialize
    if a isa Int
        b = missing
    end

    @inline function g()
        # model body starts
        b = [1, 1]
        # Can have more stuff here and multiple return statements in branches
        return 2*b
        # model body ends
    end
    out = g() # overwrite `b` if `a isa Int`

    # Return a namedtuple or the output of the model depending on the type of the input `a` (or ctx in the real example)
    if a isa Int
        return (b = b,)
    else
        return out
    end
end
@code_warntype f(1) # f(1.0)

and according to Keno Fischer on Slack, this is "a limitation that needs to be addressed properly at some point". The problem is basically the closure.

So my alternative solution is to do:

function inner_function(vi, sampler, ctx, model)
    if ctx isa DistributionContext
        s = missing
        m = missing
    end
    #function body
    model_output = nothing
    @label end_of_func
    if ctx isa DistributionContext
        return (s = s, m = m)
    else
        return model_output
    end
end

and inside #function body, we replace every return ... statement with:

model_output = ...
@goto end_of_func

and every return statement with:

@goto end_of_func

This eliminates the use of a closure and therefore should infer properly. At least, the following minimal example does:

function f(a)
    # initialize
    if a isa Int
        b = missing
    end

    # model body starts
    # can have arbitrary code here with multiple return points, some or all the random variables can be defined
    b = [1, 1]
    if sum(b) > 0.5
        out = 2*b
        @goto end_of_func
    else
        out = 3*b
        @goto end_of_func
    end
    # model body ends
    out = nothing # the function output is nothing if we made it here

    @label end_of_func

    # Return a namedtuple or the output of the model depending on the type of the input `a` (or ctx in the real example)    
    if a isa Int
        return (b = b,)
    else
        return out
    end
end
@code_warntype f(1)

[devmotion]

I'm a bit confused about the first example. It seems it's not necessary to define g and in particular to evaluate g if a is a Int? The following alternative implementation seems to work (even though I think it would be nice to define f(a::Int) separately in this case):

function f(a)
    if a isa Int
        return (b = missing,)
    else
        function g()
            b = [1, 1]
            return 2*b
        end
        out = g()
        return out
    end
end
@code_warntype f(1) # f(1.0)

I'm also a bit confused that you have to define g at all, but maybe that's due to some requirement in Turing that is not obvious from the MWE? At least in the MWE it seems more reasonable to just write

b = [1, 1]
return 2 * b

I'm also slightly confused about the last example. What's the purpose of setting b = missing for Ints there? You set b = [1, 1] anyway for all inputs below - and hence it will give different results than the first example.

[mohamed82008]

Sorry, I didn't make all the requirements clear. Here are some requirements:

• Firstly, the user must be allowed to use return to kill the current sample at any time, or to return some values from the model when the user calls the model in any context other than the DistributionContext. This is why I initially used a closure to allow the user to use return. When the context is not a DistributionContext, the output returned by the model body is returned instead of the NamedTuple. The model body is #function body above, this has all the ~ expressions.
• Returning (b = missing,) when a isa Int is not the purpose of this function. By the end of the function, b will be [1, 1] if a isa Int (the equivalent of ctx isa DistributionContext in the actual code). The idea is that inside the model, we can overwrite the random variables (e.g. imagine b ~ Normal() instead of b = [1, 1]) which are initialized to missing. But if we don't overwrite b, because we kill the sample early, then a value of missing is used for that random variable. So b = [1, 1] or b ~ Normal() may or may not be there in the actual model or closure.

In the MWE, I merely tried to make the functions as difficult to compile as the real functions I am planning to implement to see if the Julia compiler works as I expect it to or not. Please tell me if something is not clear or doesn't make sense.

[mohamed82008]

I added some comments to the MWEs to explain my intentions.

[devmotion]

Thanks, I think I understand now (better) what you want to achieve in the MWEs. I was mainly confused since I missed that @inline will actually allow b to be changed, even in the case a isa Int. However, IMO that's not a reliable approach since, as far as I know, @inline gives only "a hint to the compiler that this function is worth inlining" but does not enforce it.

[mohamed82008]

No b will be over-written whether we use @inline or not.

julia> function f()
           b = 1
           function g()
               b = 2
           end
           g()
           return b
       end
f (generic function with 1 method)

julia> f()
2

[mohamed82008]

In fact, b will be defined in the outer function whether or not a is Int! I think this is because the scope of b is decided at parse-time before the if statement optimization kicks in, at compile-time.

[devmotion]

Very interesting, thanks! 👍 It's even mentioned in the documentation: https://docs.julialang.org/en/v1/manual/variables-and-scoping/#Local-Scope-1 (together with the performance issues in https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-captured-1)

I'm not sure if I understand your last comment though. The following example fails for f(1.0) (since b is not defined) but works as expected for f(1):

function f(a)
    if a isa Int
        b = missing
    end
    function g()
        b = b === missing ? 10 : [1, 1]
        return 2 * b
    end
    out = g()
    return out
end

[mohamed82008]

This works though:

function f(a)
    if a isa Int
        b = missing
    end
    function g()
        b = [1, 1]
    end
    g()
    return b
end
f(1)

That's what I meant. In your code, the scope of b is still local to f (not g) but you check b === missing before b was defined in f.

[devmotion]

Returning (b = missing,) when a isa Int is not the purpose of this function. By the end of the function, b will be [1, 1] if a isa Int (the equivalent of ctx isa DistributionContext in the actual code). The idea is that inside the model, we can overwrite the random variables (e.g. imagine b ~ Normal() instead of b = [1, 1]) which are initialized to missing. But if we don't overwrite b, because we kill the sample early, then a value of missing is used for that random variable. So b = [1, 1] or b ~ Normal() may or may not be there in the actual model or closure.

The general problem that I don't understand yet is that in the more Turing-specific setting you don't know what the random variables are before running the model - and even if you run the model, it might be that due to some early return statements you don't sample all variables that could possibly be sampled according to the model definition. Or in the context of your example above, you don't know that you would have to set b = missing and if you should initialize some other variables with missing as well.

To me it seems we can't avoid this problem even with static analysis due to the additional branch that is based on runtime information only when determining if some tilde statement is an assumption. Maybe I'm missing something or don't fully understand the intention yet, but if I'm right that would mean that this approach of returning missing values is just not feasible.

So I'm suggesting we just return a named tuple of samples that are actually sampled without worrying about if we did not sample some random variables in a specific run. If that can happen in a model, then the user should know about it and handle it in the downstream code by, e.g., checking if the named tuple contains a sample of the variable at interest or not.

[mohamed82008]

Well any variable next to ~ and not in the model arguments can be safely initialized to missing. If a model argument is to be assumed then it will also have an initial value missing. As for partially missing variables, how to add those to the returned named tuple can be tricky. We can add a check using VarInfo that will compile away when using TypedVarInfo but will be type unstable for UntypedVarInfo. This should tell us whether we need to return some arguments or not based on whether their symbols are found in VarInfo.

[mohamed82008]

early return statements

This is why we need to replace the early return statements by goto as proposed above.

[devmotion]

These two points,

As for partially missing variables, how to add those to the returned named tuple can be tricky.

and

early return statements

are exactly the reason for why I think we should just return whatever we sample.

If the sampled variables vary due to stochastic control flow, IMO the cleanest solution is to let the user handle that who defined the model. I just think there's no reason for trying to be smart here since it seems for any "smart" algorithms there are cases where it could fail and just returning the samples without adding missing values is sufficient for any downstream analysis.