Don't unflatten inside Transition constructor

TuringLang/DynamicPPL.jl#1001

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -74,7 +74,8 @@ function DynamicPPL.initialstep(
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
 
     # Create first sample and state.
-    sample = Turing.Inference.Transition(model, vi, Q.q, nothing)
+    vi = DynamicPPL.unflatten(vi, Q.q)
+    sample = Turing.Inference.Transition(model, vi, nothing)
     state = DynamicNUTSState(ℓ, vi, Q, steps.H.κ, steps.ϵ)
 
     return sample, state
@@ -94,7 +95,8 @@ function AbstractMCMC.step(
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Create next sample and state.
-    sample = Turing.Inference.Transition(model, vi, Q.q, nothing)
+    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

@@ -137,31 +137,20 @@ struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
     stat::N
 
     """
-        Transition(model::Model, vi::AbstractVarInfo, params::AbstractVector, sampler_transition)
+        Transition(model::Model, vi::AbstractVarInfo, sampler_transition)
 
-    Construct a new `Turing.Inference.Transition` object using the outputs of a sampler step.
+    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`.
+    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.
-
-    !!! 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.
+    `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,
-        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)
+    function Transition(model::DynamicPPL.Model, vi::AbstractVarInfo, sampler_transition)
         vi = DynamicPPL.setaccs!!(
             vi,
             (
@@ -193,18 +182,16 @@ struct Transition{T,F<:AbstractFloat,N<:NamedTuple} <: AbstractTransition
             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
-        )
+        return Transition(model, DynamicPPL.typed_varinfo(untyped_vi), sampler_transition)
     end
 end
 
@@ -232,7 +219,7 @@ end
 
 getparams(::DynamicPPL.Model, t::AbstractTransition) = t.θ
 function getparams(model::DynamicPPL.Model, vi::AbstractVarInfo)
-    t = Transition(model, vi, vi[:], nothing)
+    t = Transition(model, vi, nothing)
     return getparams(model, t)
 end
 function _params_to_array(model::DynamicPPL.Model, ts::Vector)
@@ -502,16 +489,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 = [Transition(model, vi, vi[:], nothing) for vi in vis]
+    transition = [Transition(model, vi, nothing) for vi in vis]
 
     # TODO: Make compatible with immutable `AbstractVarInfo`.
     state = EmceeState(
@@ -97,7 +97,7 @@ function AbstractMCMC.step(
     # Compute the next transition and state.
     transition = map(states) do _state
         vi = DynamicPPL.unflatten(vi, _state.params)
-        return Transition(model, vi, _state.params, 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[:], nothing), 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[:], nothing), vi
+    return Transition(model, vi, nothing), vi
 end
 
 # Prior distribution of considered random variable

src/mcmc/external_sampler.jl

@@ -148,10 +148,8 @@ function AbstractMCMC.step(
     end
 
     new_parameters = getparams(f.model, state_inner)
-    return (
-        Transition(f.model, f.varinfo, new_parameters, transition_inner),
-        TuringState(state_inner, f),
-    )
+    vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
+    return (Transition(f.model, vi, transition_inner), TuringState(state_inner, f))
 end
 
 function AbstractMCMC.step(
@@ -170,8 +168,6 @@ function AbstractMCMC.step(
     )
 
     new_parameters = getparams(f.model, state_inner)
-    return (
-        Transition(f.model, f.varinfo, new_parameters, transition_inner),
-        TuringState(state_inner, f),
-    )
+    vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
+    return (Transition(f.model, vi, 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, vi[:], nothing), GibbsState(vi, states)
+    return Transition(model, 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, vi[:], nothing), GibbsState(vi, states)
+    return Transition(model, 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, vi[:], nothing), GibbsState(vi, states)
+    return Transition(model, 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, vi[:], nothing), GibbsState(vi, states)
+    return Transition(model, vi, nothing), GibbsState(vi, states)
 end
 
 """

src/mcmc/hmc.jl

@@ -234,15 +234,15 @@ function DynamicPPL.initialstep(
         )
     end
 
-    # Update VarInfo based on acceptance
+    # Update VarInfo parameters based on acceptance
     new_params = if t.stat.is_accept
         t.z.θ
     else
         theta
     end
     vi = DynamicPPL.unflatten(vi, new_params)
 
-    transition = Transition(model, vi, new_params, t)
+    transition = Transition(model, vi, t)
     state = HMCState(vi, 1, kernel, hamiltonian, t.z, adaptor)
 
     return transition, state
@@ -283,14 +283,11 @@ function AbstractMCMC.step(
     # Update variables
     vi = state.vi
     if t.stat.is_accept
-        new_params = t.z.θ
         vi = DynamicPPL.unflatten(vi, new_params)
-    else
-        new_params = vi[:]
     end
 
     # Compute next transition and state.
-    transition = Transition(model, vi, new_params, t)
+    transition = Transition(model, vi, t)
     newstate = HMCState(vi, i, kernel, hamiltonian, t.z, state.adaptor)
 
     return transition, newstate

src/mcmc/is.jl

@@ -31,14 +31,14 @@ DynamicPPL.initialsampler(sampler::Sampler{<:IS}) = sampler
 function DynamicPPL.initialstep(
     rng::AbstractRNG, model::Model, spl::Sampler{<:IS}, vi::AbstractVarInfo; kwargs...
 )
-    return Transition(model, vi, vi[:], nothing), nothing
+    return Transition(model, vi, nothing), nothing
 end
 
 function AbstractMCMC.step(
     rng::Random.AbstractRNG, model::Model, spl::Sampler{<:IS}, ::Nothing; kwargs...
 )
     vi = VarInfo(rng, model, spl)
-    return Transition(model, vi, vi[:], nothing), nothing
+    return Transition(model, vi, nothing), nothing
 end
 
 # Calculate evidence.

src/mcmc/mh.jl

@@ -364,7 +364,7 @@ function DynamicPPL.initialstep(
     # just link everything before sampling.
     vi = maybe_link!!(vi, spl, spl.alg.proposals, model)
 
-    return Transition(model, vi, vi[:], nothing), vi
+    return Transition(model, vi, nothing), vi
 end
 
 function AbstractMCMC.step(
@@ -375,7 +375,7 @@ function AbstractMCMC.step(
     # 2. A bunch of NamedTuples that specify the proposal space
     new_vi = propose!!(rng, vi, model, spl, spl.alg.proposals)
 
-    return Transition(model, new_vi, new_vi[:], nothing), new_vi
+    return Transition(model, new_vi, nothing), new_vi
 end
 
 ####

src/mcmc/prior.jl

@@ -17,7 +17,7 @@ function AbstractMCMC.step(
         model, DynamicPPL.SamplingContext(rng, DynamicPPL.SampleFromPrior(), model.context)
     )
     _, vi = DynamicPPL.evaluate!!(sampling_model, VarInfo())
-    return Transition(model, vi, vi[:], nothing), nothing
+    return Transition(model, vi, nothing), nothing
 end
 
 DynamicPPL.default_chain_type(sampler::Prior) = MCMCChains.Chains