make Hamiltonian directly an AbstractSampler

Author: penelopeysm

src/mcmc/Inference.jl

@@ -248,7 +248,7 @@ getlogevidence(transitions, sampler, state) = missing
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Union{AbstractTransition,AbstractVarInfo}},
     model_or_ldf::Union{DynamicPPL.Model,DynamicPPL.LogDensityFunction},
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
+    spl::AbstractSampler,
     state,
     chain_type::Type{MCMCChains.Chains};
     save_state=false,
@@ -316,7 +316,7 @@ end
 function AbstractMCMC.bundle_samples(
     ts::Vector{<:Union{AbstractTransition,AbstractVarInfo}},
     model_or_ldf::Union{DynamicPPL.Model,DynamicPPL.LogDensityFunction},
-    spl::Union{Sampler{<:InferenceAlgorithm},SampleFromPrior,RepeatSampler},
+    spl::AbstractSampler,
     state,
     chain_type::Type{Vector{NamedTuple}};
     kwargs...,

src/mcmc/abstractmcmc.jl

@@ -22,11 +22,7 @@
 # Because this is a pain to implement all at once, we do it for one sampler at a time.
 # This type tells us which samplers have been 'updated' to the new interface.
 
-# TODO: Eventually, we want to broaden this to InferenceAlgorithm
-const LDFCompatibleAlgorithm = Union{Hamiltonian}
-# TODO: Eventually, we want to broaden this to
-# Union{Sampler{<:InferenceAlgorithm},RepeatSampler}.
-const LDFCompatibleSampler = Union{Sampler{<:LDFCompatibleAlgorithm}}
+const LDFCompatibleSampler = Union{Hamiltonian}
 
 """
     sample(
@@ -251,54 +247,20 @@ end
 ### Everything below this is boring boilerplate for the new interface. ###
 ##########################################################################
 
-function AbstractMCMC.sample(
-    model::Model, alg::LDFCompatibleAlgorithm, N::Integer; kwargs...
-)
-    return AbstractMCMC.sample(Random.default_rng(), model, alg, N; kwargs...)
-end
-
-function AbstractMCMC.sample(
-    ldf::LogDensityFunction, alg::LDFCompatibleAlgorithm, N::Integer; kwargs...
-)
-    return AbstractMCMC.sample(Random.default_rng(), ldf, alg, N; kwargs...)
-end
-
-function AbstractMCMC.sample(
-    model::Model, spl::Sampler{<:LDFCompatibleAlgorithm}, N::Integer; kwargs...
-)
+function AbstractMCMC.sample(model::Model, spl::LDFCompatibleSampler, N::Integer; kwargs...)
     return AbstractMCMC.sample(Random.default_rng(), model, spl, N; kwargs...)
 end
 
 function AbstractMCMC.sample(
-    ldf::LogDensityFunction, spl::Sampler{<:LDFCompatibleAlgorithm}, N::Integer; kwargs...
+    ldf::LogDensityFunction, spl::LDFCompatibleSampler, N::Integer; kwargs...
 )
     return AbstractMCMC.sample(Random.default_rng(), ldf, spl, N; kwargs...)
 end
 
-function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
-    ldf::LogDensityFunction,
-    alg::LDFCompatibleAlgorithm,
-    N::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(rng, ldf, Sampler(alg), N; kwargs...)
-end
-
 function AbstractMCMC.sample(
     rng::Random.AbstractRNG,
     model::Model,
-    alg::LDFCompatibleAlgorithm,
-    N::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(rng, model, Sampler(alg), N; kwargs...)
-end
-
-function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
-    model::Model,
-    spl::Sampler{<:LDFCompatibleAlgorithm},
+    spl::LDFCompatibleSampler,
     N::Integer;
     check_model::Bool=true,
     kwargs...,
@@ -318,33 +280,7 @@ end
 
 function AbstractMCMC.sample(
     model::Model,
-    alg::LDFCompatibleAlgorithm,
-    ensemble::AbstractMCMC.AbstractMCMCEnsemble,
-    N::Integer,
-    n_chains::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(
-        Random.default_rng(), model, alg, ensemble, N, n_chains; kwargs...
-    )
-end
-
-function AbstractMCMC.sample(
-    ldf::LogDensityFunction,
-    alg::LDFCompatibleAlgorithm,
-    ensemble::AbstractMCMC.AbstractMCMCEnsemble,
-    N::Integer,
-    n_chains::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(
-        Random.default_rng(), ldf, alg, ensemble, N, n_chains; kwargs...
-    )
-end
-
-function AbstractMCMC.sample(
-    model::Model,
-    spl::Sampler{<:LDFCompatibleAlgorithm},
+    spl::LDFCompatibleSampler,
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
@@ -357,7 +293,7 @@ end
 
 function AbstractMCMC.sample(
     ldf::LogDensityFunction,
-    spl::Sampler{<:LDFCompatibleAlgorithm},
+    spl::LDFCompatibleSampler,
     ensemble::AbstractMCMC.AbstractMCMCEnsemble,
     N::Integer,
     n_chains::Integer;
@@ -368,30 +304,6 @@ function AbstractMCMC.sample(
     )
 end
 
-function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
-    ldf::LogDensityFunction,
-    alg::LDFCompatibleAlgorithm,
-    ensemble::AbstractMCMC.AbstractMCMCEnsemble,
-    N::Integer,
-    n_chains::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(rng, ldf, Sampler(alg), ensemble, N, n_chains; kwargs...)
-end
-
-function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
-    model::Model,
-    alg::LDFCompatibleAlgorithm,
-    ensemble::AbstractMCMC.AbstractMCMCEnsemble,
-    N::Integer,
-    n_chains::Integer;
-    kwargs...,
-)
-    return AbstractMCMC.sample(rng, model, Sampler(alg), ensemble, N, n_chains; kwargs...)
-end
-
 function AbstractMCMC.sample(
     rng::Random.AbstractRNG,
     model::Model,

src/mcmc/algorithm.jl

@@ -1,3 +1,4 @@
+# TODO(penelopeysm): remove
 """
     InferenceAlgorithm
 
@@ -16,41 +17,30 @@ DynamicPPL.default_chain_type(sampler::Sampler{<:InferenceAlgorithm}) = MCMCChai
 
 """
     update_sample_kwargs(spl::AbstractSampler, N::Integer, kwargs)
-    update_sample_kwargs(spl::InferenceAlgorithm, N::Integer, kwargs)
-Some InferenceAlgorithm implementations carry additional information about
-the keyword arguments that should be passed to `AbstractMCMC.sample`. This
-function provides a hook for them to update the default keyword arguments.
-The default implementation is for no changes to be made to `kwargs`.
-"""
-function update_sample_kwargs(spl::Sampler{<:InferenceAlgorithm}, N::Integer, kwargs)
-    return update_sample_kwargs(spl.alg, N, kwargs)
-end
+
+Some samplers carry additional information about the keyword arguments that
+should be passed to `AbstractMCMC.sample`. This function provides a hook for
+them to update the default keyword arguments. The default implementation is for
+no changes to be made to `kwargs`.
+"""
 update_sample_kwargs(::AbstractSampler, N::Integer, kwargs) = kwargs
-update_sample_kwargs(::InferenceAlgorithm, N::Integer, kwargs) = kwargs
 
 """
     get_adtype(spl::AbstractSampler)
-    get_adtype(spl::InferenceAlgorithm)
-Return the automatic differentiation (AD) backend to use for the sampler.
-This is needed for constructing a LogDensityFunction.
-By default, returns nothing, i.e. the LogDensityFunction that is constructed
-will not know how to calculate its gradients.
-If the sampler or algorithm requires gradient information, then this function
+
+Return the automatic differentiation (AD) backend to use for the sampler. This
+is needed for constructing a LogDensityFunction. By default, returns nothing,
+i.e. the LogDensityFunction that is constructed will not know how to calculate
+its gradients. If the sampler requires gradient information, then this function
 must return an `ADTypes.AbstractADType`.
 """
 get_adtype(::AbstractSampler) = nothing
-get_adtype(::InferenceAlgorithm) = nothing
-get_adtype(spl::Sampler{<:InferenceAlgorithm}) = get_adtype(spl.alg)
 
 """
     requires_unconstrained_space(sampler::AbstractSampler)
-    requires_unconstrained_space(sampler::InferenceAlgorithm)
+
 Return `true` if the sampler / algorithm requires unconstrained space, and
 `false` otherwise. This is used to determine whether the initial VarInfo
 should be linked. Defaults to true.
 """
 requires_unconstrained_space(::AbstractSampler) = true
-requires_unconstrained_space(::InferenceAlgorithm) = true
-function requires_unconstrained_space(spl::Sampler{<:InferenceAlgorithm})
-    return requires_unconstrained_space(spl.alg)
-end

src/mcmc/hmc.jl

@@ -1,8 +1,8 @@
-# InferenceAlgorithm interface
+# AbstractSampler interface for Turing
 
-abstract type Hamiltonian <: InferenceAlgorithm end
+abstract type Hamiltonian <: AbstractMCMC.AbstractSampler end
 
-DynamicPPL.initialsampler(::Sampler{<:Hamiltonian}) = SampleFromUniform()
+DynamicPPL.initialsampler(::Hamiltonian) = DynamicPPL.SampleFromUniform()
 requires_unconstrained_space(::Hamiltonian) = true
 # TODO(penelopeysm): This is really quite dangerous code because it implicitly
 # assumes that any concrete type that subtypes `Hamiltonian` has an adtype
@@ -152,7 +152,7 @@ end
 function AbstractMCMC.step(
     rng::AbstractRNG,
     ldf::LogDensityFunction,
-    spl::Sampler{<:Hamiltonian};
+    spl::Hamiltonian;
     initial_params=nothing,
     nadapts=0,
     kwargs...,
@@ -165,7 +165,7 @@ function AbstractMCMC.step(
     has_initial_params = initial_params !== nothing
 
     # Create a Hamiltonian.
-    metricT = getmetricT(spl.alg)
+    metricT = getmetricT(spl)
     metric = metricT(length(theta))
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
     lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
@@ -184,23 +184,23 @@ function AbstractMCMC.step(
     log_density_old = getlogp(vi)
 
     # Find good eps if not provided one
-    if iszero(spl.alg.ϵ)
+    if iszero(spl.ϵ)
         ϵ = AHMC.find_good_stepsize(rng, hamiltonian, theta)
         @info "Found initial step size" ϵ
     else
-        ϵ = spl.alg.ϵ
+        ϵ = spl.ϵ
     end
 
     # Generate a kernel.
-    kernel = make_ahmc_kernel(spl.alg, ϵ)
+    kernel = make_ahmc_kernel(spl, ϵ)
 
     # Create initial transition and state.
     # Already perform one step since otherwise we don't get any statistics.
     t = AHMC.transition(rng, hamiltonian, kernel, z)
 
     # Adaptation
-    adaptor = AHMCAdaptor(spl.alg, hamiltonian.metric; ϵ=ϵ)
-    if spl.alg isa AdaptiveHamiltonian
+    adaptor = AHMCAdaptor(spl, hamiltonian.metric; ϵ=ϵ)
+    if spl isa AdaptiveHamiltonian
         hamiltonian, kernel, _ = AHMC.adapt!(
             hamiltonian, kernel, adaptor, 1, nadapts, t.z.θ, t.stat.acceptance_rate
         )
@@ -224,7 +224,7 @@ end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     ldf::LogDensityFunction,
-    spl::Sampler{<:Hamiltonian},
+    spl::Hamiltonian,
     state::HMCState;
     nadapts=0,
     kwargs...,
@@ -236,7 +236,7 @@ function AbstractMCMC.step(
 
     # Adaptation
     i = state.i + 1
-    if spl.alg isa AdaptiveHamiltonian
+    if spl isa AdaptiveHamiltonian
         hamiltonian, kernel, _ = AHMC.adapt!(
             hamiltonian,
             state.kernel,
@@ -276,7 +276,7 @@ function get_hamiltonian(model, spl, vi, state, n)
         # using leafcontext(model.context) so could we just remove the argument
         # entirely?)
         DynamicPPL.SamplingContext(spl, DynamicPPL.leafcontext(model.context));
-        adtype=spl.alg.adtype,
+        adtype=spl.adtype,
     )
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
     lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
@@ -441,17 +441,17 @@ getmetricT(::NUTS{<:Any,metricT}) where {metricT} = metricT
 ##### HMC core functions
 #####
 
-getstepsize(sampler::Sampler{<:Hamiltonian}, state) = sampler.alg.ϵ
-getstepsize(sampler::Sampler{<:AdaptiveHamiltonian}, state) = AHMC.getϵ(state.adaptor)
+getstepsize(sampler::Hamiltonian, state) = sampler.ϵ
+getstepsize(sampler::AdaptiveHamiltonian, state) = AHMC.getϵ(state.adaptor)
 function getstepsize(
-    sampler::Sampler{<:AdaptiveHamiltonian},
+    sampler::AdaptiveHamiltonian,
     state::HMCState{TV,TKernel,THam,PhType,AHMC.Adaptation.NoAdaptation},
 ) where {TV,TKernel,THam,PhType}
     return state.kernel.τ.integrator.ϵ
 end
 
-gen_metric(dim::Int, spl::Sampler{<:Hamiltonian}, state) = AHMC.UnitEuclideanMetric(dim)
-function gen_metric(dim::Int, spl::Sampler{<:AdaptiveHamiltonian}, state)
+gen_metric(dim::Int, spl::Hamiltonian, state) = AHMC.UnitEuclideanMetric(dim)
+function gen_metric(dim::Int, spl::AdaptiveHamiltonian, state)
     return AHMC.renew(state.hamiltonian.metric, AHMC.getM⁻¹(state.adaptor.pc))
 end
 
@@ -476,13 +476,11 @@ end
 ####
 #### Compiler interface, i.e. tilde operators.
 ####
-function DynamicPPL.assume(
-    rng, ::Sampler{<:Hamiltonian}, dist::Distribution, vn::VarName, vi
-)
+function DynamicPPL.assume(rng, ::Hamiltonian, dist::Distribution, vn::VarName, vi)
     return DynamicPPL.assume(dist, vn, vi)
 end
 
-function DynamicPPL.observe(::Sampler{<:Hamiltonian}, d::Distribution, value, vi)
+function DynamicPPL.observe(::Hamiltonian, d::Distribution, value, vi)
     return DynamicPPL.observe(d, value, vi)
 end
 

test/ad.jl

@@ -245,18 +245,18 @@ end
     # the tilde-pipeline and thus change the code executed during model
     # evaluation.
     @testset "adtype=$adtype" for adtype in ADTYPES
-        @testset "alg=$alg" for alg in [
+        @testset "spl=$spl" for spl in [
             HMC(0.1, 10; adtype=adtype),
             HMCDA(0.8, 0.75; adtype=adtype),
             NUTS(1000, 0.8; adtype=adtype),
             SGHMC(; learning_rate=0.02, momentum_decay=0.5, adtype=adtype),
             SGLD(; stepsize=PolynomialStepsize(0.25), adtype=adtype),
         ]
-            @info "Testing AD for $alg"
+            @info "Testing AD for $spl"
 
             @testset "model=$(model.f)" for model in DEMO_MODELS
                 rng = StableRNG(123)
-                ctx = DynamicPPL.SamplingContext(rng, DynamicPPL.Sampler(alg))
+                ctx = DynamicPPL.SamplingContext(rng, spl)
                 @test run_ad(model, adtype; context=ctx, test=true, benchmark=false) isa Any
             end
         end
@@ -283,7 +283,7 @@ end
                     model, varnames, deepcopy(global_vi)
                 )
                 rng = StableRNG(123)
-                ctx = DynamicPPL.SamplingContext(rng, DynamicPPL.Sampler(HMC(0.1, 10)))
+                ctx = DynamicPPL.SamplingContext(rng, HMC(0.1, 10))
                 @test run_ad(model, adtype; context=ctx, test=true, benchmark=false) isa Any
             end
         end

test/mcmc/Inference.jl

@@ -76,7 +76,7 @@ using Turing
 
             # run sampler: progress logging should be disabled and
             # it should return a Chains object
-            sampler = Sampler(HMC(0.1, 7))
+            sampler = HMC(0.1, 7)
             chains = sample(StableRNG(seed), gdemo_default, sampler, MCMCThreads(), 10, 4)
             @test chains isa MCMCChains.Chains
         end