Improve performance of Prior()

Author: penelopeysm

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...,
 )
     vi = last(
@@ -19,7 +18,29 @@ function AbstractMCMC.step(
             SamplingContext(rng, DynamicPPL.SampleFromPrior(), DynamicPPL.PriorContext()),
         ),
     )
-    return vi, nothing
+    vi = DynamicPPL.typed_varinfo(vi)
+    return vi, vi
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::DynamicPPL.Model,
+    sampler::DynamicPPL.Sampler{<:Prior},
+    vi::AbstractVarInfo;
+    kwargs...,
+)
+    # TODO(DPPL0.38/penelopeysm): Use InitContext.
+    for vn in keys(vi)
+        DynamicPPL.set_flag!(vi, vn, "del")
+    end
+    vi = last(
+        DynamicPPL.evaluate!!(
+            model,
+            vi,
+            SamplingContext(rng, DynamicPPL.SampleFromPrior(), DynamicPPL.PriorContext()),
+        ),
+    )
+    return vi, vi
 end
 
 DynamicPPL.default_chain_type(sampler::Prior) = MCMCChains.Chains