Test Gibbs with different adtype in externalsampler to ensure that works

TuringLang · yebai · Jul 16, 2024 · Apr 23, 2024 · Apr 23, 2024 · May 18, 2024
commit 06357c6e8cf7c3bf3d29f7a3fc58ff919878b6e4
diff --git a/test/experimental/gibbs.jl b/test/experimental/gibbs.jl
@@ -8,6 +8,8 @@ using Random
 using Test
 using Turing
 using Turing.Inference: AdvancedHMC, AdvancedMH
+using ForwardDiff: ForwardDiff
+using ReverseDiff: ReverseDiff
 
 function check_transition_varnames(
     transition::Turing.Inference.Transition,
@@ -249,7 +251,9 @@ has_dot_assume(::Model) = true
         model = demo_gibbs_external()
         samplers_inner = [
             externalsampler(AdvancedMH.RWMH(1)),
-            externalsampler(AdvancedHMC.HMC(1e-1, 32)),
+            externalsampler(AdvancedHMC.HMC(1e-1, 32), adtype=AutoForwardDiff()),
+            externalsampler(AdvancedHMC.HMC(1e-1, 32), adtype=AutoReverseDiff()),
+            externalsampler(AdvancedHMC.HMC(1e-1, 32), adtype=AutoReverseDiff(compile=true)),
         ]
         @testset "$(sampler_inner)" for sampler_inner in samplers_inner
             sampler = Turing.Experimental.Gibbs(

diff --git a/src/experimental/gibbs.jl b/src/experimental/gibbs.jl
    # 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, value), vi
    end

    # Otherwise, falls back to the default behavior.
 )
    # 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, value), vi
    end

    # Otherwise, falls back to the default behavior.
 Return a `GibbsContext` with the values extracted from the given `varinfos` treated as conditioned.
 """
 function condition_gibbs(context::DynamicPPL.AbstractContext, varinfo::DynamicPPL.AbstractVarInfo)
    return condition_gibbs(context, DynamicPPL.values_as(varinfo, preferred_value_type(varinfo)))
 end
 function condition_gibbs(
    context::DynamicPPL.AbstractContext,
    varinfo::DynamicPPL.AbstractVarInfo,
    varinfos::DynamicPPL.AbstractVarInfo...
 )
    return condition_gibbs(condition_gibbs(context, varinfo), varinfos...)
 end
 # Allow calling this on a `DynamicPPL.Model` directly.
 function condition_gibbs(model::DynamicPPL.Model, values...)
 end

 # TODO: Remove when no longer needed.
 DynamicPPL.getspace(::Gibbs) = ()

 struct GibbsState{V<:DynamicPPL.AbstractVarInfo,S}
    vi::V

    # 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, 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;
            # 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.
    varinfos = map(last, states_and_varinfos)

    # Update the base varinfo from the first varinfo and replace it.
    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),
 end

 # TODO: Remove `rng`?
 function Turing.Inference.recompute_logprob!!(
    rng::Random.AbstractRNG,
    model::DynamicPPL.Model,
    sampler::DynamicPPL.Sampler,
        state_local,
        state_previous
    )
        state_local = Turing.Inference.recompute_logprob!!(
            rng,
            model_local,
            sampler_local,
diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl
    end
 end

 DynamicPPL.getspace(::ExternalSampler) = ()

 """
    requires_unconstrained_space(sampler::ExternalSampler)
diff --git a/src/mcmc/abstractmcmc.jl b/src/mcmc/abstractmcmc.jl

 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

 """
    setmodel(f, model[, adtype])
    `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,
    adtype::ADTypes.AbstractADType
    #    ReverseDiff.jl in compiled mode will cache the compiled tape, which means that just
    #    replacing the corresponding field with the new model won't be sufficient to obtain
    #    the correct gradients.
    return LogDensityProblemsAD.ADgradient(adtype, setmodel(parent(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


 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},
 )
    # 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, sampler.alg.adtype)
    # 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?
	# 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, value), vi
Check warning on line 81 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L81 `Added line #L81 was not covered by tests`
	end

	# Otherwise, falls back to the default behavior.
	)
	# 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, value), vi
Check warning on line 94 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L93-L94 `Added lines #L93 - L94 were not covered by tests`
	end

	# Otherwise, falls back to the default behavior.
	Return a `GibbsContext` with the values extracted from the given `varinfos` treated as conditioned.
	"""
	function condition_gibbs(context::DynamicPPL.AbstractContext, varinfo::DynamicPPL.AbstractVarInfo)
	return condition_gibbs(context, DynamicPPL.values_as(varinfo, preferred_value_type(varinfo)))
Check warning on line 147 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L147 `Added line #L147 was not covered by tests`
	end
	function condition_gibbs(
Check warning on line 149 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L149 `Added line #L149 was not covered by tests`
	context::DynamicPPL.AbstractContext,
	varinfo::DynamicPPL.AbstractVarInfo,
	varinfos::DynamicPPL.AbstractVarInfo...
	)
	return condition_gibbs(condition_gibbs(context, varinfo), varinfos...)
Check warning on line 154 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L154 `Added line #L154 was not covered by tests`
	end
	# Allow calling this on a `DynamicPPL.Model` directly.
	function condition_gibbs(model::DynamicPPL.Model, values...)
	end

	# TODO: Remove when no longer needed.
	DynamicPPL.getspace(::Gibbs) = ()
Check warning on line 242 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L242 `Added line #L242 was not covered by tests`

	struct GibbsState{V<:DynamicPPL.AbstractVarInfo,S}
	vi::V

	# 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)
Check warning on line 271 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L270-L271 `Added lines #L270 - L271 were not covered by tests`
	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))
Check warning on line 277 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L276-L277 `Added lines #L276 - L277 were not covered by tests`
	else
	# Extract from the `vi_base`, which should have the values set correctly from above.
	map(vi -> vi[:], varinfos)
Check warning on line 280 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L280 `Added line #L280 was not covered by tests`
	end

	# 2. Construct a varinfo for every vn + sampler combo.
	states_and_varinfos = map(samplers, varinfos, initial_params_all) do sampler_local, varinfo_local, initial_params_local
Check warning on line 284 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L284 `Added line #L284 was not covered by tests`
	# Construct the conditional model.
	model_local = make_conditional(model, varinfo_local, varinfos)

	# Take initial step.
	new_state_local = last(AbstractMCMC.step(
Check warning on line 289 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L289 `Added line #L289 was not covered by tests`
	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.
	varinfos = map(last, states_and_varinfos)

	# Update the base varinfo from the first varinfo and replace it.
	varinfos_new = DynamicPPL.setindex!!(varinfos, merge(vi_base, first(varinfos)), 1)
Check warning on line 311 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L311 `Added line #L311 was not covered by tests`
	# Merge the updated initial varinfo with the rest of the varinfos + update the logp.
	vi = DynamicPPL.setlogp!!(
	reduce(merge, varinfos_new),
	end

	# TODO: Remove `rng`?
	function Turing.Inference.recompute_logprob!!(
Check warning on line 392 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L392 `Added line #L392 was not covered by tests`
	rng::Random.AbstractRNG,
	model::DynamicPPL.Model,
	sampler::DynamicPPL.Sampler,
	state_local,
	state_previous
	)
	state_local = Turing.Inference.recompute_logprob!!(
Check warning on line 458 in src/experimental/gibbs.jl Codecov / codecov/patch src/experimental/gibbs.jl#L458 `Added line #L458 was not covered by tests`
	rng,
	model_local,
	sampler_local,
	end
	end

	DynamicPPL.getspace(::ExternalSampler) = ()
Check warning on line 117 in src/mcmc/Inference.jl Codecov / codecov/patch src/mcmc/Inference.jl#L117 `Added line #L117 was not covered by tests`

	"""
	requires_unconstrained_space(sampler::ExternalSampler)

	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
Check warning on line 26 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L25-L26 `Added lines #L25 - L26 were not covered by tests`

	"""
	setmodel(f, model[, adtype])
	`DynamicPPL.LogDensityFunction`, performing an update of the `model` in `f`
	might require recompilation of the gradient tape, depending on the AD backend.
	"""
	function setmodel(
Check warning on line 38 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L38 `Added line #L38 was not covered by tests`
	f::LogDensityProblemsAD.ADGradientWrapper,
	model::DynamicPPL.Model,
	adtype::ADTypes.AbstractADType
	# ReverseDiff.jl in compiled mode will cache the compiled tape, which means that just
	# replacing the corresponding field with the new model won't be sufficient to obtain
	# the correct gradients.
	return LogDensityProblemsAD.ADgradient(adtype, setmodel(parent(f), model))
Check warning on line 51 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L51 `Added line #L51 was not covered by tests`
	end
	function setmodel(f::DynamicPPL.LogDensityFunction, model::DynamicPPL.Model)
	return Accessors.@set f.model = model
Check warning on line 54 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L53-L54 `Added lines #L53 - L54 were not covered by tests`
	end

	function varinfo_from_logdensityfn(f::LogDensityProblemsAD.ADGradientWrapper)
	return varinfo_from_logdensityfn(parent(f))
Check warning on line 58 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L57-L58 `Added lines #L57 - L58 were not covered by tests`
	end
	varinfo_from_logdensityfn(f::DynamicPPL.LogDensityFunction) = f.varinfo
Check warning on line 60 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L60 `Added line #L60 was not covered by tests`

	function varinfo(state::TuringState)
	θ = getparams(getmodel(state.logdensity), state.state)
Check warning on line 63 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L62-L63 `Added lines #L62 - L63 were not covered by tests`
	# TODO: Do we need to link here first?
	return DynamicPPL.unflatten(varinfo_from_logdensityfn(state.logdensity), θ)
Check warning on line 65 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L65 `Added line #L65 was not covered by tests`
	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)
Check warning on line 73 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L72-L73 `Added lines #L72 - L73 were not covered by tests`
	end
	getstats(transition::AdvancedHMC.Transition) = transition.stat


	Recompute the log-probability of the `model` based on the given `state` and return the resulting state.
	"""
	function recompute_logprob!!(
Check warning on line 92 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L92 `Added line #L92 was not covered by tests`
	rng::Random.AbstractRNG, # TODO: Do we need the `rng` here?
	model::DynamicPPL.Model,
	sampler::DynamicPPL.Sampler{<:ExternalSampler},
	)
	# 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, sampler.alg.adtype)
Check warning on line 100 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L100 `Added line #L100 was not covered by tests`
	# Recompute the log-probability with the new `model`.
	state_inner = recompute_logprob!!(
Check warning on line 102 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L102 `Added line #L102 was not covered by tests`
	rng, AbstractMCMC.LogDensityModel(f), sampler.alg.sampler, state.state
	)
	return state_to_turing(f, state_inner)
Check warning on line 105 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L105 `Added line #L105 was not covered by tests`
	end

	function recompute_logprob!!(
Check warning on line 108 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L108 `Added line #L108 was not covered by tests`
	rng::Random.AbstractRNG,
	model::AbstractMCMC.LogDensityModel,
	sampler::AdvancedHMC.AbstractHMCSampler,
	state::AdvancedHMC.HMCState,
	)
	# Construct hamiltionian.
	hamiltonian = AdvancedHMC.Hamiltonian(state.metric, model)
Check warning on line 115 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L115 `Added line #L115 was not covered by tests`
	# Re-compute the log-probability and gradient.
	return Accessors.@set state.transition.z = AdvancedHMC.phasepoint(
Check warning on line 117 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L117 `Added line #L117 was not covered by tests`
	hamiltonian, state.transition.z.θ, state.transition.z.r
	)
	end

	function recompute_logprob!!(
Check warning on line 122 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L122 `Added line #L122 was not covered by tests`
	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)
Check warning on line 129 in src/mcmc/abstractmcmc.jl Codecov / codecov/patch src/mcmc/abstractmcmc.jl#L128-L129 `Added lines #L128 - L129 were not covered by tests`
	end

	# TODO: Do we also support `resume`, etc?