Fixes and improvements to experimental Gibbs

src/experimental/gibbs.jl

@@ -78,7 +78,7 @@
     # Short-circuits the tilde assume if `vn` is present in `context`.
     if has_conditioned_gibbs(context, vns)
         value = reconstruct_getvalue(right, get_conditioned_gibbs(context, vns))
-        return value, broadcast_logpdf(right, values), vi
+        return value, broadcast_logpdf(right, value), vi
     end
 
     # Otherwise, falls back to the default behavior.
@@ -90,8 +90,8 @@
 )
     # Short-circuits the tilde assume if `vn` is present in `context`.
     if has_conditioned_gibbs(context, vns)
-        values = reconstruct_getvalue(right, get_conditioned_gibbs(context, vns))
-        return values, broadcast_logpdf(right, values), vi
+        value = reconstruct_getvalue(right, get_conditioned_gibbs(context, vns))
+        return value, broadcast_logpdf(right, value), vi
     end
 
     # Otherwise, falls back to the default behavior.
@@ -144,14 +144,14 @@
 Return a `GibbsContext` with the values extracted from the given `varinfos` treated as conditioned.
 """
 function condition_gibbs(context::DynamicPPL.AbstractContext, varinfo::DynamicPPL.AbstractVarInfo)
-    return DynamicPPL.condition(context, DynamicPPL.values_as(varinfo, preferred_value_type(varinfo)))
+    return condition_gibbs(context, DynamicPPL.values_as(varinfo, preferred_value_type(varinfo)))
 end
-function DynamicPPL.condition(
+function condition_gibbs(
     context::DynamicPPL.AbstractContext,
     varinfo::DynamicPPL.AbstractVarInfo,
     varinfos::DynamicPPL.AbstractVarInfo...
 )
-    return DynamicPPL.condition(DynamicPPL.condition(context, varinfo), varinfos...)
+    return condition_gibbs(condition_gibbs(context, varinfo), varinfos...)
 end
 # Allow calling this on a `DynamicPPL.Model` directly.
 function condition_gibbs(model::DynamicPPL.Model, values...)
@@ -238,6 +238,9 @@
     return Gibbs(map(first, algs), map(wrap_algorithm_maybe, map(last, algs)))
 end
 
+# TODO: Remove when no longer needed.
+DynamicPPL.getspace(::Gibbs) = ()
+
 struct GibbsState{V<:DynamicPPL.AbstractVarInfo,S}
     vi::V
     states::S
@@ -252,6 +255,7 @@
     model::DynamicPPL.Model,
     spl::DynamicPPL.Sampler{<:Gibbs},
     vi_base::DynamicPPL.AbstractVarInfo;
+    initial_params=nothing,
     kwargs...,
 )
     alg = spl.alg
@@ -260,15 +264,35 @@
 
     # 1. Run the model once to get the varnames present + initial values to condition on.
     vi_base = DynamicPPL.VarInfo(model)
+
+    # Simple way of setting the initial parameters: set them in the `vi_base`
+    # if they are given so they propagate to the subset varinfos used by each sampler.
+    if initial_params !== nothing
+        vi_base = DynamicPPL.unflatten(vi_base, initial_params)
+    end
+
+    # Create the varinfos for each sampler.
     varinfos = map(Base.Fix1(DynamicPPL.subset, vi_base) ∘ _maybevec, varnames)
+    initial_params_all = if initial_params === nothing
+        fill(nothing, length(varnames))
+    else
+        # Extract from the `vi_base`, which should have the values set correctly from above.
+        map(vi -> vi[:], varinfos)
+    end
 
     # 2. Construct a varinfo for every vn + sampler combo.
-    states_and_varinfos = map(samplers, varinfos) do sampler_local, varinfo_local
+    states_and_varinfos = map(samplers, varinfos, initial_params_all) do sampler_local, varinfo_local, initial_params_local
         # Construct the conditional model.
         model_local = make_conditional(model, varinfo_local, varinfos)
 
         # Take initial step.
-        new_state_local = last(AbstractMCMC.step(rng, model_local, sampler_local; kwargs...))
+        new_state_local = last(AbstractMCMC.step(
+            rng, model_local, sampler_local;
+            # FIXME: This will cause issues if the sampler expects initial params in unconstrained space.
+            # This is not the case for any samplers in Turing.jl, but will be for external samplers, etc.
+            initial_params=initial_params_local,
+            kwargs...
+        ))
 
         # Return the new state and the invlinked `varinfo`.
         vi_local_state = Turing.Inference.varinfo(new_state_local)
@@ -284,7 +308,7 @@
     varinfos = map(last, states_and_varinfos)
 
     # Update the base varinfo from the first varinfo and replace it.
-    varinfos_new = DynamicPPL.setindex!!(varinfos, vi_base, 1)
+    varinfos_new = DynamicPPL.setindex!!(varinfos, merge(vi_base, first(varinfos)), 1)
     # Merge the updated initial varinfo with the rest of the varinfos + update the logp.
     vi = DynamicPPL.setlogp!!(
         reduce(merge, varinfos_new),
@@ -365,12 +389,7 @@
 end
 
 # TODO: Remove `rng`?
-"""
-    recompute_logprob!!(rng, model, sampler, state)
-
-Recompute the log-probability of the `model` based on the given `state` and return the resulting state.
-"""
-function recompute_logprob!!(
+function Turing.Inference.recompute_logprob!!(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
     sampler::DynamicPPL.Sampler,
@@ -436,7 +455,7 @@
         state_local,
         state_previous
     )
-        current_state_local = recompute_logprob!!(
+        state_local = Turing.Inference.recompute_logprob!!(
             rng,
             model_local,
             sampler_local,
@@ -450,7 +469,7 @@
             rng,
             model_local,
             sampler_local,
-            current_state_local;
+            state_local;
             kwargs...,
         ),
     )

src/mcmc/Inference.jl

@@ -114,6 +114,8 @@
     end
 end
 
+DynamicPPL.getspace(::ExternalSampler) = ()
+
 """
     requires_unconstrained_space(sampler::ExternalSampler)
 

src/mcmc/abstractmcmc.jl

@@ -17,10 +17,52 @@
     return transition_to_turing(parent(f), transition)
 end
 
+"""
+    getmodel(f)
+
+Return the `DynamicPPL.Model` wrapped in the given log-density function `f`.
+"""
+getmodel(f::LogDensityProblemsAD.ADGradientWrapper) = getmodel(parent(f))
+getmodel(f::DynamicPPL.LogDensityFunction) = f.model
+
+# FIXME: We'll have to overload this for every AD backend since some of the AD backends
+# will cache certain parts of a given model, e.g. the tape, which results in a discrepancy
+# between the primal (forward) and dual (backward).
+"""
+    setmodel(f, model)
+
+Set the `DynamicPPL.Model` in the given log-density function `f` to `model`.
+
+!!! warning
+    Note that if `f` is a `LogDensityProblemsAD.ADGradientWrapper` wrapping a
+    `DynamicPPL.LogDensityFunction`, performing an update of the `model` in `f`
+    might require recompilation of the gradient tape, depending on the AD backend.
+"""
+function setmodel(f::LogDensityProblemsAD.ADGradientWrapper, model::DynamicPPL.Model)
+    return Accessors.@set f.ℓ = setmodel(f.ℓ, model)
+end
+function setmodel(f::DynamicPPL.LogDensityFunction, model::DynamicPPL.Model)
+    return Accessors.@set f.model = model
+end
+
+function varinfo_from_logdensityfn(f::LogDensityProblemsAD.ADGradientWrapper)
+    return varinfo_from_logdensityfn(parent(f))
+end
+varinfo_from_logdensityfn(f::DynamicPPL.LogDensityFunction) = f.varinfo
+
+function varinfo(state::TuringState)
+    θ = getparams(getmodel(state.logdensity), state.state)
+    # TODO: Do we need to link here first?
+    return DynamicPPL.unflatten(varinfo_from_logdensityfn(state.logdensity), θ)
+end
+
 # NOTE: Only thing that depends on the underlying sampler.
 # Something similar should be part of AbstractMCMC at some point:
 # https://github.com/TuringLang/AbstractMCMC.jl/pull/86
 getparams(::DynamicPPL.Model, transition::AdvancedHMC.Transition) = transition.z.θ
+function getparams(model::DynamicPPL.Model, state::AdvancedHMC.HMCState)
+    return getparams(model, state.transition)
+end
 getstats(transition::AdvancedHMC.Transition) = transition.stat
 
 getparams(::DynamicPPL.Model, transition::AdvancedMH.Transition) = transition.params
@@ -33,13 +75,59 @@
     return Accessors.@set f.ℓ = setvarinfo(f.ℓ, varinfo)
 end
 
+"""
+    recompute_logprob!!(rng, model, sampler, state)
+
+Recompute the log-probability of the `model` based on the given `state` and return the resulting state.
+"""
+function recompute_logprob!!(
+    rng::Random.AbstractRNG,  # TODO: Do we need the `rng` here?
+    model::DynamicPPL.Model,
+    sampler::DynamicPPL.Sampler{<:ExternalSampler},
+    state,
+)
+    # Re-using the log-density function from the `state` and updating only the `model` field,
+    # since the `model` might now contain different conditioning values.
+    f = setmodel(state.logdensity, model)
+    # Recompute the log-probability with the new `model`.
+    state_inner = recompute_logprob!!(
+        rng, AbstractMCMC.LogDensityModel(f), sampler.alg.sampler, state.state
+    )
+    return state_to_turing(f, state_inner)
+end
+
+function recompute_logprob!!(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    sampler::AdvancedHMC.AbstractHMCSampler,
+    state::AdvancedHMC.HMCState,
+)
+    # Construct hamiltionian.
+    hamiltonian = AdvancedHMC.Hamiltonian(state.metric, model)
+    # Re-compute the log-probability and gradient.
+    return Accessors.@set state.transition.z = AdvancedHMC.phasepoint(
+        hamiltonian, state.transition.z.θ, state.transition.z.r
+    )
+end
+
+function recompute_logprob!!(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    sampler::AdvancedMH.MetropolisHastings,
+    state::AdvancedMH.Transition,
+)
+    logdensity = model.logdensity
+    return Accessors.@set state.lp = LogDensityProblems.logdensity(logdensity, state.params)
+end
+
+# TODO: Do we also support `resume`, etc?
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
     sampler_wrapper::Sampler{<:ExternalSampler};
     initial_state=nothing,
     initial_params=nothing,
-    kwargs...
+    kwargs...,
 )
     alg = sampler_wrapper.alg
     sampler = alg.sampler
@@ -69,7 +157,12 @@
         )
     else
         transition_inner, state_inner = AbstractMCMC.step(
-            rng, AbstractMCMC.LogDensityModel(f), sampler, initial_state; initial_params, kwargs...
+            rng,
+            AbstractMCMC.LogDensityModel(f),
+            sampler,
+            initial_state;
+            initial_params,
+            kwargs...,
         )
     end
     # Update the `state`
@@ -81,7 +174,7 @@
     model::DynamicPPL.Model,
     sampler_wrapper::Sampler{<:ExternalSampler},
     state::TuringState;
-    kwargs...
+    kwargs...,
 )
     sampler = sampler_wrapper.alg.sampler
     f = state.logdensity