fixes for MH

Author: penelopeysm

src/mcmc/mh.jl

@@ -125,7 +125,8 @@ function DynamicPPL.init(
     return DynamicPPL.UntransformedValue(rand(rng, dist))
 end
 
-function MH(vn_proposal_pairs...)
+function MH(pair1, pairs...)
+    vn_proposal_pairs = (pair1, pairs...)
     return MH(
         vnt -> begin
             proposals = DynamicPPL.VarNamedTuple()

src/mcmc/particle_mcmc.jl

@@ -334,90 +334,53 @@ function AbstractMCMC.step(
 end
 
 """
-    get_trace_local_varinfo_maybe(vi::AbstractVarInfo)
+    get_trace_local_varinfo()
 
 Get the varinfo stored in the 'taped globals' of a `Libtask.TapedTask`, if one exists. 
-If this function is not called within a TapedTask, return the provided `varinfo`.
 """
-function get_trace_local_varinfo_maybe(varinfo::AbstractVarInfo)
-    trace = try
-        Libtask.get_taped_globals(Any).other
-    catch e
-        e == KeyError(:task_variable) ? nothing : rethrow(e)
-    end
-    has_trace = trace !== nothing
-    return (has_trace ? trace.model.f.varinfo : varinfo)::AbstractVarInfo
+function get_trace_local_varinfo()
+    trace = Libtask.get_taped_globals(Any).other
+    return trace.model.f.varinfo::AbstractVarInfo
 end
 
 """
-    get_trace_local_resampled_maybe(fallback_resampled::Bool)
+    get_trace_local_resampled()
 
-Get the Boolean `resample` value stored in the 'taped globals' of a `Libtask.TapedTask`, if
-one exists. If this function is not called within a TapedTask, return the provided
-`fallback_resampled` value.
+Get the Boolean `resample` value stored in the 'taped globals' of a `Libtask.TapedTask`.
 """
-function get_trace_local_resampled_maybe(fallback_resampled::Bool)
-    trace = try
-        Libtask.get_taped_globals(Any).other
-    catch e
-        e == KeyError(:task_variable) ? nothing : rethrow(e)
-    end
-    has_trace = trace !== nothing
-    return (has_trace ? trace.model.f.resample : fallback_resampled)::Bool
+function get_trace_local_resampled()
+    trace = Libtask.get_taped_globals(Any).other
+    return trace.model.f.resample::Bool
 end
 
 """
-    get_trace_local_rng_maybe(rng::Random.AbstractRNG)
+    get_trace_local_rng()
 
-Get the `Trace` local rng if one exists.
-
-If executed within a `TapedTask`, return the `rng` stored in the "taped globals" of the
-task, otherwise return `vi`.
+Get the RNG stored in the 'taped globals' of a `Libtask.TapedTask`, if one exists.
 """
-function get_trace_local_rng_maybe(rng::Random.AbstractRNG)
-    return try
-        Libtask.get_taped_globals(Any).rng
-    catch e
-        e == KeyError(:task_variable) ? rng : rethrow(e)
-    end
+function get_trace_local_rng()
+    return Libtask.get_taped_globals(Any).rng
 end
 
 """
-    set_trace_local_varinfo_maybe(vi::AbstractVarInfo)
+    set_trace_local_varinfo(vi::AbstractVarInfo)
 
 Set the `Trace` local varinfo if executing within a `Trace`. Return `nothing`.
 
 If executed within a `TapedTask`, set the `varinfo` stored in the "taped globals" of the
 task. Otherwise do nothing.
 """
-function set_trace_local_varinfo_maybe(vi::AbstractVarInfo)
-    # TODO(mhauru) This should be done in a try-catch block, as in the commented out code.
-    # However, Libtask currently can't handle this block.
-    trace = #try
-        Libtask.get_taped_globals(Any).other
-    # catch e
-    #     e == KeyError(:task_variable) ? nothing : rethrow(e)
-    # end
-    if trace !== nothing
-        # trace isa AdvancedPS.Trace (defined in AdvancedPS src/model.jl)
-        # trace.model isa AdvancedPS.LibtaskModel (defined in AdvancedPSLibtaskExt)
-        # trace.model.f isa TracedModel (defined above)
-        trace.model.f.varinfo = vi
-    end
-    return nothing
+function set_trace_local_varinfo(vi::AbstractVarInfo)
+    trace = Libtask.get_taped_globals(Any).other
+    return trace.model.f.varinfo = vi
 end
 
 function DynamicPPL.tilde_assume!!(
-    ctx::ParticleMCMCContext,
-    dist::Distribution,
-    vn::VarName,
-    template::Any,
-    vi::AbstractVarInfo,
+    ::ParticleMCMCContext, dist::Distribution, vn::VarName, template::Any, ::AbstractVarInfo
 )
-    vi = get_trace_local_varinfo_maybe(vi)
-
-    trng = get_trace_local_rng_maybe(ctx.rng)
-    resample = get_trace_local_resampled_maybe(true)
+    vi = get_trace_local_varinfo()
+    trng = get_trace_local_rng()
+    resample = get_trace_local_resampled()
 
     dispatch_ctx = if ~haskey(vi, vn) || resample
         DynamicPPL.InitContext(trng, DynamicPPL.InitFromPrior())
@@ -426,7 +389,7 @@ function DynamicPPL.tilde_assume!!(
     end
     x, vi = DynamicPPL.tilde_assume!!(dispatch_ctx, dist, vn, template, vi)
 
-    set_trace_local_varinfo_maybe(vi)
+    set_trace_local_varinfo(vi)
     return x, vi
 end
 
@@ -437,9 +400,9 @@ function DynamicPPL.tilde_observe!!(
     vn::Union{VarName,Nothing},
     vi::AbstractVarInfo,
 )
-    vi = get_trace_local_varinfo_maybe(vi)
+    vi = get_trace_local_varinfo()
     left, vi = DynamicPPL.tilde_observe!!(DefaultContext(), right, left, vn, vi)
-    set_trace_local_varinfo_maybe(vi)
+    set_trace_local_varinfo(vi)
     return left, vi
 end
 

test/mcmc/mh.jl

@@ -70,9 +70,7 @@ GKernel(variance) = (vnt -> Normal(vnt[@varname(m)], sqrt(variance)))
                 @varname(mu1) => MH(:mu1 => GKernel(1)),
                 @varname(mu2) => MH(:mu2 => GKernel(1)),
             )
-            initial_params = InitFromParams((
-                mu1=1.0, mu2=1.0, z1=0.0, z2=0.0, z3=1.0, z4=1.0
-            ))
+            initial_params = InitFromParams((mu1=1.0, mu2=1.0, z1=0, z2=0, z3=1, z4=1))
             chain = sample(
                 StableRNG(seed),
                 MoGtest_default,
@@ -160,55 +158,6 @@ GKernel(variance) = (vnt -> Normal(vnt[@varname(m)], sqrt(variance)))
             end
         end
     end
-
-    @testset "MH link/invlink" begin
-        vi_base = DynamicPPL.VarInfo(gdemo_default)
-
-        # Don't link when no proposals are given since we're using priors
-        # as proposals.
-        vi = deepcopy(vi_base)
-        spl = MH()
-        vi = Turing.Inference.maybe_link!!(vi, spl, spl.proposals, gdemo_default)
-        @test !DynamicPPL.is_transformed(vi)
-
-        # Link if proposal is `AdvancedHM.RandomWalkProposal`
-        vi = deepcopy(vi_base)
-        d = length(vi_base[:])
-        spl = MH(AdvancedMH.RandomWalkProposal(MvNormal(zeros(d), I)))
-        vi = Turing.Inference.maybe_link!!(vi, spl, spl.proposals, gdemo_default)
-        @test DynamicPPL.is_transformed(vi)
-
-        # Link if ALL proposals are `AdvancedHM.RandomWalkProposal`.
-        vi = deepcopy(vi_base)
-        spl = MH(:s => AdvancedMH.RandomWalkProposal(Normal()))
-        vi = Turing.Inference.maybe_link!!(vi, spl, spl.proposals, gdemo_default)
-        @test DynamicPPL.is_transformed(vi)
-
-        # Don't link if at least one proposal is NOT `RandomWalkProposal`.
-        # TODO: make it so that only those that are using `RandomWalkProposal`
-        # are linked! I.e. resolve https://github.com/TuringLang/Turing.jl/issues/1583.
-        # https://github.com/TuringLang/Turing.jl/pull/1582#issuecomment-817148192
-        vi = deepcopy(vi_base)
-        spl = MH(
-            :m => AdvancedMH.StaticProposal(Normal()),
-            :s => AdvancedMH.RandomWalkProposal(Normal()),
-        )
-        vi = Turing.Inference.maybe_link!!(vi, spl, spl.proposals, gdemo_default)
-        @test !DynamicPPL.is_transformed(vi)
-    end
-
-    @testset "`filldist` proposal (issue #2180)" begin
-        @model demo_filldist_issue2180() = x ~ MvNormal(zeros(3), I)
-        chain = sample(
-            StableRNG(seed),
-            demo_filldist_issue2180(),
-            MH(AdvancedMH.RandomWalkProposal(filldist(Normal(), 3))),
-            10_000,
-        )
-        check_numerical(
-            chain, [Symbol("x[1]"), Symbol("x[2]"), Symbol("x[3]")], [0, 0, 0]; atol=0.2
-        )
-    end
 end
 
 end