Less surprising model return values

src/DynamicPPL.jl

@@ -69,7 +69,6 @@ export  AbstractVarInfo,
         getargnames,
         getdefaults,
         getgenerator,
-        runmodel!,
 # Samplers
         Sampler,
         SampleFromPrior,

src/compiler.jl

@@ -278,7 +278,7 @@ function generate_tilde(left, right, model_info)
     ctx = model_info[:main_body_names][:ctx]
     sampler = model_info[:main_body_names][:sampler]
 
-    @gensym tmpright
+    @gensym tmpright tmpleft
     top = [:($tmpright = $right),
            :($tmpright isa Union{$Distribution,AbstractVector{<:$Distribution}}
              || throw(ArgumentError($DISTMSG)))]
@@ -290,8 +290,8 @@ function generate_tilde(left, right, model_info)
         assumption = [
             :($out = $(DynamicPPL.tilde_assume)($ctx, $sampler, $tmpright, $vn, $inds,
                                                 $vi)),
-            :($left = $out[1]),
-            :($(DynamicPPL.acclogp!)($vi, $out[2]))
+            :($(DynamicPPL.acclogp!)($vi, $out[2])),
+            :($left = $out[1])
         ]
 
         # It can only be an observation if the LHS is an argument of the model
@@ -303,11 +303,13 @@ function generate_tilde(left, right, model_info)
                 if $isassumption
                     $(assumption...)
                 else
+                    $tmpleft = $left
                     $(DynamicPPL.acclogp!)(
                         $vi,
-                        $(DynamicPPL.tilde_observe)($ctx, $sampler, $tmpright, $left, $vn,
+                        $(DynamicPPL.tilde_observe)($ctx, $sampler, $tmpright, $tmpleft, $vn,
                                                     $inds, $vi)
                     )
+                    $tmpleft
                 end
             end
         end
@@ -321,10 +323,12 @@ function generate_tilde(left, right, model_info)
     # If the LHS is a literal, it is always an observation
     return quote
         $(top...)
+        $tmpleft = $left
         $(DynamicPPL.acclogp!)(
             $vi,
-            $(DynamicPPL.tilde_observe)($ctx, $sampler, $tmpright, $left, $vi)
+            $(DynamicPPL.tilde_observe)($ctx, $sampler, $tmpright, $tmpleft, $vi)
         )
+        $tmpleft
     end
 end
 
@@ -341,7 +345,7 @@ function generate_dot_tilde(left, right, model_info)
     ctx = model_info[:main_body_names][:ctx]
     sampler = model_info[:main_body_names][:sampler]
 
-    @gensym tmpright
+    @gensym tmpright tmpleft
     top = [:($tmpright = $right),
            :($tmpright isa Union{$Distribution,AbstractVector{<:$Distribution}}
              || throw(ArgumentError($DISTMSG)))]
@@ -353,8 +357,8 @@ function generate_dot_tilde(left, right, model_info)
         assumption = [
             :($out = $(DynamicPPL.dot_tilde_assume)($ctx, $sampler, $tmpright, $left,
                                                     $vn, $inds, $vi)),
-            :($left .= $out[1]),
-            :($(DynamicPPL.acclogp!)($vi, $out[2]))
+            :($(DynamicPPL.acclogp!)($vi, $out[2])),
+            :($left .= $out[1])
         ]
 
         # It can only be an observation if the LHS is an argument of the model
@@ -366,11 +370,13 @@ function generate_dot_tilde(left, right, model_info)
                 if $isassumption
                     $(assumption...)
                 else
+                    $tmpleft = $left
                     $(DynamicPPL.acclogp!)(
                         $vi,
-                        $(DynamicPPL.dot_tilde_observe)($ctx, $sampler, $tmpright, $left,
+                        $(DynamicPPL.dot_tilde_observe)($ctx, $sampler, $tmpright, $tmpleft,
                                                         $vn, $inds, $vi)
                     )
+                    $tmpleft
                 end
             end
         end
@@ -384,10 +390,12 @@ function generate_dot_tilde(left, right, model_info)
     # If the LHS is a literal, it is always an observation
     return quote
         $(top...)
+        $tmpleft = $left
         $(DynamicPPL.acclogp!)(
             $vi,
-            $(DynamicPPL.dot_tilde_observe)($ctx, $sampler, $tmpright, $left, $vi)
+            $(DynamicPPL.dot_tilde_observe)($ctx, $sampler, $tmpright, $tmpleft, $vi)
         )
+        $tmpleft
     end
 end
 
@@ -443,7 +451,6 @@ function build_output(model_info)
             $ctx::$(DynamicPPL.AbstractContext),
         )
             $unwrap_data_expr
-            $(DynamicPPL.resetlogp!)($vi)
             $main_body
         end
 

src/model.jl

@@ -109,34 +109,34 @@ function Model{missings}(
     return Model{missings}(model.f, args, modelgen)
 end
 
+"""
+    (model::Model)([spl = SampleFromPrior(), ctx = DefaultContext()])
 
+Sample from `model` using the sampler `spl`.
+"""
 function (model::Model)(
-    vi::AbstractVarInfo=VarInfo(),
     spl::AbstractSampler=SampleFromPrior(),
     ctx::AbstractContext=DefaultContext()
 )
-    return model.f(model, vi, spl, ctx)
+    return model(VarInfo(), spl, ctx)
 end
 
-
 """
-    runmodel!(model::Model, vi::AbstractVarInfo[, spl::AbstractSampler, ctx::AbstractContext])
+    (model::Model)(vi::AbstractVarInfo[, spl = SampleFromPrior(), ctx = DefaultContext()])
 
 Sample from `model` using the sampler `spl` storing the sample and log joint probability in `vi`.
 Resets the `vi` and increases `spl`s `state.eval_num`.
 """
-function runmodel!(
-    model::Model,
+function (model::Model)(
     vi::AbstractVarInfo,
     spl::AbstractSampler=SampleFromPrior(),
     ctx::AbstractContext=DefaultContext()
 )
-    setlogp!(vi, 0)
+    resetlogp!(vi)
     if has_eval_num(spl)
         spl.state.eval_num += 1
     end
-    model(vi, spl, ctx)
-    return vi
+    return model.f(model, vi, spl, ctx)
 end
 
 

src/varinfo.jl

@@ -107,7 +107,7 @@ const TypedVarInfo = VarInfo{<:NamedTuple}
 
 function VarInfo(model::Model, ctx = DefaultContext())
     vi = VarInfo()
-    runmodel!(model, vi, SampleFromPrior(), ctx)
+    model(vi, SampleFromPrior(), ctx)
     return TypedVarInfo(vi)
 end
 

test/Turing/Turing.jl

@@ -16,7 +16,7 @@ using Markdown, Libtask, MacroTools
 using Tracker: Tracker
 
 import Base: ~, ==, convert, hash, promote_rule, rand, getindex, setindex!
-import DynamicPPL: getspace, runmodel!
+import DynamicPPL: getspace
 
 const PROGRESS = Ref(true)
 function turnprogress(switch::Bool)

test/Turing/contrib/inference/dynamichmc.jl

@@ -60,11 +60,11 @@ function AbstractMCMC.sample_init!(
         gradient_logp(x, spl.state.vi, model, spl)
     end
 
-    runmodel!(model, spl.state.vi, SampleFromUniform())
+    model(spl.state.vi, SampleFromUniform())
 
     if spl.selector.tag == :default
         link!(spl.state.vi, spl)
-        runmodel!(model, spl.state.vi, spl)
+        model(spl.state.vi, spl)
     end
 
     # Set the parameters to a starting value.
@@ -145,4 +145,4 @@ function AbstractMCMC.psample(
     end
     return AbstractMCMC.psample(rng, model, Sampler(alg, model), N, n_chains;
                                 chain_type=chain_type, progress=false, kwargs...)
-end
+end

test/Turing/contrib/inference/sghmc.jl

@@ -85,7 +85,7 @@ function step(
     Turing.DEBUG && @debug "X-> R..."
     if spl.selector.tag != :default
         link!(vi, spl)
-        runmodel!(model, vi, spl)
+        model(vi, spl)
     end
 
     Turing.DEBUG && @debug "recording old variables..."
@@ -198,7 +198,7 @@ function step(
     Turing.DEBUG && @debug "X-> R..."
     if spl.selector.tag != :default
         link!(vi, spl)
-        runmodel!(model, vi, spl)
+        model(vi, spl)
     end
 
     Turing.DEBUG && @debug "recording old variables..."

test/Turing/core/Core.jl

@@ -6,7 +6,7 @@ using Distributions, LinearAlgebra
 using ..Utilities, Reexport
 using Tracker: Tracker
 using ..Turing: Turing
-using DynamicPPL: Model, runmodel!,
+using DynamicPPL: Model,
     AbstractSampler, Sampler, SampleFromPrior
 using LinearAlgebra: copytri!
 using Bijectors: PDMatDistribution

test/Turing/core/ad.jl

@@ -92,7 +92,8 @@ function gradient_logp(
     logp_old = getlogp(vi)
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        logp = getlogp(runmodel!(model, new_vi, sampler))
+        model(new_vi, sampler)
+        logp = getlogp(new_vi)
         setlogp!(vi, ForwardDiff.value(logp))
         return logp
     end
@@ -119,7 +120,8 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        return getlogp(runmodel!(model, new_vi, sampler))
+        model(new_vi, sampler)
+        return getlogp(new_vi)
     end
 
     # Compute forward and reverse passes.

test/Turing/core/compat/zygote.jl

@@ -16,7 +16,8 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        return getlogp(runmodel!(model, new_vi, sampler))
+        model(new_vi, sampler)
+        return getlogp(new_vi)
     end
 
     # Compute forward and reverse passes.

test/Turing/inference/Inference.jl

@@ -4,7 +4,7 @@ using ..Core, ..Utilities
 using DynamicPPL: Metadata, _tail, VarInfo, TypedVarInfo, 
     islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize, 
     settrans!, _getvns, getdist, CACHERESET, AbstractSampler,
-    Model, runmodel!, Sampler, SampleFromPrior, SampleFromUniform,
+    Model, Sampler, SampleFromPrior, SampleFromUniform,
     Selector, AbstractSamplerState, DefaultContext, PriorContext,
     LikelihoodContext, MiniBatchContext, set_flag!, unset_flag!, NamedDist, NoDist
 using Distributions, Libtask, Bijectors

test/Turing/inference/ess.jl

@@ -76,7 +76,7 @@ function AbstractMCMC.step!(
 
     # recompute log-likelihood in logp
     if spl.selector.tag !== :default
-        runmodel!(model, vi, spl)
+        model(vi, spl)
     end
 
     # define previous sampler state
@@ -117,7 +117,7 @@ function EllipticalSliceSampling.sample_prior(rng::Random.AbstractRNG, model::ES
     vi = spl.state.vi
     vns = _getvns(vi, spl)
     set_flag!(vi, vns[1][1], "del")
-    runmodel!(model.model, vi, spl)
+    model.model(vi, spl)
     return vi[spl]
 end
 
@@ -140,7 +140,7 @@ function Distributions.loglikelihood(model::ESSModel, f)
     spl = model.spl
     vi = spl.state.vi
     vi[spl] = f
-    runmodel!(model.model, vi, spl)
+    model.model(vi, spl)
     getlogp(vi)
 end
 

test/Turing/inference/hmc.jl

@@ -137,7 +137,7 @@ function AbstractMCMC.sample_init!(
     # non-Gibbs sampling.
     if !islinked(spl.state.vi, spl) && spl.selector.tag == :default
         link!(spl.state.vi, spl)
-        runmodel!(model, spl.state.vi, spl)
+        model(spl.state.vi, spl)
     end
 end
 
@@ -343,7 +343,7 @@ function AbstractMCMC.step!(
         # Transform the space
         Turing.DEBUG && @debug "X-> R..."
         link!(spl.state.vi, spl)
-        runmodel!(model, spl.state.vi, spl)
+        model(spl.state.vi, spl)
         # Update Hamiltonian
         metric = gen_metric(length(spl.state.vi[spl]), spl)
         ∂logπ∂θ = gen_∂logπ∂θ(spl.state.vi, spl, model)
@@ -413,7 +413,7 @@ function gen_logπ(vi::VarInfo, spl::Sampler, model)
     function logπ(x)::Float64
         x_old, lj_old = vi[spl], getlogp(vi)
         vi[spl] = x
-        runmodel!(model, vi, spl)
+        model(vi, spl)
         lj = getlogp(vi)
         vi[spl] = x_old
         setlogp!(vi, lj_old)

test/Turing/inference/mh.jl

@@ -105,7 +105,7 @@ function gen_logπ_mh(spl::Sampler, model)
         x_old, lj_old = vi[spl], getlogp(vi)
         # vi[spl] = x
         set_namedtuple!(vi, x)
-        runmodel!(model, vi)
+        model(vi)
         lj = getlogp(vi)
         vi[spl] = x_old
         setlogp!(vi, lj_old)
@@ -231,7 +231,7 @@ function AbstractMCMC.step!(
     kwargs...
 )
     if spl.selector.rerun # Recompute joint in logp
-        runmodel!(model, spl.state.vi)
+        model(spl.state.vi)
     end
 
     # Retrieve distribution and value NamedTuples.

test/independence.jl

@@ -7,8 +7,5 @@
         end
     end
     model = coinflip([1,1,0])
-
-    vi = VarInfo()
-
-    runmodel!(model, vi, SampleFromPrior(), LikelihoodContext())
+    model(SampleFromPrior(), LikelihoodContext())
 end

test/varinfo.jl

@@ -1,7 +1,7 @@
 using .Turing, Random
 using AbstractMCMC: step!
 using DynamicPPL: Selector, reconstruct, invlink, CACHERESET,
-    SampleFromPrior, Sampler, runmodel!, SampleFromUniform,
+    SampleFromPrior, Sampler, SampleFromUniform,
     _getidcs, set_retained_vns_del_by_spl!, is_flagged,
     set_flag!, unset_flag!, VarInfo, TypedVarInfo,
     getlogp, setlogp!, resetlogp!, acclogp!, vectorize,
@@ -114,19 +114,18 @@ include(dir*"/test/test_utils/AllUtils.jl")
         test_base!(vi)
         test_base!(empty!(TypedVarInfo(vi)))
     end
-    @testset "runmodel!" begin
-        # Test that eval_num is incremented when calling runmodel!
+    @testset "in-place" begin
+        # Test that eval_num is incremented when running the model
         @model testmodel() = begin
             x ~ Normal()
         end
         alg = HMC(0.1, 5)
         spl = Sampler(alg, testmodel())
-        vi = VarInfo()
         m = testmodel()
-        m(vi)
-        runmodel!(m, vi, spl)
+        vi = VarInfo(m)
+        m(vi, spl)
         @test spl.state.eval_num == 1
-        runmodel!(m, vi, spl)
+        m(vi, spl)
         @test spl.state.eval_num == 2
     end
     @testset "flags" begin