Use templating info in MH

Author: penelopeysm

src/mcmc/mh.jl

@@ -64,9 +64,10 @@ spl = MH(
 
 You can also use a callable to define a proposal that is conditional on the current values.
 The callable must accept a single argument, which is a `DynamicPPL.VarNamedTuple` that holds
-all the current values of the parameters. You can obtain the value of a specific parameter
-by indexing into this `VarNamedTuple` using a `VarName` (note that symbol indexing is not
-supported). The callable must then return a `Distribution` from which to draw the proposal.
+all the values of the parameters from the previous step. You can obtain the value of a
+specific parameter by indexing into this `VarNamedTuple` using a `VarName` (note that symbol
+indexing is not supported). The callable must then return a `Distribution` from which to
+draw the proposal.
 
 !!! note
     In general, there is no way for Turing to reliably detect whether a proposal is meant to
@@ -125,9 +126,14 @@ function DynamicPPL.init(
     return DynamicPPL.UntransformedValue(rand(rng, dist))
 end
 
-function MH(pair1, pairs...)
+const SymOrVNPair = Pair{<:Union{Symbol,VarName},<:Any}
+
+function MH(pair1::SymOrVNPair, pairs::Vararg{SymOrVNPair})
     vn_proposal_pairs = (pair1, pairs...)
     return MH(
+        # It is assumed that `vnt` is a VarNamedTuple that has all the variables' values
+        # already set. NOTE: It doesn't store `AbstractTransformedValue`s, but the actual
+        # raw values.
         vnt -> begin
             proposals = DynamicPPL.VarNamedTuple()
             for pair in vn_proposal_pairs
@@ -143,16 +149,15 @@ function MH(pair1, pairs...)
                     )
                 end
                 # Check whether the proposal is a Distribution.
-                # TODO(penelopeysm): We don't have templating information here. :(
-                # We *could* steal the templating info from the `vnt` argument. It would
-                # be quite type unstable since you would be mixing and matching
-                # distributions with values. Need to check if it's worth it.
-                if proposal isa Distribution
-                    proposals = BangBang.setindex!!(proposals, proposal, vn)
+                proposal_dist = if proposal isa Distribution
+                    proposal
                 else
-                    # Assume it's a callable.
-                    proposals = BangBang.setindex!!(proposals, proposal(vnt), vn)
+                    # It's a callable that takes `vnt` and returns a distribution.
+                    proposal(vnt)
                 end
+                proposals = DynamicPPL.templated_setindex!!(
+                    proposals, proposal_dist, vn, vnt.data[AbstractPPL.getsym(vn)]
+                )
             end
             return InitFromProposals(proposals, Dict{VarName,Distribution}())
         end,