Use typed varinfo for Prior

Author: penelopeysm

HISTORY.md

@@ -1,6 +1,11 @@
 # 0.40.0
 
-[...]
+TODO
+
+  - DynamicPPL 0.37 stuff
+
+  - pMCMC and Gibbs?
+  - Prior is faster
 
 # 0.39.9
 

src/mcmc/prior.jl

@@ -8,8 +8,7 @@ struct Prior <: InferenceAlgorithm end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
-    sampler::DynamicPPL.Sampler{<:Prior},
-    state=nothing;
+    sampler::DynamicPPL.Sampler{<:Prior};
     kwargs...,
 )
     # TODO(DPPL0.38/penelopeysm): replace with init!!
@@ -26,5 +25,27 @@ function AbstractMCMC.step(
         ),
     )
     _, vi = DynamicPPL.evaluate!!(sampling_model, vi)
-    return Transition(model, vi, nothing; reevaluate=false), nothing
+    vi = DynamicPPL.typed_varinfo(vi)
+    return Transition(model, vi, nothing; reevaluate=false), vi
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::DynamicPPL.Sampler{<:Prior},
+    vi::AbstractVarInfo;
+    kwargs...,
+)
+    # TODO(DPPL0.38/penelopeysm): replace the entire thing with init!!
+    # `vi` is a VarInfo from the previous step so already has all the
+    # right accumulators and stuff. The only thing we need to change is
+    # to make sure that the old values are overwritten...
+    for vn in keys(vi)
+        DynamicPPL.set_flag!(vi, vn, "del")
+    end
+    sampling_model = DynamicPPL.contextualize(
+        model, DynamicPPL.SamplingContext(rng, DynamicPPL.SampleFromPrior(), model.context)
+    )
+    _, vi = DynamicPPL.evaluate!!(sampling_model, vi)
+    return Transition(model, vi, nothing; reevaluate=false), vi
 end