diff --git a/Project.toml b/Project.toml
index adfc433092..4657baa0f1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -45,7 +45,7 @@ DataStructures = "0.18"
 Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8"
-DynamicPPL = "0.16"
+DynamicPPL = "0.17.2"
 EllipticalSliceSampling = "0.4"
 ForwardDiff = "0.10.3"
 Libtask = "0.4, 0.5.3"
diff --git a/docs/src/for-developers/how_turing_implements_abstractmcmc.md b/docs/src/for-developers/how_turing_implements_abstractmcmc.md
index f29d0ff782..880c954dfe 100644
--- a/docs/src/for-developers/how_turing_implements_abstractmcmc.md
+++ b/docs/src/for-developers/how_turing_implements_abstractmcmc.md
@@ -197,7 +197,7 @@ Consider an instance `m` of `Model` and a sampler `spl`, with associated `VarInf
   * recall that the code for `m.f(vi, ...)` is automatically generated by compilation of the `@model` macro
   * for every tilde statement in the `@model` declaration, this code contains a call to `assume(vi, ...)` if the variable on the LHS of the tilde is a **model parameter to infer**, and `observe(vi, ...)` if the variable on the LHS of the tilde is an **observation**
   * in the file corresponding to your sampling method (ie in `Turing.jl/src/inference/<your_method>.jl`), you have **overloaded** `assume` and `observe`, so that they can modify `vi` to include the information and samples that you care about!
-  * at a minimum, `assume` and `observe` return the log density `lp` of the sample or observation. the model evaluation function then immediately calls `acclogp!(vi, lp)`, which adds `lp` to the value of the log joint density stored in `vi`.
+  * at a minimum, `assume` and `observe` return the log density `lp` of the sample or observation. the model evaluation function then immediately calls `acclogp!!(vi, lp)`, which adds `lp` to the value of the log joint density stored in `vi`.
 
 Here's what `assume` looks like for Importance Sampling:
 
@@ -226,10 +226,10 @@ It simply returns the density (in the discrete case, the probability) of the obs
 We focus on the AbstractMCMC functions that are overriden in `is.jl` and executed inside `mcmcsample`: `step!`, which is called `n_samples` times, and `sample_end!`, which is executed once after those `n_samples` iterations.
 
 * During the \$\$i\$\$-th iteration, `step!` does 3 things:
-  * `empty!(spl.state.vi)`: remove information about the previous sample from the sampler's `VarInfo`
+  * `empty!!(spl.state.vi)`: remove information about the previous sample from the sampler's `VarInfo`
   * `model(rng, spl.state.vi, spl)`: call the model evaluation function
     * calls to `assume` add the samples from the prior \$\$s\_i\$\$ and \$\$m\_i\$\$ to `spl.state.vi`
-    * calls to both `assume` or `observe` are followed by the line `acclogp!(vi, lp)`, where `lp` is an output of `assume` and `observe`
+    * calls to both `assume` or `observe` are followed by the line `acclogp!!(vi, lp)`, where `lp` is an output of `assume` and `observe`
     * `lp` is set to 0 after `assume`, and to the value of the density at the observation after `observe`
     * when all the tilde statements have been covered, `spl.state.vi.logp[]` is the sum of the `lp`, ie the likelihood \$\$\log p(x, y \mid s\_i, m\_i) = \log p(x \mid s\_i, m\_i) + \log p(y \mid s\_i, m\_i)\$\$ of the observations given the latent variable samples \$\$s\_i\$\$ and \$\$m\_i\$\$.
   * `return Transition(spl)`: build a transition from the sampler, and return that transition
diff --git a/src/contrib/inference/dynamichmc.jl b/src/contrib/inference/dynamichmc.jl
index 9a1cadd3a6..151fee42fd 100644
--- a/src/contrib/inference/dynamichmc.jl
+++ b/src/contrib/inference/dynamichmc.jl
@@ -97,7 +97,7 @@ function DynamicPPL.initialstep(
 
     # Update the variables.
     vi[spl] = Q.q
-    DynamicPPL.setlogp!(vi, Q.ℓq)
+    DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create first sample and state.
     sample = Transition(vi)
@@ -127,7 +127,7 @@ function AbstractMCMC.step(
 
     # Update the variables.
     vi[spl] = Q.q
-    DynamicPPL.setlogp!(vi, Q.ℓq)
+    DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create next sample and state.
     sample = Transition(vi)
diff --git a/src/core/ad.jl b/src/core/ad.jl
index 0fca502107..d915edb4e4 100644
--- a/src/core/ad.jl
+++ b/src/core/ad.jl
@@ -108,9 +108,11 @@ function gradient_logp(
     logp_old = getlogp(vi)
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        model(new_vi, sampler, ctx)
+        new_vi = last(DynamicPPL.evaluate!!(model, new_vi, sampler, ctx))
         logp = getlogp(new_vi)
-        setlogp!(vi, ForwardDiff.value(logp))
+        # Don't need to capture the resulting `vi` since this is only
+        # needed if `vi` is mutable.
+        setlogp!!(vi, ForwardDiff.value(logp))
         return logp
     end
 
@@ -120,7 +122,7 @@ function gradient_logp(
     config = ForwardDiff.GradientConfig(f, θ, chunk)
     ∂l∂θ = ForwardDiff.gradient!(similar(θ), f, θ, config)
     l = getlogp(vi)
-    setlogp!(vi, logp_old)
+    setlogp!!(vi, logp_old)
 
     return l, ∂l∂θ
 end
@@ -137,7 +139,7 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        model(new_vi, sampler, ctx)
+        new_vi = last(DynamicPPL.evaluate!!(model, new_vi, sampler, ctx))
         return getlogp(new_vi)
     end
 
@@ -162,7 +164,7 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        model(new_vi, sampler, context)
+        new_vi = last(DynamicPPL.evaluate!!(model, new_vi, sampler, context))
         return getlogp(new_vi)
     end
 
diff --git a/src/core/container.jl b/src/core/container.jl
index d1ba1d6da7..1c135e9f20 100644
--- a/src/core/container.jl
+++ b/src/core/container.jl
@@ -43,7 +43,7 @@ function AdvancedPS.reset_model(f::TracedModel)
 end
 
 function AdvancedPS.reset_logprob!(f::TracedModel)
-    DynamicPPL.resetlogp!(f.varinfo)
+    DynamicPPL.resetlogp!!(f.varinfo)
     return
 end
 
diff --git a/src/inference/AdvancedSMC.jl b/src/inference/AdvancedSMC.jl
index e7cf135862..236084d89f 100644
--- a/src/inference/AdvancedSMC.jl
+++ b/src/inference/AdvancedSMC.jl
@@ -110,8 +110,8 @@ function DynamicPPL.initialstep(
     # Reset the VarInfo.
     reset_num_produce!(vi)
     set_retained_vns_del_by_spl!(vi, spl)
-    resetlogp!(vi)
-    empty!(vi)
+    resetlogp!!(vi)
+    empty!!(vi)
 
     # Create a new set of particles.
     particles = AdvancedPS.ParticleContainer(
@@ -249,7 +249,7 @@ function DynamicPPL.initialstep(
     # Reset the VarInfo before new sweep
     reset_num_produce!(vi)
     set_retained_vns_del_by_spl!(vi, spl)
-    resetlogp!(vi)
+    resetlogp!!(vi)
 
     # Create a new set of particles
     num_particles = spl.alg.nparticles
@@ -281,7 +281,7 @@ function AbstractMCMC.step(
 )
     # Reset the VarInfo before new sweep.
     reset_num_produce!(vi)
-    resetlogp!(vi)
+    resetlogp!!(vi)
 
     # Create reference particle for which the samples will be retained.
     reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi))
@@ -315,6 +315,8 @@ function AbstractMCMC.step(
     return transition, _vi
 end
 
+DynamicPPL.use_threadsafe_eval(::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, ::AbstractVarInfo) = false
+
 function DynamicPPL.assume(
     rng,
     spl::Sampler{<:Union{PG,SMC}},
@@ -326,7 +328,7 @@ function DynamicPPL.assume(
     if inspace(vn, spl)
         if ~haskey(vi, vn)
             r = rand(rng, dist)
-            push!(vi, vn, r, dist, spl)
+            push!!(vi, vn, r, dist, spl)
         elseif is_flagged(vi, vn, "del")
             unset_flag!(vi, vn, "del")
             r = rand(rng, dist)
@@ -342,17 +344,17 @@ function DynamicPPL.assume(
             r = vi[vn]
         else
             r = rand(rng, dist)
-            push!(vi, vn, r, dist, DynamicPPL.Selector(:invalid))
+            push!!(vi, vn, r, dist, DynamicPPL.Selector(:invalid))
         end
         lp = logpdf_with_trans(dist, r, istrans(vi, vn))
-        acclogp!(vi, lp)
+        acclogp!!(vi, lp)
     end
-    return r, 0
+    return r, 0, vi
 end
 
 function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
     produce(logpdf(dist, value))
-    return 0
+    return 0, vi
 end
 
 # Convenient constructor
diff --git a/src/inference/Inference.jl b/src/inference/Inference.jl
index 7ec21dcec4..49efcc5f4b 100644
--- a/src/inference/Inference.jl
+++ b/src/inference/Inference.jl
@@ -3,7 +3,10 @@ module Inference
 using ..Core
 using ..Utilities
 using DynamicPPL: Metadata, VarInfo, TypedVarInfo, 
-    islinked, invlink!, getlogp, tonamedtuple, VarName, getsym, vectorize, 
+    islinked, invlink!, link!,
+    setindex!!, push!!,
+    setlogp!!, getlogp,
+    tonamedtuple, VarName, getsym, vectorize,
     settrans!, _getvns, getdist,
     Model, Sampler, SampleFromPrior, SampleFromUniform,
     DefaultContext, PriorContext,
diff --git a/src/inference/emcee.jl b/src/inference/emcee.jl
index b3f1579a30..83415c191f 100644
--- a/src/inference/emcee.jl
+++ b/src/inference/emcee.jl
@@ -38,22 +38,21 @@ function AbstractMCMC.step(
 
     # Update the parameters if provided.
     if haskey(kwargs, :init_params)
-        for vi in vis
-            initialize_parameters!(vi, kwargs[:init_params], spl)
+        vis = map(vis) do vi
+            vi = initialize_parameters!!(vi, kwargs[:init_params], spl)
 
             # Update log joint probability.
-            model(rng, vi, SampleFromPrior())
+            last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromPrior()))
         end
     end
 
     # Compute initial transition and states.
-    transition = map(vis) do vi
-        Transition(vi)
-    end
+    transition = map(Transition, vis)
+
+    # TODO: Make compatible with immutable `AbstractVarInfo`.
     state = EmceeState(
         vis[1],
         map(vis) do vi
-            # Transform to unconstrained space.
             DynamicPPL.link!(vi, spl)
             AMH.Transition(vi[spl], getlogp(vi))
         end
@@ -78,7 +77,7 @@ function AbstractMCMC.step(
 
     # Compute the next transition and state.
     transition = map(states) do _state
-        vi[spl] = _state.params
+        vi = setindex!!(vi, _state.params, spl)
         DynamicPPL.invlink!(vi, spl)
         t = Transition(tonamedtuple(vi), _state.lp)
         DynamicPPL.link!(vi, spl)
diff --git a/src/inference/ess.jl b/src/inference/ess.jl
index 13c70bda17..f724f3c536 100644
--- a/src/inference/ess.jl
+++ b/src/inference/ess.jl
@@ -71,8 +71,8 @@ function AbstractMCMC.step(
     )
 
     # update sample and log-likelihood
-    vi[spl] = sample
-    setlogp!(vi, state.loglikelihood)
+    vi = setindex!!(vi, sample, spl)
+    vi = setlogp!!(vi, state.loglikelihood)
 
     return Transition(vi), vi
 end
@@ -111,6 +111,7 @@ EllipticalSliceSampling.isgaussian(::Type{<:ESSPrior}) = true
 function Base.rand(rng::Random.AbstractRNG, p::ESSPrior)
     sampler = p.sampler
     varinfo = p.varinfo
+    # TODO: Surely there's a better way of doing this now that we have `SamplingContext`?
     vns = _getvns(varinfo, sampler)
     for vn in Iterators.flatten(values(vns))
         set_flag!(varinfo, vn, "del")
@@ -131,17 +132,16 @@ end
 
 function (ℓ::ESSLogLikelihood)(f)
     sampler = ℓ.sampler
-    varinfo = ℓ.varinfo
-    varinfo[sampler] = f
-    ℓ.model(varinfo, sampler)
+    varinfo = setindex!!(ℓ.varinfo, f, sampler)
+    varinfo = last(DynamicPPL.evaluate!!(ℓ.model, varinfo, sampler))
     return getlogp(varinfo)
 end
 
 function DynamicPPL.tilde_assume(rng::Random.AbstractRNG, ctx::DefaultContext, sampler::Sampler{<:ESS}, right, vn, vi)
-    if inspace(vn, sampler)
-        return DynamicPPL.tilde_assume(rng, LikelihoodContext(), SampleFromPrior(), right, vn, vi)
+    return if inspace(vn, sampler)
+        DynamicPPL.tilde_assume(rng, LikelihoodContext(), SampleFromPrior(), right, vn, vi)
     else
-        return DynamicPPL.tilde_assume(rng, ctx, SampleFromPrior(), right, vn, vi)
+        DynamicPPL.tilde_assume(rng, ctx, SampleFromPrior(), right, vn, vi)
     end
 end
 
@@ -151,10 +151,10 @@ end
 
 function DynamicPPL.dot_tilde_assume(rng::Random.AbstractRNG, ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vns, vi)
     # TODO: Or should we do `all(Base.Fix2(inspace, sampler), vns)`?
-    if inspace(first(vns), sampler)
-        return DynamicPPL.dot_tilde_assume(rng, LikelihoodContext(), SampleFromPrior(), right, left, vns, vi)
+    return if inspace(first(vns), sampler)
+        DynamicPPL.dot_tilde_assume(rng, LikelihoodContext(), SampleFromPrior(), right, left, vns, vi)
     else
-        return DynamicPPL.dot_tilde_assume(rng, ctx, SampleFromPrior(), right, left, vns, vi)
+        DynamicPPL.dot_tilde_assume(rng, ctx, SampleFromPrior(), right, left, vns, vi)
     end
 end
 
diff --git a/src/inference/gibbs.jl b/src/inference/gibbs.jl
index 1d4fa6f64e..58fe814f49 100644
--- a/src/inference/gibbs.jl
+++ b/src/inference/gibbs.jl
@@ -163,6 +163,9 @@ function DynamicPPL.initialstep(
     vi::AbstractVarInfo;
     kwargs...
 )
+    # TODO: Technically this only works for `VarInfo` or `ThreadSafeVarInfo{<:VarInfo}`.
+    # Should we enforce this?
+
     # Create tuple of samplers
     algs = spl.alg.algs
     i = 0
@@ -230,7 +233,7 @@ function AbstractMCMC.step(
     states = map(samplers, state.states) do _sampler, _state
         # Recompute `vi.logp` if needed.
         if _sampler.selector.rerun
-            model(rng, vi, _sampler)
+            vi = last(DynamicPPL.evaluate!!(model, rng, vi, _sampler))
         end
 
         # Update state of current sampler with updated `VarInfo` object.
diff --git a/src/inference/gibbs_conditional.jl b/src/inference/gibbs_conditional.jl
index f923d5f866..8dd44904b0 100644
--- a/src/inference/gibbs_conditional.jl
+++ b/src/inference/gibbs_conditional.jl
@@ -79,8 +79,10 @@ function AbstractMCMC.step(
     condvals = conditioned(tonamedtuple(vi))
     conddist = spl.alg.conditional(condvals)
     updated = rand(rng, conddist)
-    vi[spl] = [updated;]  # setindex allows only vectors in this case...
-    model(rng, vi, SampleFromPrior()) # update log joint probability
+    # Setindex allows only vectors in this case.
+    vi = setindex!!(vi, [updated;], spl)
+    # Update log joint probability.
+    vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromPrior()))
 
     return nothing, vi
 end
diff --git a/src/inference/hmc.jl b/src/inference/hmc.jl
index fee124525d..883ae960e3 100644
--- a/src/inference/hmc.jl
+++ b/src/inference/hmc.jl
@@ -151,7 +151,7 @@ function DynamicPPL.initialstep(
 )
     # Transform the samples to unconstrained space and compute the joint log probability.
     link!(vi, spl)
-    model(rng, vi, spl)
+    vi = last(DynamicPPL.evaluate!!(model, rng, vi, spl))
 
     # Extract parameters.
     theta = vi[spl]
@@ -170,7 +170,7 @@ function DynamicPPL.initialstep(
     # and its gradient are finite.
     if init_params === nothing
         while !isfinite(z)
-            model(rng, vi, SampleFromUniform())
+            vi = last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromUniform()))
             link!(vi, spl)
             theta = vi[spl]
 
@@ -207,11 +207,11 @@ function DynamicPPL.initialstep(
 
     # Update `vi` based on acceptance
     if t.stat.is_accept
-        vi[spl] = t.z.θ
-        setlogp!(vi, t.stat.log_density)
+        vi = setindex!!(vi, t.z.θ, spl)
+        vi = setlogp!!(vi, t.stat.log_density)
     else
-        vi[spl] = theta
-        setlogp!(vi, log_density_old)
+        vi = setindex!!(vi, theta, spl)
+        vi = setlogp!!(vi, log_density_old)
     end
 
     transition = HMCTransition(vi, t)
@@ -249,8 +249,8 @@ function AbstractMCMC.step(
     # Update variables
     vi = state.vi
     if t.stat.is_accept
-        vi[spl] = t.z.θ
-        setlogp!(vi, t.stat.log_density)
+        vi = setindex!!(vi, t.z.θ, spl)
+        vi = setlogp!!(vi, t.stat.log_density)
     end
 
     # Compute next transition and state.
@@ -441,14 +441,17 @@ end
 Generate a function that takes `θ` and returns logpdf at `θ` for the model specified by
 `(vi, spl, model)`.
 """
-function gen_logπ(vi, spl::AbstractSampler, model)
+function gen_logπ(vi_base, spl::AbstractSampler, model)
     function logπ(x)::Float64
+        vi = vi_base
         x_old, lj_old = vi[spl], getlogp(vi)
-        vi[spl] = x
-        model(vi, spl)
+        vi = setindex!!(vi, x, spl)
+        vi = last(DynamicPPL.evaluate!!(model, vi, spl))
         lj = getlogp(vi)
-        vi[spl] = x_old
-        setlogp!(vi, lj_old)
+        # Don't really need to capture these will only be
+        # necessary if `vi` is indeed mutable.
+        setindex!!(vi, x_old, spl)
+        setlogp!!(vi, lj_old)
         return lj
     end
     return logπ
@@ -478,10 +481,7 @@ function DynamicPPL.assume(
     vi,
 )
     DynamicPPL.updategid!(vi, vn, spl)
-    r = vi[vn]
-    # acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn)))
-    # r
-    return r, logpdf_with_trans(dist, r, istrans(vi, vn))
+    return DynamicPPL.assume(dist, vn, vi)
 end
 
 function DynamicPPL.dot_assume(
@@ -492,11 +492,8 @@ function DynamicPPL.dot_assume(
     var::AbstractMatrix,
     vi,
 )
-    @assert length(dist) == size(var, 1)
     DynamicPPL.updategid!.(Ref(vi), vns, Ref(spl))
-    r = vi[vns]
-    var .= r
-    return var, sum(logpdf_with_trans(dist, r, istrans(vi, vns[1])))
+    return DynamicPPL.dot_assume(dist, var, vns, vi)
 end
 function DynamicPPL.dot_assume(
     rng,
@@ -506,10 +503,8 @@ function DynamicPPL.dot_assume(
     var::AbstractArray,
     vi,
 )
-    DynamicPPL.updategid!.(Ref(vi), vns, Ref(spl))
-    r = reshape(vi[vec(vns)], size(var))
-    var .= r
-    return var, sum(logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
+    DynamicPPL.updategid!(Ref(vi), vns, Ref(spl))
+    return DynamicPPL.dot_assume(dists, var, vns, vi)
 end
 
 function DynamicPPL.observe(
@@ -518,7 +513,7 @@ function DynamicPPL.observe(
     value,
     vi,
 )
-    return DynamicPPL.observe(SampleFromPrior(), d, value, vi)
+    return DynamicPPL.observe(d, value, vi)
 end
 
 function DynamicPPL.dot_observe(
@@ -527,7 +522,7 @@ function DynamicPPL.dot_observe(
     value::AbstractArray,
     vi,
 )
-    return DynamicPPL.dot_observe(SampleFromPrior(), ds, value, vi)
+    return DynamicPPL.dot_observe(ds, value, vi)
 end
 
 ####
@@ -566,7 +561,9 @@ function HMCState(
     kwargs...
 )
     # Link everything if needed.
-    !islinked(vi, spl) && link!(vi, spl)
+    if !islinked(vi, spl)
+        link!(vi, spl)
+    end
 
     # Get the initial log pdf and gradient functions.
     ∂logπ∂θ = gen_∂logπ∂θ(vi, spl, model)
diff --git a/src/inference/is.jl b/src/inference/is.jl
index eed48ffe74..f79e29d344 100644
--- a/src/inference/is.jl
+++ b/src/inference/is.jl
@@ -64,11 +64,11 @@ function DynamicPPL.assume(rng, spl::Sampler{<:IS}, dist::Distribution, vn::VarN
         r = vi[vn]
     else
         r = rand(rng, dist)
-        push!(vi, vn, r, dist, spl)
+        vi = push!!(vi, vn, r, dist, spl)
     end
-    return r, 0
+    return r, 0, vi
 end
 
 function DynamicPPL.observe(spl::Sampler{<:IS}, dist::Distribution, value, vi)
-    return logpdf(dist, value)
+    return logpdf(dist, value), vi
 end
diff --git a/src/inference/mh.jl b/src/inference/mh.jl
index 8fea1cf8bc..62748cdf49 100644
--- a/src/inference/mh.jl
+++ b/src/inference/mh.jl
@@ -254,10 +254,12 @@ function (f::MHLogDensityFunction)(x)
 
     x_old, lj_old = vi[sampler], getlogp(vi)
     set_namedtuple!(vi, x)
-    f.model(vi)
-    lj = getlogp(vi)
-    vi[sampler] = x_old
-    setlogp!(vi, lj_old)
+    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, DynamicPPL.DefaultContext()))
+    lj = getlogp(vi_new)
+
+    # Reset old `vi`.
+    setindex!!(vi, x_old, sampler)
+    setlogp!!(vi, lj_old)
 
     return lj
 end
@@ -352,8 +354,10 @@ function should_link(
 end
 
 function maybe_link!(varinfo, sampler, proposal)
-    should_link(varinfo, sampler, proposal) && link!(varinfo, sampler)
-    return varinfo
+    if should_link(varinfo, sampler, proposal)
+        link!(varinfo, sampler)
+    end
+    return nothing
 end
 
 # Make a proposal if we don't have a covariance proposal matrix (the default).
@@ -375,11 +379,12 @@ function propose!(
     densitymodel = AMH.DensityModel(MHLogDensityFunction(model, spl, vi))
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
 
+    # TODO: Make this compatible with immmutable `VarInfo`.
     # Update the values in the VarInfo.
     set_namedtuple!(vi, trans.params)
-    setlogp!(vi, trans.lp)
+    setlogp!!(vi, trans.lp)
 
-    return
+    return vi
 end
 
 # Make a proposal if we DO have a covariance proposal matrix.
@@ -402,11 +407,12 @@ function propose!(
     densitymodel = AMH.DensityModel(gen_logπ(vi, spl, model))
     trans, _ = AbstractMCMC.step(rng, densitymodel, mh_sampler, prev_trans)
 
+    # TODO: Make this compatible with immmutable `VarInfo`.
     # Update the values in the VarInfo.
-    vi[spl] = trans.params
-    setlogp!(vi, trans.lp)
+    setindex!!(vi, trans.params, spl)
+    setlogp!!(vi, trans.lp)
 
-    return
+    return vi
 end
 
 function DynamicPPL.initialstep(
@@ -450,7 +456,7 @@ function DynamicPPL.assume(
 )
     DynamicPPL.updategid!(vi, vn, spl)
     r = vi[vn]
-    return r, logpdf_with_trans(dist, r, istrans(vi, vn))
+    return r, logpdf_with_trans(dist, r, istrans(vi, vn)), vi
 end
 
 function DynamicPPL.dot_assume(
@@ -467,7 +473,7 @@ function DynamicPPL.dot_assume(
     DynamicPPL.updategid!.(Ref(vi), vns, Ref(spl))
     r = vi[vns]
     var .= r
-    return var, sum(logpdf_with_trans(dist, r, istrans(vi, vns[1])))
+    return var, sum(logpdf_with_trans(dist, r, istrans(vi, vns[1]))), vi
 end
 function DynamicPPL.dot_assume(
     rng,
@@ -482,7 +488,7 @@ function DynamicPPL.dot_assume(
     DynamicPPL.updategid!.(Ref(vi), vns, Ref(spl))
     r = reshape(vi[vec(vns)], size(var))
     var .= r
-    return var, sum(logpdf_with_trans.(dists, r, istrans(vi, vns[1])))
+    return var, sum(logpdf_with_trans.(dists, r, istrans(vi, vns[1]))), vi
 end
 
 function DynamicPPL.observe(
diff --git a/src/modes/ModeEstimation.jl b/src/modes/ModeEstimation.jl
index 50d32121db..12fed783a8 100644
--- a/src/modes/ModeEstimation.jl
+++ b/src/modes/ModeEstimation.jl
@@ -48,12 +48,12 @@ end
 
 function DynamicPPL.tilde_assume(ctx::OptimizationContext{<:LikelihoodContext}, spl, dist, vn, vi)
     r = vi[vn]
-    return r, 0
+    return r, 0, vi
 end
 
 function DynamicPPL.tilde_assume(ctx::OptimizationContext, spl, dist, vn, vi)
     r = vi[vn]
-    return r, Distributions.logpdf(dist, r)
+    return r, Distributions.logpdf(dist, r), vi
 end
 
 # dot assume
@@ -65,14 +65,14 @@ function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext{<:LikelihoodContex
     # Values should be set and we're using `SampleFromPrior`, hence the `rng` argument shouldn't
     # affect anything.
     r = DynamicPPL.get_and_set_val!(Random.GLOBAL_RNG, vi, vns, right, sampler)
-    return r, 0
+    return r, 0, vi
 end
 
 function DynamicPPL.dot_tilde_assume(ctx::OptimizationContext, sampler::SampleFromPrior, right, left, vns, vi)
     # Values should be set and we're using `SampleFromPrior`, hence the `rng` argument shouldn't
     # affect anything.
     r = DynamicPPL.get_and_set_val!(Random.GLOBAL_RNG, vi, vns, right, sampler)
-    return r, loglikelihood(right, r)
+    return r, loglikelihood(right, r), vi
 end
 
 """
@@ -158,26 +158,26 @@ end
 #################################################
 
 function transform!(f::OptimLogDensity)
-  spl = DynamicPPL.SampleFromPrior()
+    spl = DynamicPPL.SampleFromPrior()
 
-  ## Check link status of vi in OptimLogDensity
-  linked = DynamicPPL.islinked(f.vi, spl) 
+    ## Check link status of vi in OptimLogDensity
+    linked = DynamicPPL.islinked(f.vi, spl) 
 
-  ## transform into constrained or unconstrained space depending on current state of vi
-  if !linked
-    DynamicPPL.link!(f.vi, spl)
-  else
-    DynamicPPL.invlink!(f.vi, spl)
-  end
+    ## transform into constrained or unconstrained space depending on current state of vi
+    if !linked
+        DynamicPPL.link!(f.vi, spl)
+    else
+        DynamicPPL.invlink!(f.vi, spl)
+    end
 
-  return nothing
+    return nothing
 end
 
 function transform!(p::AbstractArray, vi::DynamicPPL.VarInfo, ::constrained_space{true})
     spl = DynamicPPL.SampleFromPrior()
 
     linked = DynamicPPL.islinked(vi, spl)
-  
+    
     # !linked && DynamicPPL.link!(vi, spl)
     !linked && return identity(p) 
     vi[spl] = p
@@ -190,22 +190,22 @@ function transform!(p::AbstractArray, vi::DynamicPPL.VarInfo, ::constrained_spac
 end
 
 function transform!(p::AbstractArray, vi::DynamicPPL.VarInfo, ::constrained_space{false})
-  spl = DynamicPPL.SampleFromPrior()
+    spl = DynamicPPL.SampleFromPrior()
 
-  linked = DynamicPPL.islinked(vi, spl)
-  linked && DynamicPPL.invlink!(vi, spl)
-  vi[spl] = p
-  DynamicPPL.link!(vi, spl)
-  p .= vi[spl]
-  !linked && DynamicPPL.invlink!(vi, spl)
+    linked = DynamicPPL.islinked(vi, spl)
+    linked && DynamicPPL.invlink!(vi, spl)
+    vi[spl] = p
+    DynamicPPL.link!(vi, spl)
+    p .= vi[spl]
+    !linked && DynamicPPL.invlink!(vi, spl)
 
-  return nothing
+    return nothing
 end
 
 function transform(p::AbstractArray, vi::DynamicPPL.VarInfo, con::constrained_space)
     tp = copy(p)
     transform!(tp, vi, con)
-  return tp
+    return tp
 end
 
 abstract type AbstractTransform end
@@ -228,13 +228,13 @@ function (t::Init)()
     return t.vi[DynamicPPL.SampleFromPrior()]
 end 
 
- function get_parameter_bounds(model::DynamicPPL.Model)
+function get_parameter_bounds(model::DynamicPPL.Model)
     vi = DynamicPPL.VarInfo(model)
     spl = DynamicPPL.SampleFromPrior()
 
     ## Check link status of vi
     linked = DynamicPPL.islinked(vi, spl) 
-  
+    
     ## transform into unconstrained
     !linked && DynamicPPL.link!(vi, spl)
     
@@ -245,34 +245,34 @@ end
 end
 
 function _optim_objective(model::DynamicPPL.Model, ::MAP, ::constrained_space{false})
-  ctx = OptimizationContext(DynamicPPL.DefaultContext())
-  obj = OptimLogDensity(model, ctx)
+    ctx = OptimizationContext(DynamicPPL.DefaultContext())
+    obj = OptimLogDensity(model, ctx)
 
-  transform!(obj)
-  init = Init(obj.vi, constrained_space{false}())
-  t = ParameterTransform(obj.vi, constrained_space{true}())
+    transform!(obj)
+    init = Init(obj.vi, constrained_space{false}())
+    t = ParameterTransform(obj.vi, constrained_space{true}())
 
-  return (obj=obj, init = init, transform=t)
+    return (obj=obj, init = init, transform=t)
 end
 
 function _optim_objective(model::DynamicPPL.Model, ::MAP, ::constrained_space{true})
     ctx = OptimizationContext(DynamicPPL.DefaultContext())
     obj = OptimLogDensity(model, ctx)
-  
+    
     init = Init(obj.vi, constrained_space{true}())
     t = ParameterTransform(obj.vi, constrained_space{true}())
-      
+    
     return (obj=obj, init = init, transform=t)
-  end
+end
 
 function _optim_objective(model::DynamicPPL.Model, ::MLE,  ::constrained_space{false})
     ctx = OptimizationContext(DynamicPPL.LikelihoodContext())
     obj = OptimLogDensity(model, ctx)
-  
+    
     transform!(obj)
     init = Init(obj.vi, constrained_space{false}())
     t = ParameterTransform(obj.vi, constrained_space{true}())
-  
+    
     return (obj=obj, init = init, transform=t)
 end
 
@@ -293,10 +293,10 @@ end
 
 function optim_function(model::DynamicPPL.Model, estimator::Union{MLE, MAP}; constrained::Bool=true, autoad::Union{Nothing, AbstractADType}=nothing)
     obj, init, t = optim_objective(model, estimator; constrained=constrained)
-  
+    
     l(x,p) = obj(x)
     f = isa(autoad, AbstractADType) ? OptimizationFunction(l, autoad) : OptimizationFunction(l; grad = (G,x,p) -> obj(nothing, G, nothing, x), hess = (H,x,p) -> obj(nothing, nothing, H, x))
-  
+    
     return (func=f, init=init, transform = t)
 end
 
diff --git a/test/Project.toml b/test/Project.toml
index 0705adf1d0..14277bffe5 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -38,7 +38,7 @@ CmdStan = "6.0.8"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-DynamicPPL = "0.16"
+DynamicPPL = "0.17.2"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12"
 GalacticOptim = "0.4, 1, 2"
diff --git a/test/inference/gibbs.jl b/test/inference/gibbs.jl
index 9077aecbe8..e0d827afdb 100644
--- a/test/inference/gibbs.jl
+++ b/test/inference/gibbs.jl
@@ -108,7 +108,8 @@
             end
         end
         model = imm(randn(100), 1.0);
-        sample(model, Gibbs(MH(:z), HMC(0.01, 4, :m)), 100);
+        # https://github.com/TuringLang/Turing.jl/issues/1725
+        # sample(model, Gibbs(MH(:z), HMC(0.01, 4, :m)), 100);
         sample(model, Gibbs(PG(10, :z), HMC(0.01, 4, :m)), 100);
     end
 end
diff --git a/test/test_utils/models.jl b/test/test_utils/models.jl
index cec6a7cb9b..af207621bb 100644
--- a/test/test_utils/models.jl
+++ b/test/test_utils/models.jl
@@ -51,94 +51,3 @@ MoGtest_default = MoGtest([1.0 1.0 4.0 4.0])
 
 # Declare empty model to make the Sampler constructor work.
 @model empty_model() = begin x = 1; end
-
-# A collection of models for which the mean-of-means for the posterior should
-# be same.
-@model function gdemo1(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
-    # `dot_assume` and `observe`
-    m = TV(undef, length(x))
-    m .~ Normal()
-    x ~ MvNormal(m, 0.25 * I)
-end
-
-@model function gdemo2(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
-    # `assume` with indexing and `observe`
-    m = TV(undef, length(x))
-    for i in eachindex(m)
-        m[i] ~ Normal()
-    end
-    x ~ MvNormal(m, 0.25 * I)
-end
-
-@model function gdemo3(x = 10 * ones(2))
-    # Multivariate `assume` and `observe`
-    m ~ MvNormal(zero(x), I)
-    x ~ MvNormal(m, 0.25 * I)
-end
-
-@model function gdemo4(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
-    # `dot_assume` and `observe` with indexing
-    m = TV(undef, length(x))
-    m .~ Normal()
-    for i in eachindex(x)
-        x[i] ~ Normal(m[i], 0.5)
-    end
-end
-
-# Using vector of `length` 1 here so the posterior of `m` is the same
-# as the others.
-@model function gdemo5(x = 10 * ones(1))
-    # `assume` and `dot_observe`
-    m ~ Normal()
-    x .~ Normal(m, 0.5)
-end
-
-@model function gdemo6()
-    # `assume` and literal `observe`
-    m ~ MvNormal(zeros(2), I)
-    [10.0, 10.0] ~ MvNormal(m, 0.25 * I)
-end
-
-@model function gdemo7(::Type{TV} = Vector{Float64}) where {TV}
-    # `dot_assume` and literal `observe` with indexing
-    m = TV(undef, 2)
-    m .~ Normal()
-    for i in eachindex(m)
-        10.0 ~ Normal(m[i], 0.5)
-    end
-end
-
-@model function gdemo8()
-    # `assume` and literal `dot_observe`
-    m ~ Normal()
-    [10.0, ] .~ Normal(m, 0.5)
-end
-
-@model function _prior_dot_assume(::Type{TV} = Vector{Float64}) where {TV}
-    m = TV(undef, 2)
-    m .~ Normal()
-
-    return m
-end
-
-@model function gdemo9()
-    # Submodel prior
-    m = @submodel _prior_dot_assume()
-    for i in eachindex(m)
-        10.0 ~ Normal(m[i], 0.5)
-    end
-end
-
-@model function _likelihood_dot_observe(m, x)
-    x ~ MvNormal(m, 0.25 * I)
-end
-
-@model function gdemo10(x = 10 * ones(2), ::Type{TV} = Vector{Float64}) where {TV}
-    m = TV(undef, length(x))
-    m .~ Normal()
-
-    # Submodel likelihood
-    @submodel _likelihood_dot_observe(m, x)
-end
-
-const gdemo_models = (gdemo1(), gdemo2(), gdemo3(), gdemo4(), gdemo5(), gdemo6(), gdemo7(), gdemo8(), gdemo9(), gdemo10())
diff --git a/test/test_utils/numerical_tests.jl b/test/test_utils/numerical_tests.jl
index c3f29fbffa..ea8a0d8d5d 100644
--- a/test/test_utils/numerical_tests.jl
+++ b/test/test_utils/numerical_tests.jl
@@ -66,7 +66,7 @@ function check_MoGtest_default(chain; atol=0.2, rtol=0.0)
 end
 
 function check_gdemo_models(alg, nsamples, args...; atol=0.0, rtol=0.2, kwargs...)
-    @testset "$(alg) on $(m.name)" for m in gdemo_models
+    @testset "$(alg) on $(m.name)" for m in DynamicPPL.TestUtils.DEMO_MODELS
         # Log this so that if something goes wrong, we can identify the
         # algorithm and model.
         μ = mean(Array(sample(m, alg, nsamples, args...; kwargs...)))