Enforce re-evaluation when constructing Transition

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -73,14 +73,8 @@ 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)
+    sample = Turing.Inference.Transition(model, vi, Q.q, nothing)
     state = DynamicNUTSState(ℓ, vi, Q, steps.H.κ, steps.ϵ)
 
     return sample, state
@@ -99,12 +93,8 @@ 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)
+    sample = Turing.Inference.Transition(model, vi, Q.q, nothing)
     newstate = DynamicNUTSState(ℓ, vi, Q, state.metric, state.stepsize)
 
     return sample, newstate

src/mcmc/Inference.jl

@@ -124,85 +124,124 @@ 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
-
-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))
+    logprior::F
+    loglikelihood::F
+    stat::N
+
+    """
+        Transition(model::Model, vi::AbstractVarInfo, params::AbstractVector, sampler_transition)
+
+    Construct a new `Turing.Inference.Transition` object using the outputs of a sampler step.
+
+    Here, `vi` represents a VarInfo which in general may have junk contents (both
+    parameters and accumulators). The role of this method is to re-evaluate `model` by inserting
+    the new `params` (provided by the sampler) into the VarInfo `vi`.
+
+    `sampler_transition` is the transition object returned by the sampler itself and is only used
+    to extract statistics of interest.
+
+    !!! warning "Parameters must match varinfo linking status"
+        It is mandatory that the vector of parameters provided line up exactly with how the
+        VarInfo `vi` is linked. Otherwise, this can silently produce incorrect results.
+    """
+    function Transition(
+        model::DynamicPPL.Model,
+        vi::AbstractVarInfo,
+        parameters::AbstractVector,
+        sampler_transition,
+    )
+        # To be safe...
+        vi = deepcopy(vi)
+        # Set the parameters and re-evaluate with the appropriate accumulators
+        vi = DynamicPPL.unflatten(vi, parameters)
+        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)
+
+        # Convert values to the format needed (with individual VarNames split up).
+        # TODO(penelopeysm): This wouldn't be necessary if not for MCMCChains's poor
+        # representation...
+        iters = map(
+            DynamicPPL.varname_and_value_leaves,
+            keys(vals_as_in_model),
+            values(vals_as_in_model),
+        )
+        values_split = mapreduce(collect, vcat, iters)
+
+        # Get additional statistics
+        stats = getstats(sampler_transition)
+        return new{typeof(values_split),typeof(logprior),typeof(stats)}(
+            values_split, logprior, loglikelihood, stats
+        )
+    end
+    function Transition(
+        model::DynamicPPL.Model,
+        untyped_vi::DynamicPPL.VarInfo{<:DynamicPPL.Metadata},
+        parameters::AbstractVector,
+        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), parameters, 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`.
-    vals = DynamicPPL.values_as_in_model(model, true, deepcopy(vi))
-
-    # Obtain an iterator over the flattened parameter names and values.
-    iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
-
-    # Materialize the iterators and concatenate.
-    return mapreduce(collect, vcat, iters)
+getparams(::DynamicPPL.Model, t::AbstractTransition) = t.θ
+function getparams(model::DynamicPPL.Model, vi::AbstractVarInfo)
+    t = Transition(model, vi, 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 float(Real)[]
-end
-
 function _params_to_array(model::DynamicPPL.Model, ts::Vector)
     names_set = OrderedSet{VarName}()
     # Extract the parameter names and values from each transition.
     dicts = map(ts) do t
+        # TODO(penelopeysm): Get rid of AbstractVarInfo transitions. see 
+        # https://github.com/TuringLang/Turing.jl/issues/2631. That would
+        # allow us to just use t.θ here.
         nms_and_vs = getparams(model, t)
         nms = map(first, nms_and_vs)
         vs = map(last, nms_and_vs)
@@ -221,7 +260,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
 

src/mcmc/emcee.jl

@@ -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, _state.params, 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, 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, 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,26 +83,14 @@ 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}
+struct TuringState{S,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.
-    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::TuringState) = state.ldf.varinfo
 varinfo(state::AbstractVarInfo) = state
 
 getparams(::DynamicPPL.Model, transition::AdvancedHMC.Transition) = transition.z.θ
@@ -115,27 +101,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 +147,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)
     return (
-        Transition(f.model, new_varinfo, transition_inner),
-        TuringState(state_inner, new_varinfo, f),
+        Transition(f.model, varinfo, new_parameters, transition_inner),
+        TuringState(state_inner, f),
     )
 end
 
@@ -207,12 +169,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)
     return (
-        Transition(f.model, new_varinfo, transition_inner),
-        TuringState(state_inner, new_varinfo, f),
+        Transition(f.model, varinfo, new_parameters, transition_inner),
+        TuringState(state_inner, f),
     )
 end

src/mcmc/gibbs.jl

@@ -389,7 +389,7 @@ function AbstractMCMC.step(
         initial_params=initial_params,
         kwargs...,
     )
-    return Transition(model, vi), GibbsState(vi, states)
+    return Transition(model, vi, vi[:], nothing), GibbsState(vi, states)
 end
 
 function AbstractMCMC.step_warmup(
@@ -414,7 +414,7 @@ function AbstractMCMC.step_warmup(
         initial_params=initial_params,
         kwargs...,
     )
-    return Transition(model, vi), GibbsState(vi, states)
+    return Transition(model, vi, vi[:], nothing), GibbsState(vi, states)
 end
 
 """
@@ -502,7 +502,7 @@ function AbstractMCMC.step(
     vi, states = gibbs_step_recursive(
         rng, model, AbstractMCMC.step, varnames, samplers, states, vi; kwargs...
     )
-    return Transition(model, vi), GibbsState(vi, states)
+    return Transition(model, vi, vi[:], nothing), GibbsState(vi, states)
 end
 
 function AbstractMCMC.step_warmup(
@@ -522,7 +522,7 @@ function AbstractMCMC.step_warmup(
     vi, states = gibbs_step_recursive(
         rng, model, AbstractMCMC.step_warmup, varnames, samplers, states, vi; kwargs...
     )
-    return Transition(model, vi), GibbsState(vi, states)
+    return Transition(model, vi, vi[:], nothing), GibbsState(vi, states)
 end
 
 """