More test fixes

Author: penelopeysm

src/mcmc/emcee.jl

@@ -65,7 +65,7 @@ function AbstractMCMC.step(
     end
 
     # Compute initial transition and states.
-    transition = map(Base.Fix1(Transition, model), vis)
+    transition = [Transition(model, vi, vi[:], nothing) for vi in vis]
 
     # TODO: Make compatible with immutable `AbstractVarInfo`.
     state = EmceeState(

src/mcmc/is.jl

@@ -43,7 +43,7 @@ end
 
 # Calculate evidence.
 function getlogevidence(samples::Vector{<:Transition}, ::Sampler{<:IS}, state)
-    return logsumexp(map(x -> x.lp, samples)) - log(length(samples))
+    return logsumexp(map(x -> x.loglikelihood, samples)) - log(length(samples))
 end
 
 function DynamicPPL.assume(rng, ::Sampler{<:IS}, dist::Distribution, vn::VarName, vi)

src/mcmc/mh.jl

@@ -317,14 +317,9 @@ function propose!!(
         ),
     )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
-
-    # TODO: Make this compatible with immutable `VarInfo`.
-    # Update the values in the VarInfo.
-    # TODO(DPPL0.37/penelopeysm): This is obviously incorrect. We need to
-    # re-evaluate the model.
+    # trans.params isa NamedTuple
     set_namedtuple!(vi, trans.params)
-    vi = DynamicPPL.setloglikelihood!!(vi, trans.lp)
-    return DynamicPPL.setlogprior!!(vi, 0.0)
+    return vi
 end
 
 # Make a proposal if we DO have a covariance proposal matrix.
@@ -354,12 +349,8 @@ function propose!!(
         ),
     )
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
-
-    # TODO(DPPL0.37/penelopeysm): This is obviously incorrect. We need to
-    # re-evaluate the model.
-    vi = DynamicPPL.unflatten(vi, trans.params)
-    vi = DynamicPPL.setloglikelihood!!(vi, trans.lp)
-    return DynamicPPL.setlogprior!!(vi, 0.0)
+    # trans.params isa AbstractVector
+    return DynamicPPL.unflatten(vi, trans.params)
 end
 
 function DynamicPPL.initialstep(
@@ -382,9 +373,9 @@ function AbstractMCMC.step(
     # Cases:
     # 1. A covariance proposal matrix
     # 2. A bunch of NamedTuples that specify the proposal space
-    vi = propose!!(rng, vi, model, spl, spl.alg.proposals)
+    new_vi = propose!!(rng, vi, model, spl, spl.alg.proposals)
 
-    return Transition(model, vi, vi[:], nothing), vi
+    return Transition(model, new_vi, new_vi[:], nothing), new_vi
 end
 
 ####

src/mcmc/sghmc.jl

@@ -67,10 +67,7 @@ function DynamicPPL.initialstep(
     # Compute initial sample and state.
     sample = Transition(model, vi, vi[:], nothing)
     ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
+        model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.alg.adtype
     )
     state = SGHMCState(ℓ, vi, zero(vi[:]))
 
@@ -233,10 +230,7 @@ function DynamicPPL.initialstep(
     # Create first sample and state.
     sample = SGLDTransition(model, vi, zero(spl.alg.stepsize(0)))
     ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
+        model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.alg.adtype
     )
     state = SGLDState(ℓ, vi, 1)