fix externalsampler in gibbs

Author: penelopeysm

src/mcmc/external_sampler.jl

@@ -83,14 +83,20 @@ function externalsampler(
     return ExternalSampler(sampler, adtype, Val(unconstrained))
 end
 
-struct TuringState{S,M,V}
+# TODO(penelopeysm): Can't we clean this up somehow?
+struct TuringState{S,V1,M,V}
     state::S
+    # 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
 
-get_varinfo(state::TuringState) = state.ldf.varinfo
+# 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.θ
@@ -148,8 +154,10 @@ function AbstractMCMC.step(
     end
 
     new_parameters = getparams(f.model, state_inner)
-    vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
-    return (Transition(f.model, vi, transition_inner), TuringState(state_inner, f))
+    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
+    return (
+        Transition(f.model, new_vi, transition_inner), TuringState(state_inner, new_vi, f)
+    )
 end
 
 function AbstractMCMC.step(
@@ -168,6 +176,8 @@ function AbstractMCMC.step(
     )
 
     new_parameters = getparams(f.model, state_inner)
-    vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
-    return (Transition(f.model, vi, transition_inner), TuringState(state_inner, f))
+    new_vi = DynamicPPL.unflatten(f.varinfo, new_parameters)
+    return (
+        Transition(f.model, new_vi, transition_inner), TuringState(state_inner, new_vi, f)
+    )
 end

src/mcmc/gibbs.jl

@@ -573,7 +573,7 @@ function setparams_varinfo!!(
     new_inner_state = setparams_varinfo!!(
         AbstractMCMC.LogDensityModel(logdensity), sampler, state.state, params
     )
-    return TuringState(new_inner_state, logdensity)
+    return TuringState(new_inner_state, params, logdensity)
 end
 
 function setparams_varinfo!!(