Do not re-evaluate model for Prior (#2644)

* Allow Prior to skip model re-evaluation

* remove unneeded `default_chain_type` method

* add a test

* add a likelihood term too

* why not test correctness while we're at it

Author: penelopeysm

src/mcmc/Inference.jl

@@ -136,7 +136,7 @@ struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
     stat::N
 
     """
-        Transition(model::Model, vi::AbstractVarInfo, sampler_transition)
+        Transition(model::Model, vi::AbstractVarInfo, sampler_transition; reevaluate=true)
 
     Construct a new `Turing.Inference.Transition` object using the outputs of a
     sampler step.
@@ -148,17 +148,38 @@ struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
 
     `sampler_transition` is the transition object returned by the sampler
     itself and is only used to extract statistics of interest.
+
+    By default, the model is re-evaluated in order to obtain values of:
+      - the values of the parameters as per user parameterisation (`vals_as_in_model`)
+      - the various components of the log joint probability (`logprior`, `loglikelihood`)
+    that are guaranteed to be correct.
+
+    If you **know** for a fact that the VarInfo `vi` already contains this information,
+    then you can set `reevaluate=false` to skip the re-evaluation step.
+
+    !!! warning
+        Note that in general this is unsafe and may lead to wrong results.
+
+    If `reevaluate` is set to `false`, it is the caller's responsibility to ensure that
+    the `VarInfo` passed in has `ValuesAsInModelAccumulator`, `LogPriorAccumulator`,
+    and `LogLikelihoodAccumulator` set up with the correct values. Note that the
+    `ValuesAsInModelAccumulator` must also have `include_colon_eq == true`, i.e. it
+    must be set up to track `x := y` statements.
     """
-    function Transition(model::DynamicPPL.Model, vi::AbstractVarInfo, sampler_transition)
-        vi = DynamicPPL.setaccs!!(
-            vi,
-            (
-                DynamicPPL.ValuesAsInModelAccumulator(true),
-                DynamicPPL.LogPriorAccumulator(),
-                DynamicPPL.LogLikelihoodAccumulator(),
-            ),
-        )
-        _, vi = DynamicPPL.evaluate!!(model, vi)
+    function Transition(
+        model::DynamicPPL.Model, vi::AbstractVarInfo, sampler_transition; reevaluate=true
+    )
+        if reevaluate
+            vi = DynamicPPL.setaccs!!(
+                vi,
+                (
+                    DynamicPPL.ValuesAsInModelAccumulator(true),
+                    DynamicPPL.LogPriorAccumulator(),
+                    DynamicPPL.LogLikelihoodAccumulator(),
+                ),
+            )
+            _, vi = DynamicPPL.evaluate!!(model, vi)
+        end
 
         # Extract all the information we need
         vals_as_in_model = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
@@ -175,12 +196,18 @@ struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
     function Transition(
         model::DynamicPPL.Model,
         untyped_vi::DynamicPPL.VarInfo{<:DynamicPPL.Metadata},
-        sampler_transition,
+        sampler_transition;
+        reevaluate=true,
     )
         # Re-evaluating the model is unconscionably slow for untyped VarInfo. It's
         # much faster to convert it to a typed varinfo first, hence this method.
         # https://github.com/TuringLang/Turing.jl/issues/2604
-        return Transition(model, DynamicPPL.typed_varinfo(untyped_vi), sampler_transition)
+        return Transition(
+            model,
+            DynamicPPL.typed_varinfo(untyped_vi),
+            sampler_transition;
+            reevaluate=reevaluate,
+        )
     end
 end
 

src/mcmc/prior.jl

@@ -16,8 +16,15 @@ function AbstractMCMC.step(
     sampling_model = DynamicPPL.contextualize(
         model, DynamicPPL.SamplingContext(rng, DynamicPPL.SampleFromPrior(), model.context)
     )
-    _, vi = DynamicPPL.evaluate!!(sampling_model, VarInfo())
-    return Transition(model, vi, nothing), nothing
+    vi = VarInfo()
+    vi = DynamicPPL.setaccs!!(
+        vi,
+        (
+            DynamicPPL.ValuesAsInModelAccumulator(true),
+            DynamicPPL.LogPriorAccumulator(),
+            DynamicPPL.LogLikelihoodAccumulator(),
+        ),
+    )
+    _, vi = DynamicPPL.evaluate!!(sampling_model, vi)
+    return Transition(model, vi, nothing; reevaluate=false), nothing
 end
-
-DynamicPPL.default_chain_type(sampler::Prior) = MCMCChains.Chains

test/mcmc/Inference.jl

@@ -142,6 +142,29 @@ using Turing
             @test mean(x[:s][1] for x in chains) ≈ 3 atol = 0.11
             @test mean(x[:m][1] for x in chains) ≈ 0 atol = 0.1
         end
+
+        @testset "accumulators are set correctly" begin
+            # Prior() uses `reevaluate=false` when constructing a
+            # `Turing.Inference.Transition`, so we had better make sure that it
+            # does capture colon-eq statements, as we can't rely on the default
+            # `Transition` constructor to do this for us.
+            @model function coloneq()
+                x ~ Normal()
+                10.0 ~ Normal(x)
+                z := 1.0
+                return nothing
+            end
+            chain = sample(coloneq(), Prior(), N)
+            @test chain isa MCMCChains.Chains
+            @test all(x -> x == 1.0, chain[:z])
+            # And for the same reason we should also make sure that the logp
+            # components are correctly calculated.
+            @test isapprox(chain[:logprior], logpdf.(Normal(), chain[:x]))
+            @test isapprox(chain[:loglikelihood], logpdf.(Normal.(chain[:x]), 10.0))
+            @test isapprox(chain[:lp], chain[:logprior] .+ chain[:loglikelihood])
+            # And that the outcome is not influenced by the likelihood
+            @test mean(chain, :x) ≈ 0.0 atol = 0.1
+        end
     end
 
     @testset "chain ordering" begin