New Turing-libtask integration

[yebai]

This PR introduces a simpler way of evaluating models. It is to avoid the current complicated calling stack encountered in #1744. The benefit is that we no longer need to trace many boiler-plate functions such as evaluate!!. For example, model(...) ==> evaluate!! ==> _evaluate!! becomes

julia> using Turing, DynamicPPL, AdvancedPS

julia> @model gdemo(x, y) = begin
         # Assumptions
         σ ~ InverseGamma(2,3)
         μ ~ Normal(0,sqrt(σ))
         # Observations
         x ~ Normal(μ, sqrt(σ))
         y ~ Normal(μ, sqrt(σ))
       end
gdemo (generic function with 2 methods)

# Case 1: Sample from the prior. 

julia> m = Turing.Core.TracedModel(gdemo(1.5, 2.), SampleFromPrior(), VarInfo()) 

julia> f = m.evaluator[1];

julia> args = m.evaluator[2:end];

julia> f(args...) 
(2.0, VarInfo (2 variables (μ, σ), dimension 2; logp: -6.162))


using Libtask

t = Libtask.CTask(f, args...) 

consume(t) # work fine!


# Case 2: SMC sampler 

julia> m = Turing.Core.TracedModel(gdemo(1.5, 2.), Sampler(SMC(50)), VarInfo());

julia> t = Libtask.CTask(m.evaluator[1], m.evaluator[2:end]...);

julia> consume(t)
ERROR: KeyError: key :__trace not found
Stacktrace:
 [1] consume(ttask::TapedTask)
   @ Libtask ~/.julia/dev/Libtask/src/tapedtask.jl:164
 [2] top-level scope
   @ REPL[58]:1

So, with this PR, the integration between Turing-Libtask should be exactly the same as that for AdvancedPS-Libtask.

More tests (not needed until above example is fixed)

using Turing, Test, AbstractMCMC, DynamicPPL, Random

import AbstractMCMC.AbstractSampler

function check_numerical(chain,
    symbols::Vector,
    exact_vals::Vector;
    atol=0.2,
    rtol=0.0
)
    for (sym, val) in zip(symbols, exact_vals)
    E = val isa Real ?
    mean(chain[sym]) :
    vec(mean(chain[sym], dims=1))
    @info (symbol=sym, exact=val, evaluated=E)
    @test E ≈ val atol=atol rtol=rtol
    end
end

function check_MoGtest_default(chain; atol=0.2, rtol=0.0)
    check_numerical(chain,
        [:z1, :z2, :z3, :z4, :mu1, :mu2],
        [1.0, 1.0, 2.0, 2.0, 1.0, 4.0],
        atol=atol, rtol=rtol)
end

@model gdemo_d(x, y) = begin
    s ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s))
    x ~ Normal(m, sqrt(s))
    y ~ Normal(m, sqrt(s))
    return s, m
end

alg = CSMC(15)
chain = sample(gdemo_d(1.5, 2.0), alg, 5_00)
check_numerical(chain, [:s, :m], [49/24, 7/6], atol=0.1)
  
@model MoGtest(D) = begin
    mu1 ~ Normal(1, 1)
    mu2 ~ Normal(4, 1)
    z1 ~ Categorical(2)
    if z1 == 1
        D[1] ~ Normal(mu1, 1)
    else
        D[1] ~ Normal(mu2, 1)
    end
    z2 ~ Categorical(2)
    if z2 == 1
        D[2] ~ Normal(mu1, 1)
    else
        D[2] ~ Normal(mu2, 1)
    end
    z3 ~ Categorical(2)
    if z3 == 1
        D[3] ~ Normal(mu1, 1)
    else
        D[3] ~ Normal(mu2, 1)
    end
    z4 ~ Categorical(2)
    if z4 == 1
        D[4] ~ Normal(mu1, 1)
    else
        D[4] ~ Normal(mu2, 1)
    end
    z1, z2, z3, z4, mu1, mu2
end
  
MoGtest_default = MoGtest([1.0 1.0 4.0 4.0])

gibbs = Gibbs(PG(15, :z1, :z2, :z3, :z4), HMC(0.15, 3, :mu1, :mu2))
chain = sample(MoGtest_default, gibbs, 5_00);
check_MoGtest_default(chain, atol=0.15)

[yebai]

@KDr2 Could you take a look at why case 2 is broken?

[KDr2]

t = Libtask.CTask(m.evaluator[1], m.evaluator[2:end]...);

Below is the real bt of the error:

┌ Error: TapedTask Error:
│   exception =
│    KeyError: key :__trace not found
│    Stacktrace:
│      [1] getindex(d::IdDict{Any, Any}, key::Any)
│        @ Base ./iddict.jl:108
│      [2] current_trace()
│        @ AdvancedPS ~/.julia/packages/AdvancedPS/XPI6n/src/container.jl:64
│      [3] assume(rng::Random._GLOBAL_RNG, spl::Sampler{SMC{(), AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Int64}}}, dist::InverseGamma{Float64}, vn::VarName{:σ, Setfield.IdentityLens}, #unused#::UntypedVarInfo{DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}, Float64})
│        @ Turing.Inference ~/Work/julia/Turing.jl/src/inference/AdvancedSMC.jl:327
│      [4] tilde_assume
│        @ ~/.julia/packages/DynamicPPL/c8MjC/src/context_implementations.jl:49 [inlined]
│      [5] tilde_assume
│        @ ~/.julia/packages/DynamicPPL/c8MjC/src/context_implementations.jl:46 [inlined]
│      [6] tilde_assume
│        @ ~/.julia/packages/DynamicPPL/c8MjC/src/context_implementations.jl:31 [inlined]
│      [7] tilde_assume!!(context::SamplingContext{Sampler{SMC{(), AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Int64}}}, DefaultContext, Random._GLOBAL_RNG}, right::InverseGamma{Float64}, vn::VarName{:σ, Setfield.IdentityLens}, vi::UntypedVarInfo{DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}, Float64})
│        @ DynamicPPL ~/.julia/packages/DynamicPPL/c8MjC/src/context_implementations.jl:117
│      [8] (::Libtask.Instruction{typeof(Core._apply_iterate)})()
│        @ Libtask ~/Work/julia/Libtask.jl/src/tapedfunction.jl:84
│      [9] step_in(t::Libtask.Tape, args::Tuple{Model{typeof(gdemo), (:x, :y), (), (), Tuple{Float64, Float64}, Tuple{}, DefaultContext}, UntypedVarInfo{DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}, Float64}, SamplingContext{Sampler{SMC{(), AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Int64}}}, DefaultContext, Random._GLOBAL_RNG}, Float64, Float64})
│        @ Libtask ~/Work/julia/Libtask.jl/src/tapedtask.jl:66
│     [10] (::Libtask.var"#11#12"{Libtask.TapedFunction, Tuple{Model{typeof(gdemo), (:x, :y), (), (), Tuple{Float64, Float64}, Tuple{}, DefaultContext}, UntypedVarInfo{DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}, Float64}, SamplingContext{Sampler{SMC{(), AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Int64}}}, DefaultContext, Random._GLOBAL_RNG}, Float64, Float64}, Channel{Int64}, Channel{Any}})()
│        @ Libtask ./task.jl:123
└ @ Libtask ~/Work/julia/Libtask.jl/src/tapedtask.jl:31

[KDr2]

@yebai Do we need to copy the local storage of tasks when we copy CTasks?

[yebai]

I found the issue - we need to construct an AdvancedPS.Trace object so that task storage is correctly initialised with a backward trace reference. With a bit changes to AdvancedPS, I can now run case 2 successfully.

julia> using Turing, DynamicPPL, AdvancedPS, Libtask

julia> @model gdemo(x, y) = begin
         # Assumptions
         σ ~ InverseGamma(2,3)
         μ ~ Normal(0,sqrt(σ))
         # Observations
         x ~ Normal(μ, sqrt(σ))
         y ~ Normal(μ, sqrt(σ))
       end
gdemo (generic function with 2 methods)

julia> m = Turing.Core.TracedModel(gdemo(1.5, 2.), Sampler(SMC(50)), VarInfo());

julia> t = AdvancedPS.Trace(m);

julia> consume(t.ctask)
-4.143262128020936

julia> consume(t.ctask)
-5.481572280981258

# This also works

julia> t = Libtask.CTask(m.evaluator[1], m.evaluator[2:end]...);

julia> consume(t)
-2.5527145711306765

julia> consume(t)
-3.498088347524718

Now the remaining pieces of complete Turing-Libtask integration are to make sure t::Advanced.Trace can be copied (i.e. fork and forkr work correctly), then the original and copied traces can be consumed, in a similar way to CTask.

[KDr2]

Now the remaining pieces of complete Turing-Libtask integration are to make sure t::Advanced.Trace can be copied (i.e. fork and forkr work correctly), then the original and copied traces can be consumed, in a similar way to CTask.

I checked the old version of Libtask, it copies the local storage when a copy occurs. This commit added the same behavior: TuringLang/Libtask.jl@2fd124c

Can this ensure the Trace object to be correctly copied?

[yebai]

The code now runs correctly locally using TuringLang/AdvancedPS.jl#40 and TuringLang/Libtask.jl#102. More numerical fixes might be needed though.

[yebai]

@KDr2 There are only 4 errors left. After further debugging, I found that 3 errors are due to unsupported control flow from the same model (d2d315b), which can be ignored now. However, one error can not be attributed to known factors. It might be related to some unknown bug in the intercept function. Could you look into it, please? Here is an MWE.

julia> using Turing, Test, AbstractMCMC, DynamicPPL, Random, Turing.RandomMeasures, Libtask

julia> @model infiniteGMM(x) = begin
           # Hyper-parameters, i.e. concentration parameter and parameters of H.
           α = 1.0
           μ0 = 0.0
           σ0 = 1.0
           
           # Define random measure, e.g. Dirichlet process.
           rpm = DirichletProcess(α)
           
           # Define the base distribution, i.e. expected value of the Dirichlet process.
           H = Normal(μ0, σ0)
           
           # Latent assignment.
           z = tzeros(Int, length(x))
               
           # Locations of the infinitely many clusters.
           μ = tzeros(Float64, 0)
           
           for i in 1:length(x)
               
               # Number of clusters.
               K = maximum(z)
               nk = Vector{Int}(map(k -> sum(z .== k), 1:K))

               # Draw the latent assignment.
               z[i] ~ ChineseRestaurantProcess(rpm, nk)
               
               # Create a new cluster?
               if z[i] > K
                   push!(μ, 0.0)

                   # Draw location of new cluster.
                   μ[z[i]] ~ H
               end
                       
               # Draw observation.
               x[i] ~ Normal(μ[z[i]], 1.0)
           end
       end
infiniteGMM (generic function with 2 methods)

julia> # Generate some test data.
       Random.seed!(1);

julia> data = vcat(randn(10), randn(10) .- 5, randn(10) .+ 10);

julia> data .-= mean(data);

julia> data /= std(data);

julia> # MCMC sampling
       Random.seed!(2);

julia> iterations = 500;

julia> model_fun = infiniteGMM(data);

julia> m = Turing.Core.TracedModel(model_fun, Sampler(SMC(50)), VarInfo());

julia> f = m.evaluator[1];

julia> args = m.evaluator[2:end];

t = Libtask.CTask(f, args...) 

ERROR: TypeError: in new, expected Core.Box, got a value of type Libtask.Box{Core.Box}
Stacktrace:
  [1] macro expansion
    @ ./REPL[2]:23 [inlined]
  [2] macro expansion
    @ ~/.julia/packages/DynamicPPL/c8MjC/src/compiler.jl:529 [inlined]
  [3] ##300
    @ ./REPL[2]:19 [inlined]
  [4] var"##300"(arg1::typeof(infiniteGMM), arg2::Model{typeof(infiniteGMM), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DefaultContext}, arg3::UntypedVarInfo{DynamicPPL.Metadata{Dict{VarName, Int64}, Vector{Distribution}, Vector{VarName}, Vector{Real}, Vector{Set{DynamicPPL.Selector}}}, Float64}, arg4::SamplingContext{Sampler{SMC{(), AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Int64}}}, DefaultContext, Random._GLOBAL_RNG}, arg5::Vector{Float64})
    @ IRTools.Inner ~/.julia/packages/IRTools/isLV2/src/eval.jl:0
  [5] invokelatest(::Any, ::Any, ::Vararg{Any}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Base ./essentials.jl:716
  [6] invokelatest(::Any, ::Any, ::Vararg{Any})
    @ Base ./essentials.jl:714
  [7] evalir(::Module, ::IRTools.Inner.IR, ::Any, ::Vararg{Any})
    @ IRTools.Inner ~/.julia/packages/IRTools/isLV2/src/eval.jl:25
  [8] evalir(::IRTools.Inner.IR, ::Any, ::Any, ::Vararg{Any})
    @ IRTools.Inner ~/.julia/packages/IRTools/isLV2/src/eval.jl:26
  [9] (::Libtask.TapedFunction)(::Model{typeof(infiniteGMM), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DefaultContext}, ::Vararg{Any})
    @ Libtask ~/.julia/dev/Libtask/src/tapedfunction.jl:212
 [10] TapedTask(::Libtask.TapedFunction, ::Model{typeof(infiniteGMM), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DefaultContext}, ::Vararg{Any})
    @ Libtask ~/.julia/dev/Libtask/src/tapedtask.jl:31
 [11] TapedTask(::Function, ::Model{typeof(infiniteGMM), (:x,), (), (), Tuple{Vector{Float64}}, Tuple{}, DefaultContext}, ::Vararg{Any})
    @ Libtask ~/.julia/dev/Libtask/src/tapedtask.jl:64
 [12] top-level scope
    @ REPL[13]:1

[KDr2]

ERROR: TypeError: in new, expected Core.Box, got a value of type Libtask.Box{Core.Box}

I am on it already, it is caused by the new operator, I haven't figured out a way to put the new operator in an instruction yet...

[yebai]

I fixed the last CI error in 7c81350. Local tests are passing now - hopefully CI will also pass now!

[yebai]

All CI tests run without errors now (there're still numerical failures) - with TuringLang/Libtask.jl#107, we should be able to fix the remaining numerical failures associated with control flow soon.