Fix more tests

Author: penelopeysm

src/mcmc/is.jl

@@ -46,9 +46,10 @@ function getlogevidence(samples::Vector{<:Transition}, ::Sampler{<:IS}, state)
     return logsumexp(map(x -> x.loglikelihood, samples)) - log(length(samples))
 end
 
-struct ISContext{R<:AbstractRNG}
+struct ISContext{R<:AbstractRNG} <: DynamicPPL.AbstractContext
     rng::R
 end
+DynamicPPL.NodeTrait(::ISContext) = DynamicPPL.IsLeaf()
 
 function DynamicPPL.tilde_assume!!(ctx::ISContext, dist::Distribution, vn::VarName, vi)
     if haskey(vi, vn)

src/mcmc/mh.jl

@@ -415,3 +415,8 @@ function DynamicPPL.tilde_assume!!(
     end
     return DynamicPPL.tilde_assume!!(dispatch_ctx, right, vn, vi)
 end
+function DynamicPPL.tilde_observe!!(
+    ::MHContext, right::Distribution, left, vn::Union{VarName,Nothing}, vi::AbstractVarInfo
+)
+    return DynamicPPL.tilde_observe!!(DefaultContext(), right, left, vn, vi)
+end

test/ad.jl

@@ -154,31 +154,23 @@ end
 # context, and then call check_adtype on the result before returning the results from the
 # child context.
 
-function DynamicPPL.tilde_assume(context::ADTypeCheckContext, right, vn, vi)
-    value, vi = DynamicPPL.tilde_assume(DynamicPPL.childcontext(context), right, vn, vi)
-    check_adtype(context, vi)
-    return value, vi
-end
-
-function DynamicPPL.tilde_assume(
-    rng::Random.AbstractRNG, context::ADTypeCheckContext, sampler, right, vn, vi
+function DynamicPPL.tilde_assume!!(
+    context::ADTypeCheckContext, right::Distribution, vn::VarName, vi::AbstractVarInfo
 )
-    value, vi = DynamicPPL.tilde_assume(
-        rng, DynamicPPL.childcontext(context), sampler, right, vn, vi
-    )
+    value, vi = DynamicPPL.tilde_assume!!(DynamicPPL.childcontext(context), right, vn, vi)
     check_adtype(context, vi)
     return value, vi
 end
 
-function DynamicPPL.tilde_observe!!(context::ADTypeCheckContext, right, left, vi)
-    left, vi = DynamicPPL.tilde_observe!!(DynamicPPL.childcontext(context), right, left, vi)
-    check_adtype(context, vi)
-    return left, vi
-end
-
-function DynamicPPL.tilde_observe!!(context::ADTypeCheckContext, sampler, right, left, vi)
+function DynamicPPL.tilde_observe!!(
+    context::ADTypeCheckContext,
+    right::Distribution,
+    left,
+    vn::Union{VarName,Nothing},
+    vi::AbstractVarInfo,
+)
     left, vi = DynamicPPL.tilde_observe!!(
-        DynamicPPL.childcontext(context), sampler, right, left, vi
+        DynamicPPL.childcontext(context), right, left, vn, vi
     )
     check_adtype(context, vi)
     return left, vi

test/mcmc/gibbs.jl

@@ -695,7 +695,7 @@ end
                 # Determine initial parameters to make comparison as fair as possible.
                 # posterior_mean returns a NamedTuple so we can plug it in directly.
                 posterior_mean = DynamicPPL.TestUtils.posterior_mean(model)
-                initial_params = fill(InitFromParams(initial_params), num_chains)
+                initial_params = fill(InitFromParams(posterior_mean), num_chains)
 
                 # Sampler to use for Gibbs components.
                 hmc = HMC(0.1, 32)