Improve maybe_link! for MH sampler

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.15.13"
+version = "0.15.14"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/inference/mh.jl

@@ -334,9 +334,26 @@ end
 end
 
 # Utility functions to link
-maybe_link!(varinfo, sampler, proposal) = nothing
-function maybe_link!(varinfo, sampler, proposal::AdvancedMH.RandomWalkProposal)
-    link!(varinfo, sampler)
+should_link(varinfo, sampler, proposal) = false
+function should_link(varinfo, sampler, proposal::NamedTuple{(), Tuple{}})
+    # If it's an empty `NamedTuple`, we're using the priors as proposals
+    # in which case we shouldn't link.
+    return false
+end
+function should_link(varinfo, sampler, proposal::AdvancedMH.RandomWalkProposal)
+    return true
+end
+function should_link(
+    varinfo,
+    sampler,
+    proposal::NamedTuple{names, vals}
+) where {names, vals<:NTuple{<:Any, <:AdvancedMH.RandomWalkProposal}}
+    return true
+end
+
+function maybe_link!(varinfo, sampler, proposal)
+    should_link(varinfo, sampler, proposal) && link!(varinfo, sampler)
+    return varinfo
 end
 
 # Make a proposal if we don't have a covariance proposal matrix (the default).

test/inference/mh.jl

@@ -130,13 +130,17 @@
         # with small-valued VC matrix to check if we only see very small steps
         vc_μ = convert(Array, 1e-4*I(2))
         vc_σ = convert(Array, 1e-4*I(2))
-
-        chn = sample(mod,
-            Gibbs(
-                MH((:μ, vc_μ)),
-                MH((:σ, vc_σ)),
-                ), 3_000 # draws
-            )
+
+        alg = Gibbs(
+            MH((:μ, vc_μ)),
+            MH((:σ, vc_σ)),
+        )
+
+        chn = sample(
+            mod,
+            alg,
+            3_000 # draws
+        )
 
 
         chn2 = sample(mod, MH(), 3_000)
@@ -168,4 +172,44 @@
             @test mean(chain, "T[$j][$i]") ≈ 0.2 atol=0.01
         end
     end
+
+    @turing_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)
+        alg = MH()
+        spl = DynamicPPL.Sampler(alg)
+        Turing.Inference.maybe_link!(vi, spl, alg.proposals)
+        @test !DynamicPPL.islinked(vi, spl)
+
+        # Link if proposal is `AdvancedHM.RandomWalkProposal`
+        vi = deepcopy(vi_base)
+        d = length(vi_base[DynamicPPL.SampleFromPrior()])
+        alg = MH(AdvancedMH.RandomWalkProposal(MvNormal(d, 1.0)))
+        spl = DynamicPPL.Sampler(alg)
+        Turing.Inference.maybe_link!(vi, spl, alg.proposals)
+        @test DynamicPPL.islinked(vi, spl)
+
+        # Link if ALL proposals are `AdvancedHM.RandomWalkProposal`.
+        vi = deepcopy(vi_base)
+        alg = MH(:s => AdvancedMH.RandomWalkProposal(Normal()))
+        spl = DynamicPPL.Sampler(alg)
+        Turing.Inference.maybe_link!(vi, spl, alg.proposals)
+        @test DynamicPPL.islinked(vi, spl)
+
+        # 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)
+        alg = MH(
+            :m => AdvancedMH.StaticProposal(Normal()),
+            :s => AdvancedMH.RandomWalkProposal(Normal())
+        )
+        spl = DynamicPPL.Sampler(alg)
+        Turing.Inference.maybe_link!(vi, spl, alg.proposals)
+        @test !DynamicPPL.islinked(vi, spl)
+    end
 end