fix externalsampler bug

Author: penelopeysm

src/mcmc/external_sampler.jl

@@ -122,17 +122,20 @@ function externalsampler(
 end
 
 # TODO(penelopeysm): Can't we clean this up somehow?
-struct TuringState{S,V,L<:DynamicPPL.LogDensityFunction}
+struct TuringState{S,V,P<:AbstractVector,L<:DynamicPPL.LogDensityFunction}
     state::S
     # Note that this varinfo is used only for structure. Its parameters and other info do
     # not need to be accurate
     varinfo::V
+    # These are the actual parameters that this state is at
+    params::P
     ldf::L
 end
 
-# 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 must return something from which the correct parameters can be obtained
+function get_varinfo(state::TuringState)
+    return DynamicPPL.unflatten(state.varinfo, state.params)
+end
 get_varinfo(state::AbstractVarInfo) = state
 
 function AbstractMCMC.step(
@@ -188,7 +191,7 @@ function AbstractMCMC.step(
     new_stats = AbstractMCMC.getstats(state_inner)
     return (
         DynamicPPL.ParamsWithStats(new_parameters, f, new_stats),
-        TuringState(state_inner, varinfo, f),
+        TuringState(state_inner, varinfo, new_parameters, f),
     )
 end
 
@@ -211,6 +214,6 @@ function AbstractMCMC.step(
     new_stats = AbstractMCMC.getstats(state_inner)
     return (
         DynamicPPL.ParamsWithStats(new_parameters, f, new_stats),
-        TuringState(state_inner, state.varinfo, f),
+        TuringState(state_inner, state.varinfo, new_parameters, f),
     )
 end

src/mcmc/gibbs.jl

@@ -456,14 +456,11 @@ function AbstractMCMC.step_warmup(
 end
 
 """
-    setparams_varinfo!!(model, sampler::AbstractSampler, state, params::AbstractVarInfo)
+    setparams_varinfo!!(model::DynamicPPL.Model, sampler::AbstractSampler, state, params::AbstractVarInfo)
 
 A lot like AbstractMCMC.setparams!!, but instead of taking a vector of parameters, takes an
 `AbstractVarInfo` object. Also takes the `sampler` as an argument. By default, falls back to
 `AbstractMCMC.setparams!!(model, state, params[:])`.
-
-`model` is typically a `DynamicPPL.Model`, but can also be e.g. an
-`AbstractMCMC.LogDensityModel`.
 """
 function setparams_varinfo!!(
     model::DynamicPPL.Model, ::AbstractSampler, state, params::AbstractVarInfo
@@ -488,12 +485,18 @@ function setparams_varinfo!!(
 end
 
 function setparams_varinfo!!(
-    ::DynamicPPL.Model, ::ExternalSampler, state::TuringState, params::AbstractVarInfo
+    model::DynamicPPL.Model,
+    sampler::ExternalSampler,
+    state::TuringState,
+    params::AbstractVarInfo,
 )
+    new_ldf = DynamicPPL.LogDensityFunction(
+        model, DynamicPPL.getlogjoint_internal, params; adtype=sampler.adtype
+    )
     new_inner_state = AbstractMCMC.setparams!!(
-        AbstractMCMC.LogDensityModel(state.ldf), state.state, params[:]
+        AbstractMCMC.LogDensityModel(new_ldf), state.state, params[:]
     )
-    return TuringState(new_inner_state, params, state.ldf)
+    return TuringState(new_inner_state, params, params[:], new_ldf)
 end
 
 function setparams_varinfo!!(