Unwrap the other Hamiltonians to make them AbstractSamplers

Author: penelopeysm

ext/TuringDynamicHMCExt.jl

@@ -43,15 +43,12 @@ struct DynamicNUTSState{V<:DynamicPPL.AbstractVarInfo,C,M,S}
     stepsize::S
 end
 
-function DynamicPPL.initialsampler(::DynamicPPL.Sampler{<:DynamicNUTS})
+function DynamicPPL.initialsampler(::DynamicNUTS)
     return DynamicPPL.SampleFromUniform()
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    ldf::DynamicPPL.LogDensityFunction,
-    spl::DynamicPPL.Sampler{<:DynamicNUTS};
-    kwargs...,
+    rng::Random.AbstractRNG, ldf::DynamicPPL.LogDensityFunction, spl::DynamicNUTS; kwargs...
 )
     vi = ldf.varinfo
 
@@ -76,13 +73,13 @@ end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     ldf::DynamicPPL.LogDensityFunction,
-    spl::DynamicPPL.Sampler{<:DynamicNUTS},
+    spl::DynamicNUTS,
     state::DynamicNUTSState;
     kwargs...,
 )
     # Compute next sample.
     vi = state.vi
-    steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ldf, state.stepsize)
+    steps = DynamicHMC.mcmc_steps(rng, spl.sampler, state.metric, ldf, state.stepsize)
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Update the variables.

src/mcmc/sghmc.jl

@@ -51,10 +51,7 @@ struct SGHMCState{V<:AbstractVarInfo,T<:AbstractVector{<:Real}}
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    ldf::DynamicPPL.LogDensityFunction,
-    spl::Sampler{<:SGHMC};
-    kwargs...,
+    rng::Random.AbstractRNG, ldf::DynamicPPL.LogDensityFunction, spl::SGHMC; kwargs...
 )
     vi = ldf.varinfo
 
@@ -68,7 +65,7 @@ end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     ldf::DynamicPPL.LogDensityFunction,
-    spl::Sampler{<:SGHMC},
+    spl::SGHMC,
     state::SGHMCState;
     kwargs...,
 )
@@ -81,8 +78,8 @@ function AbstractMCMC.step(
     # equation (15) of Chen et al. (2014)
     v = state.velocity
     θ .+= v
-    η = spl.alg.learning_rate
-    α = spl.alg.momentum_decay
+    η = spl.learning_rate
+    α = spl.momentum_decay
     newv = (1 - α) .* v .+ η .* grad .+ sqrt(2 * η * α) .* randn(rng, eltype(v), length(v))
 
     # Save new variables and recompute log density.
@@ -201,31 +198,24 @@ struct SGLDState{V<:AbstractVarInfo}
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    ldf::DynamicPPL.LogDensityFunction,
-    spl::Sampler{<:SGLD};
-    kwargs...,
+    rng::Random.AbstractRNG, ldf::DynamicPPL.LogDensityFunction, spl::SGLD; kwargs...
 )
     # Create first sample and state.
     vi = ldf.varinfo
-    sample = SGLDTransition(ldf.model, vi, zero(spl.alg.stepsize(0)))
+    sample = SGLDTransition(ldf.model, vi, zero(spl.stepsize(0)))
     state = SGLDState(vi, 1)
     return sample, state
 end
 
 function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    ldf::LogDensityFunction,
-    spl::Sampler{<:SGLD},
-    state::SGLDState;
-    kwargs...,
+    rng::Random.AbstractRNG, ldf::LogDensityFunction, spl::SGLD, state::SGLDState; kwargs...
 )
     # Perform gradient step.
     vi = state.vi
     θ = vi[:]
     grad = last(LogDensityProblems.logdensity_and_gradient(ldf, θ))
     step = state.step
-    stepsize = spl.alg.stepsize(step)
+    stepsize = spl.stepsize(step)
     θ .+= (stepsize / 2) .* grad .+ sqrt(stepsize) .* randn(rng, eltype(θ), length(θ))
 
     # Save new variables and recompute log density.

test/ext/dynamichmc.jl

@@ -14,10 +14,6 @@ using Turing
 @testset "TuringDynamicHMCExt" begin
     spl = externalsampler(DynamicHMC.NUTS())
 
-    @testset "alg_str" begin
-        @test DynamicPPL.alg_str(Sampler(spl)) == "DynamicNUTS"
-    end
-
     @testset "sample() interface" begin
         @model function demo_normal(x)
             a ~ Normal()

test/mcmc/repeat_sampler.jl

@@ -14,7 +14,7 @@ using Turing
     num_chains = 2
 
     rng = StableRNG(0)
-    for sampler in [MH(), DynamicPPL.Sampler(HMC(0.01, 4))]
+    for sampler in [MH(), HMC(0.01, 4)]
         model_or_ldf = if sampler isa MH
             gdemo_default
         else

test/mcmc/sghmc.jl

@@ -79,13 +79,6 @@ end
     @testset "sghmc constructor" begin
         alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1)
         @test alg isa SGHMC
-        sampler = DynamicPPL.Sampler(alg)
-        @test sampler isa DynamicPPL.Sampler{<:SGHMC}
-
-        alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1)
-        @test alg isa SGHMC
-        sampler = DynamicPPL.Sampler(alg)
-        @test sampler isa DynamicPPL.Sampler{<:SGHMC}
     end
 
     @testset "sghmc inference" begin
@@ -100,13 +93,6 @@ end
     @testset "sgld constructor" begin
         alg = SGLD(; stepsize=PolynomialStepsize(0.25))
         @test alg isa SGLD
-        sampler = DynamicPPL.Sampler(alg)
-        @test sampler isa DynamicPPL.Sampler{<:SGLD}
-
-        alg = SGLD(; stepsize=PolynomialStepsize(0.25))
-        @test alg isa SGLD
-        sampler = DynamicPPL.Sampler(alg)
-        @test sampler isa DynamicPPL.Sampler{<:SGLD}
     end
     @testset "sgld inference" begin
         rng = StableRNG(1)