more fix fix fix

Author: penelopeysm

src/mcmc/is.jl

@@ -31,14 +31,20 @@ DynamicPPL.initialsampler(sampler::Sampler{<:IS}) = sampler
 function DynamicPPL.initialstep(
     rng::AbstractRNG, model::Model, spl::Sampler{<:IS}, vi::AbstractVarInfo; kwargs...
 )
-    return Transition(model, vi), nothing
+    # Need to manually construct the Transition here because we only
+    # want to use the likelihood.
+    xs = Turing.Inference.getparams(model, vi)
+    lp = DynamicPPL.getloglikelihood(vi)
+    return Transition(xs, lp, 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), nothing
+    xs = Turing.Inference.getparams(model, vi)
+    lp = DynamicPPL.getloglikelihood(vi)
+    return Transition(xs, lp, nothing), nothing
 end
 
 # Calculate evidence.
@@ -53,5 +59,6 @@ function DynamicPPL.assume(rng, ::Sampler{<:IS}, dist::Distribution, vn::VarName
         r = rand(rng, dist)
         vi = push!!(vi, vn, r, dist)
     end
-    return r, 0, vi
+    vi = accumulate_assume!!(vi, r, 0.0, vn, dist)
+    return r, vi
 end

src/mcmc/particle_mcmc.jl

@@ -35,7 +35,9 @@ function AdvancedPS.advance!(
     trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false
 )
     # Make sure we load/reset the rng in the new replaying mechanism
-    trace.model.f.varinfo = DynamicPPL.increment_num_produce!!(trace.model.f.varinfo)
+    trace = Accessors.@set trace.model.f.varinfo = DynamicPPL.increment_num_produce!!(
+        trace.model.f.varinfo
+    )
     isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng)
     score = consume(trace.model.ctask)
     if score === nothing
@@ -51,13 +53,11 @@ function AdvancedPS.delete_retained!(trace::TracedModel)
 end
 
 function AdvancedPS.reset_model(trace::TracedModel)
-    DynamicPPL.reset_num_produce!(trace.varinfo)
-    return trace
+    return Accessors.@set trace.varinfo = DynamicPPL.reset_num_produce!!(trace.varinfo)
 end
 
 function AdvancedPS.reset_logprob!(trace::TracedModel)
-    DynamicPPL.resetlogp!!(trace.model.varinfo)
-    return trace
+    return Accessors.@set trace.model.varinfo = DynamicPPL.resetlogp!!(trace.model.varinfo)
 end
 
 function AdvancedPS.update_rng!(