Enforce re-evaluation when constructing Transition

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -73,14 +73,9 @@ function DynamicPPL.initialstep(
     steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
 
-    # Update the variables.
-    vi = DynamicPPL.unflatten(vi, Q.q)
-    # TODO(DPPL0.37/penelopeysm): This is obviously incorrect. Fix this.
-    vi = DynamicPPL.setloglikelihood!!(vi, Q.ℓq)
-    vi = DynamicPPL.setlogprior!!(vi, 0.0)
-
     # Create first sample and state.
-    sample = Turing.Inference.Transition(model, vi)
+    vi = DynamicPPL.unflatten(vi, Q.q)
+    sample = Turing.Inference.Transition(model, vi, nothing)
     state = DynamicNUTSState(ℓ, vi, Q, steps.H.κ, steps.ϵ)
 
     return sample, state
@@ -99,12 +94,9 @@ function AbstractMCMC.step(
     steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ℓ, state.stepsize)
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
-    # Update the variables.
-    vi = DynamicPPL.unflatten(vi, Q.q)
-    vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
-
     # Create next sample and state.
-    sample = Turing.Inference.Transition(model, vi)
+    vi = DynamicPPL.unflatten(vi, Q.q)
+    sample = Turing.Inference.Transition(model, vi, nothing)
     newstate = DynamicNUTSState(ℓ, vi, Q, state.metric, state.stepsize)
 
     return sample, newstate

src/mcmc/Inference.jl

@@ -124,75 +124,94 @@ end
 ######################
 # Default Transition #
 ######################
-# Default
-getstats(t) = nothing
+getstats(::Any) = NamedTuple()
 
+# TODO(penelopeysm): Remove this abstract type by converting SGLDTransition,
+# SMCTransition, and PGTransition to Turing.Inference.Transition instead.
 abstract type AbstractTransition end
 
-struct Transition{T,F<:AbstractFloat,S<:Union{NamedTuple,Nothing}} <: AbstractTransition
+struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
     θ::T
-    lp::F # TODO: merge `lp` with `stat`
-    stat::S
-end
+    logprior::F
+    loglikelihood::F
+    stat::N
+
+    """
+        Transition(model::Model, vi::AbstractVarInfo, sampler_transition)
+
+    Construct a new `Turing.Inference.Transition` object using the outputs of a
+    sampler step.
+
+    Here, `vi` represents a VarInfo _for which the appropriate parameters have
+    already been set_. However, the accumulators (e.g. logp) may in general
+    have junk contents. The role of this method is to re-evaluate `model` and
+    thus set the accumulators to the correct values.
+
+    `sampler_transition` is the transition object returned by the sampler
+    itself and is only used to extract statistics of interest.
+    """
+    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)
+
+        # Extract all the information we need
+        vals_as_in_model = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
+        logprior = DynamicPPL.getlogprior(vi)
+        loglikelihood = DynamicPPL.getloglikelihood(vi)
+
+        # Get additional statistics
+        stats = getstats(sampler_transition)
+        return new{typeof(vals_as_in_model),typeof(logprior),typeof(stats)}(
+            vals_as_in_model, logprior, loglikelihood, stats
+        )
+    end
 
-Transition(θ, lp) = Transition(θ, lp, nothing)
-function Transition(model::DynamicPPL.Model, vi::AbstractVarInfo, t)
-    # TODO(DPPL0.37/penelopeysm): Fix this
-    θ = getparams(model, vi)
-    lp = getlogjoint_internal(vi)
-    return Transition(θ, lp, getstats(t))
+    function Transition(
+        model::DynamicPPL.Model,
+        untyped_vi::DynamicPPL.VarInfo{<:DynamicPPL.Metadata},
+        sampler_transition,
+    )
+        # 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)
+    end
 end
 
-# TODO(DPPL0.37/penelopeysm): Add log-prior and log-likelihood terms as well
 function metadata(t::Transition)
-    stat = t.stat
-    if stat === nothing
-        return (lp=t.lp,)
-    else
-        return merge((lp=t.lp,), stat)
-    end
+    return merge(
+        t.stat,
+        (
+            lp=t.logprior + t.loglikelihood,
+            logprior=t.logprior,
+            loglikelihood=t.loglikelihood,
+        ),
+    )
+end
+function metadata(vi::AbstractVarInfo)
+    return (
+        lp=DynamicPPL.getlogjoint(vi),
+        logprior=DynamicPPL.getlogp(vi),
+        loglikelihood=DynamicPPL.getloglikelihood(vi),
+    )
 end
-
-# TODO(DPPL0.37/penelopeysm): Fix this
-DynamicPPL.getlogjoint(t::Transition) = t.lp
-
-# Metadata of VarInfo object
-# TODO(DPPL0.37/penelopeysm): Add log-prior and log-likelihood terms as well
-metadata(vi::AbstractVarInfo) = (lp=getlogjoint(vi),)
 
 ##########################
 # Chain making utilities #
 ##########################
 
-"""
-    getparams(model, t)
-
-Return a named tuple of parameters.
-"""
-getparams(model, t) = t.θ
-function getparams(model::DynamicPPL.Model, vi::DynamicPPL.VarInfo)
-    # NOTE: In the past, `invlink(vi, model)` + `values_as(vi, OrderedDict)` was used.
-    # Unfortunately, using `invlink` can cause issues in scenarios where the constraints
-    # of the parameters change depending on the realizations. Hence we have to use
-    # `values_as_in_model`, which re-runs the model and extracts the parameters
-    # as they are seen in the model, i.e. in the constrained space. Moreover,
-    # this means that the code below will work both of linked and invlinked `vi`.
-    # Ref: https://github.com/TuringLang/Turing.jl/issues/2195
-    # NOTE: We need to `deepcopy` here to avoid modifying the original `vi`.
-    return DynamicPPL.values_as_in_model(model, true, deepcopy(vi))
+getparams(::DynamicPPL.Model, t::AbstractTransition) = t.θ
+function getparams(model::DynamicPPL.Model, vi::AbstractVarInfo)
+    t = Transition(model, vi, nothing)
+    return getparams(model, t)
 end
-function getparams(
-    model::DynamicPPL.Model, untyped_vi::DynamicPPL.VarInfo{<:DynamicPPL.Metadata}
-)
-    # values_as_in_model is unconscionably slow for untyped VarInfo. It's
-    # much faster to convert it to a typed varinfo before calling getparams.
-    # https://github.com/TuringLang/Turing.jl/issues/2604
-    return getparams(model, DynamicPPL.typed_varinfo(untyped_vi))
-end
-function getparams(::DynamicPPL.Model, ::DynamicPPL.VarInfo{NamedTuple{(),Tuple{}}})
-    return Dict{VarName,Any}()
-end
-
 function _params_to_array(model::DynamicPPL.Model, ts::Vector)
     names_set = OrderedSet{VarName}()
     # Extract the parameter names and values from each transition.
@@ -208,7 +227,6 @@ function _params_to_array(model::DynamicPPL.Model, ts::Vector)
             iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
             mapreduce(collect, vcat, iters)
         end
-
         nms = map(first, nms_and_vs)
         vs = map(last, nms_and_vs)
         for nm in nms
@@ -224,7 +242,7 @@ function _params_to_array(model::DynamicPPL.Model, ts::Vector)
 end
 
 function get_transition_extras(ts::AbstractVector{<:VarInfo})
-    valmat = reshape([getlogjoint(t) for t in ts], :, 1)
+    valmat = reshape([DynamicPPL.getlogjoint(t) for t in ts], :, 1)
     return [:lp], valmat
 end
 
@@ -466,16 +484,17 @@ function transitions_from_chain(
     chain::MCMCChains.Chains;
     sampler=DynamicPPL.SampleFromPrior(),
 )
-    vi = Turing.VarInfo(model)
+    vi = VarInfo(model)
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     transitions = map(iters) do (sample_idx, chain_idx)
         # Set variables present in `chain` and mark those NOT present in chain to be resampled.
+        # TODO(DPPL0.37/penelopeysm): Aargh! setval_and_resample!!!! Burn this!!!
         DynamicPPL.setval_and_resample!(vi, chain, sample_idx, chain_idx)
         model(rng, vi, sampler)
 
         # Convert `VarInfo` into `NamedTuple` and save.
-        Transition(model, vi)
+        Transition(model, vi, nothing)
     end
 
     return transitions

src/mcmc/emcee.jl

@@ -65,7 +65,7 @@ function AbstractMCMC.step(
     end
 
     # Compute initial transition and states.
-    transition = map(Base.Fix1(Transition, model), vis)
+    transition = [Transition(model, vi, nothing) for vi in vis]
 
     # TODO: Make compatible with immutable `AbstractVarInfo`.
     state = EmceeState(
@@ -92,13 +92,12 @@ function AbstractMCMC.step(
     )
 
     # Compute the next states.
-    states = last(AbstractMCMC.step(rng, densitymodel, spl.alg.ensemble, state.states))
+    t, states = AbstractMCMC.step(rng, densitymodel, spl.alg.ensemble, state.states)
 
     # Compute the next transition and state.
     transition = map(states) do _state
         vi = DynamicPPL.unflatten(vi, _state.params)
-        t = Transition(getparams(model, vi), _state.lp)
-        return t
+        return Transition(model, vi, t)
     end
     newstate = EmceeState(vi, states)
 

src/mcmc/ess.jl

@@ -31,7 +31,7 @@ function DynamicPPL.initialstep(
         EllipticalSliceSampling.isgaussian(typeof(dist)) ||
             error("ESS only supports Gaussian prior distributions")
     end
-    return Transition(model, vi), vi
+    return Transition(model, vi, nothing), vi
 end
 
 function AbstractMCMC.step(
@@ -56,7 +56,7 @@ function AbstractMCMC.step(
     vi = DynamicPPL.unflatten(vi, sample)
     vi = DynamicPPL.setloglikelihood!!(vi, state.loglikelihood)
 
-    return Transition(model, vi), vi
+    return Transition(model, vi, nothing), vi
 end
 
 # Prior distribution of considered random variable

src/mcmc/external_sampler.jl

@@ -22,8 +22,6 @@ There are a few more optional functions which you can implement to improve the i
 - `Turing.Inference.isgibbscomponent(::MySampler)`: If you want your sampler to function as a component in Turing's Gibbs sampler, you should make this evaluate to `true`.
 
 - `Turing.Inference.requires_unconstrained_space(::MySampler)`: If your sampler requires unconstrained space, you should return `true`. This tells Turing to perform linking on the VarInfo before evaluation, and ensures that the parameter values passed to your sampler will always be in unconstrained (Euclidean) space.
-
-- `Turing.Inference.getlogp_external(external_transition, external_state)`: Tell Turing how to extract the log probability density associated with this transition (and state). If you do not specify these, Turing will simply re-evaluate the model with the parameters obtained from `getparams`, which can be inefficient. It is therefore recommended to store the log probability density in either the transition or the state (or both) and override this method.
 """
 struct ExternalSampler{S<:AbstractSampler,AD<:ADTypes.AbstractADType,Unconstrained} <:
        InferenceAlgorithm
@@ -85,27 +83,21 @@ function externalsampler(
     return ExternalSampler(sampler, adtype, Val(unconstrained))
 end
 
-"""
-    getlogp_external(external_transition, external_state)
-
-Get the log probability density associated with the external sampler's
-transition and state. Returns `missing` by default; in this case, an extra
-model evaluation will be needed to calculate the correct log density.
-"""
-getlogp_external(::Any, ::Any) = missing
-getlogp_external(mh::AdvancedMH.Transition, ::AdvancedMH.Transition) = mh.lp
-getlogp_external(hmc::AdvancedHMC.Transition, ::AdvancedHMC.HMCState) = hmc.stat.log_density
-
-struct TuringState{S,V1<:AbstractVarInfo,M,V}
+# TODO(penelopeysm): Can't we clean this up somehow?
+struct TuringState{S,V1,M,V}
     state::S
-    # Note that this varinfo has the correct parameters and logp obtained from
-    # the state, whereas `ldf.varinfo` will in general have junk inside it.
+    # Note that this varinfo must have the correct parameters set; but logp
+    # does not matter as it will be re-evaluated
     varinfo::V1
+    # Note that in general the VarInfo inside this LogDensityFunction will have
+    # junk parameters and logp. It only exists to provide structure
     ldf::DynamicPPL.LogDensityFunction{M,V}
 end
 
-varinfo(state::TuringState) = state.varinfo
-varinfo(state::AbstractVarInfo) = state
+# get_varinfo should return something from which the correct parameters can be
+# obtained, hence we use state.varinfo rather than state.ldf.varinfo
+get_varinfo(state::TuringState) = state.varinfo
+get_varinfo(state::AbstractVarInfo) = state
 
 getparams(::DynamicPPL.Model, transition::AdvancedHMC.Transition) = transition.z.θ
 function getparams(model::DynamicPPL.Model, state::AdvancedHMC.HMCState)
@@ -115,27 +107,6 @@ getstats(transition::AdvancedHMC.Transition) = transition.stat
 
 getparams(::DynamicPPL.Model, transition::AdvancedMH.Transition) = transition.params
 
-function make_updated_varinfo(
-    f::DynamicPPL.LogDensityFunction, external_transition, external_state
-)
-    # Set the parameters.
-    new_parameters = getparams(f.model, external_state)
-    new_varinfo = DynamicPPL.unflatten(f.varinfo, new_parameters)
-    # Set (or recalculate, if needed) the log density.
-    new_logp = getlogp_external(external_transition, external_state)
-    return if ismissing(new_logp)
-        last(DynamicPPL.evaluate!!(f.model, new_varinfo, f.context))
-    else
-        # TODO(DPPL0.37/penelopeysm) This is obviously wrong. Note that we
-        # have the same problem here as in HMC in that the sampler doesn't
-        # tell us about how logp is broken down into prior and likelihood.
-        # We should probably just re-evaluate unconditionally. A bit
-        # unfortunate.
-        DynamicPPL.setlogprior!!(new_varinfo, 0.0)
-        DynamicPPL.setloglikelihood!!(new_varinfo, new_logp)
-    end
-end
-
 # TODO: Do we also support `resume`, etc?
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
@@ -182,13 +153,10 @@ function AbstractMCMC.step(
         )
     end
 
-    # Get the parameters and log density, and set them in the varinfo.
-    new_varinfo = make_updated_varinfo(f, transition_inner, state_inner)
-
-    # Update the `state`
+    new_parameters = getparams(f.model, state_inner)
+    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
     return (
-        Transition(f.model, new_varinfo, transition_inner),
-        TuringState(state_inner, new_varinfo, f),
+        Transition(f.model, new_vi, transition_inner), TuringState(state_inner, new_vi, f)
     )
 end
 
@@ -207,12 +175,9 @@ function AbstractMCMC.step(
         rng, AbstractMCMC.LogDensityModel(f), sampler, state.state; kwargs...
     )
 
-    # Get the parameters and log density, and set them in the varinfo.
-    new_varinfo = make_updated_varinfo(f, transition_inner, state_inner)
-
-    # Update the `state`
+    new_parameters = getparams(f.model, state_inner)
+    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
     return (
-        Transition(f.model, new_varinfo, transition_inner),
-        TuringState(state_inner, new_varinfo, f),
+        Transition(f.model, new_vi, transition_inner), TuringState(state_inner, new_vi, f)
     )
 end