Merge remote-tracking branch 'origin/master' into dw/enzyme

Author: penelopeysm

HISTORY.md

@@ -4,13 +4,11 @@
 
 0.36.0 introduces a new Gibbs sampler. It's been included in several previous releases as `Turing.Experimental.Gibbs`, but now takes over the old Gibbs sampler, which gets removed completely.
 
-The new Gibbs sampler supports the same user-facing interface as the old one. However, given
-that the internals of it having been completely rewritten in a very different manner, there
-may be accidental breakage that we haven't anticipated. Please report any you find.
+The new Gibbs sampler currently supports the same user-facing interface as the old one, but the old constructors have been deprecated, and will be removed in the future. Also, given that the internals have been completely rewritten in a very different manner, there may be accidental breakage that we haven't anticipated. Please report any you find.
 
 `GibbsConditional` has also been removed. It was never very user-facing, but it was exported, so technically this is breaking.
 
-The old Gibbs constructor relied on being called with several subsamplers, and each of the constructors of the subsamplers would take as arguments the symbols for the variables that they are to sample, e.g. `Gibbs(HMC(:x), MH(:y))`. This constructor has been deprecated, and will be removed in the future. The new constructor works by assigning samplers to either symbols or `VarNames`, e.g. `Gibbs(; x=HMC(), y=MH())` or `Gibbs(@varname(x) => HMC(), @varname(y) => MH())`. This allows more granular specification of which sampler to use for which variable.
+The old Gibbs constructor relied on being called with several subsamplers, and each of the constructors of the subsamplers would take as arguments the symbols for the variables that they are to sample, e.g. `Gibbs(HMC(:x), MH(:y))`. This constructor has been deprecated, and will be removed in the future. The new constructor works by mapping symbols, `VarName`s, or iterables thereof to samplers, e.g. `Gibbs(x=>HMC(), y=>MH())`, `Gibbs(@varname(x) => HMC(), @varname(y) => MH())`, `Gibbs((:x, :y) => NUTS(), :z => MH())`. This allows more granular specification of which sampler to use for which variable.
 
 Likewise, the old constructor for calling one subsampler more often than another, `Gibbs((HMC(0.01, 4, :x), 2), (MH(:y), 1))` has been deprecated. The new way to do this is to use `RepeatSampler`, also introduced at this version: `Gibbs(@varname(x) => RepeatSampler(HMC(0.01, 4), 2), @varname(y) => MH())`.
 

docs/src/api.md

@@ -58,7 +58,6 @@ even though [`Prior()`](@ref) is actually defined in the `Turing.Inference` modu
 | `Emcee`              | [`Turing.Inference.Emcee`](@ref)              | Affine-invariant ensemble sampler                                   |
 | `ESS`                | [`Turing.Inference.ESS`](@ref)                | Elliptical slice sampling                                           |
 | `Gibbs`              | [`Turing.Inference.Gibbs`](@ref)              | Gibbs sampling                                                      |
-| `GibbsConditional`   | [`Turing.Inference.GibbsConditional`](@ref)   | A "pseudo-sampler" to provide analytical conditionals to `Gibbs`    |
 | `HMC`                | [`Turing.Inference.HMC`](@ref)                | Hamiltonian Monte Carlo                                             |
 | `SGLD`               | [`Turing.Inference.SGLD`](@ref)               | Stochastic gradient Langevin dynamics                               |
 | `SGHMC`              | [`Turing.Inference.SGHMC`](@ref)              | Stochastic gradient Hamiltonian Monte Carlo                         |

src/mcmc/gibbs.jl

@@ -292,15 +292,40 @@ function set_selector(x::RepeatSampler)
 end
 set_selector(x::InferenceAlgorithm) = DynamicPPL.Sampler(x, DynamicPPL.Selector(0))
 
+to_varname_list(x::Union{VarName,Symbol}) = [VarName(x)]
+# Any other value is assumed to be an iterable of VarNames and Symbols.
+to_varname_list(t) = collect(map(VarName, t))
+
 """
     Gibbs
 
 A type representing a Gibbs sampler.
 
+# Constructors
+
+`Gibbs` needs to be given a set of pairs of variable names and samplers. Instead of a single
+variable name per sampler, one can also give an iterable of variables, all of which are
+sampled by the same component sampler.
+
+Each variable name can be given as either a `Symbol` or a `VarName`.
+
+Some examples of valid constructors are:
+```julia
+Gibbs(:x => NUTS(), :y => MH())
+Gibbs(@varname(x) => NUTS(), @varname(y) => MH())
+Gibbs((@varname(x), :y) => NUTS(), :z => MH())
+```
+
+Currently only variable names without indexing are supported, so for instance
+`Gibbs(@varname(x[1]) => NUTS())` does not work. This will hopefully change in the future.
+
 # Fields
 $(TYPEDFIELDS)
 """
-struct Gibbs{V,A} <: InferenceAlgorithm
+struct Gibbs{N,V<:NTuple{N,AbstractVector{<:VarName}},A<:NTuple{N,Any}} <:
+       InferenceAlgorithm
+    # TODO(mhauru) Revisit whether A should have a fixed element type once
+    # InferenceAlgorithm/Sampler types have been cleaned up.
     "varnames representing variables for each sampler"
     varnames::V
     "samplers for each entry in `varnames`"
@@ -310,40 +335,30 @@ struct Gibbs{V,A} <: InferenceAlgorithm
         if length(varnames) != length(samplers)
             throw(ArgumentError("Number of varnames and samplers must match."))
         end
+
         for spl in samplers
             if !isgibbscomponent(spl)
                 msg = "All samplers must be valid Gibbs components, $(spl) is not."
                 throw(ArgumentError(msg))
             end
         end
-        return new{typeof(varnames),typeof(samplers)}(varnames, samplers)
-    end
-end
 
-to_varname(vn::VarName) = vn
-to_varname(s::Symbol) = VarName{s}()
-# Any other value is assumed to be an iterable.
-to_varname(t) = map(to_varname, collect(t))
-
-# NamedTuple
-Gibbs(; algs...) = Gibbs(NamedTuple(algs))
-function Gibbs(algs::NamedTuple)
-    return Gibbs(map(to_varname, keys(algs)), map(set_selector ∘ drop_space, values(algs)))
+        # Ensure that samplers have the same selector, and that varnames are lists of
+        # VarNames.
+        samplers = tuple(map(set_selector ∘ drop_space, samplers)...)
+        varnames = tuple(map(to_varname_list, varnames)...)
+        return new{length(samplers),typeof(varnames),typeof(samplers)}(varnames, samplers)
+    end
 end
 
-# AbstractDict
-function Gibbs(algs::AbstractDict)
-    return Gibbs(
-        map(to_varname, collect(keys(algs))), map(set_selector ∘ drop_space, values(algs))
-    )
-end
 function Gibbs(algs::Pair...)
-    return Gibbs(map(to_varname ∘ first, algs), map(set_selector ∘ drop_space ∘ last, algs))
+    return Gibbs(map(first, algs), map(last, algs))
 end
 
 # The below two constructors only provide backwards compatibility with the constructor of
 # the old Gibbs sampler. They are deprecated and will be removed in the future.
-function Gibbs(algs::InferenceAlgorithm...)
+function Gibbs(alg1::InferenceAlgorithm, other_algs::InferenceAlgorithm...)
+    algs = [alg1, other_algs...]
     varnames = map(algs) do alg
         space = getspace(alg)
         if (space isa VarName)
@@ -365,7 +380,11 @@ function Gibbs(algs::InferenceAlgorithm...)
     return Gibbs(varnames, map(set_selector ∘ drop_space, algs))
 end
 
-function Gibbs(algs_with_iters::Tuple{<:InferenceAlgorithm,Int}...)
+function Gibbs(
+    alg_with_iters1::Tuple{<:InferenceAlgorithm,Int},
+    other_algs_with_iters::Tuple{<:InferenceAlgorithm,Int}...,
+)
+    algs_with_iters = [alg_with_iters1, other_algs_with_iters...]
     algs = Iterators.map(first, algs_with_iters)
     iters = Iterators.map(last, algs_with_iters)
     algs_duplicated = Iterators.flatten((
@@ -384,64 +403,80 @@ struct GibbsState{V<:DynamicPPL.AbstractVarInfo,S}
     states::S
 end
 
-_maybevec(x) = vec(x)  # assume it's iterable
-_maybevec(x::Tuple) = [x...]
-_maybevec(x::VarName) = [x]
-_maybevec(x::Symbol) = [x]
-
 varinfo(state::GibbsState) = state.vi
 
 function DynamicPPL.initialstep(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
     spl::DynamicPPL.Sampler{<:Gibbs},
-    vi_base::DynamicPPL.AbstractVarInfo;
+    vi::DynamicPPL.AbstractVarInfo;
     initial_params=nothing,
     kwargs...,
 )
     alg = spl.alg
     varnames = alg.varnames
     samplers = alg.samplers
 
-    # Run the model once to get the varnames present + initial values to condition on.
-    vi = DynamicPPL.VarInfo(rng, model)
-    if initial_params !== nothing
-        vi = DynamicPPL.unflatten(vi, initial_params)
-    end
+    vi, states = gibbs_initialstep_recursive(
+        rng, model, varnames, samplers, vi; initial_params=initial_params, kwargs...
+    )
+    return Transition(model, vi), GibbsState(vi, states)
+end
 
-    # Initialise each component sampler in turn, collect all their states.
-    states = []
-    for (varnames_local, sampler_local) in zip(varnames, samplers)
-        varnames_local = _maybevec(varnames_local)
-        # Get the initial values for this component sampler.
-        initial_params_local = if initial_params === nothing
-            nothing
-        else
-            DynamicPPL.subset(vi, varnames_local)[:]
-        end
+"""
+Take the first step of MCMC for the first component sampler, and call the same function
+recursively on the remaining samplers, until no samplers remain. Return the global VarInfo
+and a tuple of initial states for all component samplers.
+"""
+function gibbs_initialstep_recursive(
+    rng, model, varname_vecs, samplers, vi, states=(); initial_params=nothing, kwargs...
+)
+    # End recursion
+    if isempty(varname_vecs) && isempty(samplers)
+        return vi, states
+    end
 
-        # Construct the conditioned model.
-        model_local, context_local = make_conditional(model, varnames_local, vi)
+    varnames, varname_vecs_tail... = varname_vecs
+    sampler, samplers_tail... = samplers
 
-        # Take initial step.
-        _, new_state_local = 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...,
-        )
-        new_vi_local = varinfo(new_state_local)
-        # Merge in any new variables that were introduced during the step, but that
-        # were not in the domain of the current sampler.
-        vi = merge(vi, get_global_varinfo(context_local))
-        # Merge the new values for all the variables sampled by the current sampler.
-        vi = merge(vi, new_vi_local)
-        push!(states, new_state_local)
+    # Get the initial values for this component sampler.
+    initial_params_local = if initial_params === nothing
+        nothing
+    else
+        DynamicPPL.subset(vi, varnames)[:]
     end
-    return Transition(model, vi), GibbsState(vi, states)
+
+    # Construct the conditioned model.
+    conditioned_model, context = make_conditional(model, varnames, vi)
+
+    # Take initial step with the current sampler.
+    _, new_state = AbstractMCMC.step(
+        rng,
+        conditioned_model,
+        sampler;
+        # 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...,
+    )
+    new_vi_local = varinfo(new_state)
+    # Merge in any new variables that were introduced during the step, but that
+    # were not in the domain of the current sampler.
+    vi = merge(vi, get_global_varinfo(context))
+    # Merge the new values for all the variables sampled by the current sampler.
+    vi = merge(vi, new_vi_local)
+
+    states = (states..., new_state)
+    return gibbs_initialstep_recursive(
+        rng,
+        model,
+        varname_vecs_tail,
+        samplers_tail,
+        vi,
+        states;
+        initial_params=initial_params,
+        kwargs...,
+    )
 end
 
 function AbstractMCMC.step(
@@ -458,17 +493,7 @@ function AbstractMCMC.step(
     states = state.states
     @assert length(samplers) == length(state.states)
 
-    # TODO: move this into a recursive function so we can unroll when reasonable?
-    for index in 1:length(samplers)
-        # Take the inner step.
-        sampler_local = samplers[index]
-        state_local = states[index]
-        varnames_local = _maybevec(varnames[index])
-        vi, new_state_local = gibbs_step_inner(
-            rng, model, varnames_local, sampler_local, state_local, vi; kwargs...
-        )
-        states = Accessors.setindex(states, new_state_local, index)
-    end
+    vi, states = gibbs_step_recursive(rng, model, varnames, samplers, states, vi; kwargs...)
     return Transition(model, vi), GibbsState(vi, states)
 end
 
@@ -592,19 +617,33 @@ function match_linking!!(varinfo_local, prev_state_local, model)
     return varinfo_local
 end
 
-function gibbs_step_inner(
+"""
+Run a Gibbs step for the first varname/sampler/state tuple, and recursively call the same
+function on the tail, until there are no more samplers left.
+"""
+function gibbs_step_recursive(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
-    varnames_local,
-    sampler_local,
-    state_local,
-    global_vi;
+    varname_vecs,
+    samplers,
+    states,
+    global_vi,
+    new_states=();
     kwargs...,
 )
+    # End recursion.
+    if isempty(varname_vecs) && isempty(samplers) && isempty(states)
+        return global_vi, new_states
+    end
+
+    varnames, varname_vecs_tail... = varname_vecs
+    sampler, samplers_tail... = samplers
+    state, states_tail... = states
+
     # Construct the conditional model and the varinfo that this sampler should use.
-    model_local, context_local = make_conditional(model, varnames_local, global_vi)
-    varinfo_local = subset(global_vi, varnames_local)
-    varinfo_local = match_linking!!(varinfo_local, state_local, model)
+    conditioned_model, context = make_conditional(model, varnames, global_vi)
+    vi = subset(global_vi, varnames)
+    vi = match_linking!!(vi, state, model)
 
     # TODO(mhauru) The below may be overkill. If the varnames for this sampler are not
     # sampled by other samplers, we don't need to `setparams`, but could rather simply
@@ -615,18 +654,25 @@ function gibbs_step_inner(
     # going to be a significant expense anyway.
     # Set the state of the current sampler, accounting for any changes made by other
     # samplers.
-    state_local = setparams_varinfo!!(
-        model_local, sampler_local, state_local, varinfo_local
-    )
+    state = setparams_varinfo!!(conditioned_model, sampler, state, vi)
 
     # Take a step with the local sampler.
-    new_state_local = last(
-        AbstractMCMC.step(rng, model_local, sampler_local, state_local; kwargs...)
-    )
+    new_state = last(AbstractMCMC.step(rng, conditioned_model, sampler, state; kwargs...))
 
-    new_vi_local = varinfo(new_state_local)
+    new_vi_local = varinfo(new_state)
     # Merge the latest values for all the variables in the current sampler.
-    new_global_vi = merge(get_global_varinfo(context_local), new_vi_local)
+    new_global_vi = merge(get_global_varinfo(context), new_vi_local)
     new_global_vi = setlogp!!(new_global_vi, getlogp(new_vi_local))
-    return new_global_vi, new_state_local
+
+    new_states = (new_states..., new_state)
+    return gibbs_step_recursive(
+        rng,
+        model,
+        varname_vecs_tail,
+        samplers_tail,
+        states_tail,
+        new_global_vi,
+        new_states;
+        kwargs...,
+    )
 end

test/dynamicppl/compiler.jl

@@ -54,7 +54,7 @@ const gdemo_default = gdemo_d()
 
         smc = SMC()
         pg = PG(10)
-        gibbs = Gibbs(; p=HMC(0.2, 3), x=PG(10))
+        gibbs = Gibbs(:p => HMC(0.2, 3), :x => PG(10))
 
         chn_s = sample(testbb(obs), smc, 1000)
         chn_p = sample(testbb(obs), pg, 2000)
@@ -81,7 +81,7 @@ const gdemo_default = gdemo_d()
             return s, m
         end
 
-        gibbs = Gibbs(; s=PG(10), m=HMC(0.4, 8))
+        gibbs = Gibbs(:s => PG(10), :m => HMC(0.4, 8))
         chain = sample(fggibbstest(xs), gibbs, 2)
     end
     @testset "new grammar" begin
@@ -177,7 +177,7 @@ const gdemo_default = gdemo_d()
     end
 
     @testset "sample" begin
-        alg = Gibbs(; m=HMC(0.2, 3), s=PG(10))
+        alg = Gibbs(:m => HMC(0.2, 3), :s => PG(10))
         chn = sample(gdemo_default, alg, 1000)
     end