More minor bugfixes

Author: mhauru

src/mcmc/Inference.jl

@@ -173,7 +173,7 @@ end
 getADType(alg::Hamiltonian) = alg.adtype
 
 function LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)
-    return LogDensityProblemsAD.ADgradient(getADType(DynamicPPL.getcontext(ℓ)), ℓ)
+    return LogDensityProblemsAD.ADgradient(getADType(ℓ.context), ℓ)
 end
 
 function LogDensityProblems.logdensity(

src/mcmc/abstractmcmc.jl

@@ -45,7 +45,10 @@ function getvarinfo(f::LogDensityProblemsAD.ADGradientWrapper)
     return getvarinfo(LogDensityProblemsAD.parent(f))
 end
 
-setvarinfo(f::DynamicPPL.LogDensityFunction, varinfo) = Accessors.@set f.varinfo = varinfo
+function setvarinfo(f::DynamicPPL.LogDensityFunction, varinfo)
+    return DynamicPPL.LogDensityFunction(f.model, varinfo, f.context; adtype=f.adtype)
+end
+
 function setvarinfo(
     f::LogDensityProblemsAD.ADGradientWrapper, varinfo, adtype::ADTypes.AbstractADType
 )

src/mcmc/emcee.jl

@@ -53,7 +53,7 @@ function AbstractMCMC.step(
         length(initial_params) == n ||
             throw(ArgumentError("initial parameters have to be specified for each walker"))
         vis = map(vis, initial_params) do vi, init
-            vi = DynamicPPL.initialize_parameters!!(vi, init, spl, model)
+            vi = DynamicPPL.initialize_parameters!!(vi, init, model)
 
             # Update log joint probability.
             last(DynamicPPL.evaluate!!(model, rng, vi, SampleFromPrior()))

src/mcmc/ess.jl

@@ -119,16 +119,10 @@ end
 Distributions.mean(p::ESSPrior) = p.μ
 
 # Evaluate log-likelihood of proposals
-const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} = Turing.LogDensityFunction{
-    V,M,<:DynamicPPL.SamplingContext{<:S}
-}
-
-function (ℓ::ESSLogLikelihood)(f::AbstractVector)
-    sampler = DynamicPPL.getsampler(ℓ)
-    varinfo = DynamicPPL.unflatten(ℓ.varinfo, f)
-    varinfo = last(DynamicPPL.evaluate!!(ℓ.model, varinfo, sampler))
-    return getlogp(varinfo)
-end
+const ESSLogLikelihood{M<:Model,S<:Sampler{<:ESS},V<:AbstractVarInfo} =
+    Turing.LogDensityFunction{M,V,<:DynamicPPL.SamplingContext{<:S},AD} where {AD}
+
+(ℓ::ESSLogLikelihood)(f::AbstractVector) = LogDensityProblems.logdensity(ℓ, f)
 
 function DynamicPPL.tilde_assume(
     rng::Random.AbstractRNG, ::DefaultContext, ::Sampler{<:ESS}, right, vn, vi
@@ -138,8 +132,6 @@ function DynamicPPL.tilde_assume(
     )
 end
 
-function DynamicPPL.tilde_observe(
-    ctx::DefaultContext, sampler::Sampler{<:ESS}, right, left, vi
-)
+function DynamicPPL.tilde_observe(ctx::DefaultContext, ::Sampler{<:ESS}, right, left, vi)
     return DynamicPPL.tilde_observe(ctx, SampleFromPrior(), right, left, vi)
 end

src/mcmc/is.jl

@@ -46,16 +46,16 @@ function getlogevidence(samples::Vector{<:Transition}, ::Sampler{<:IS}, state)
     return logsumexp(map(x -> x.lp, samples)) - log(length(samples))
 end
 
-function DynamicPPL.assume(rng, spl::Sampler{<:IS}, dist::Distribution, vn::VarName, vi)
+function DynamicPPL.assume(rng, ::Sampler{<:IS}, dist::Distribution, vn::VarName, vi)
     if haskey(vi, vn)
         r = vi[vn]
     else
         r = rand(rng, dist)
-        vi = push!!(vi, vn, r, dist, spl)
+        vi = push!!(vi, vn, r, dist)
     end
     return r, 0, vi
 end
 
-function DynamicPPL.observe(spl::Sampler{<:IS}, dist::Distribution, value, vi)
+function DynamicPPL.observe(::Sampler{<:IS}, dist::Distribution, value, vi)
     return logpdf(dist, value), vi
 end

src/mcmc/mh.jl

@@ -188,14 +188,13 @@ A log density function for the MH sampler.
 
 This variant uses the `set_namedtuple!` function to update the `VarInfo`.
 """
-const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} = Turing.LogDensityFunction{
-    V,M,<:DynamicPPL.SamplingContext{<:S}
-}
+const MHLogDensityFunction{M<:Model,S<:Sampler{<:MH},V<:AbstractVarInfo} =
+    Turing.LogDensityFunction{M,V,<:DynamicPPL.SamplingContext{<:S},AD} where {AD}
 
 function LogDensityProblems.logdensity(f::MHLogDensityFunction, x::NamedTuple)
     vi = deepcopy(f.varinfo)
     set_namedtuple!(vi, x)
-    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, DynamicPPL.getcontext(f)))
+    vi_new = last(DynamicPPL.evaluate!!(f.model, vi, f.context))
     lj = getlogp(vi_new)
     return lj
 end

src/mcmc/particle_mcmc.jl

@@ -365,7 +365,7 @@ function DynamicPPL.assume(
 
     if ~haskey(vi, vn)
         r = rand(trng, dist)
-        push!!(vi, vn, r, dist, spl)
+        push!!(vi, vn, r, dist)
     elseif is_flagged(vi, vn, "del")
         unset_flag!(vi, vn, "del") # Reference particle parent
         r = rand(trng, dist)

src/optimisation/Optimisation.jl

@@ -99,8 +99,8 @@ end
 
 A struct that stores the negative log density function of a `DynamicPPL` model.
 """
-const OptimLogDensity{M<:DynamicPPL.Model,C<:OptimizationContext,V<:DynamicPPL.VarInfo} = Turing.LogDensityFunction{
-    V,M,C
+const OptimLogDensity{M<:DynamicPPL.Model,C<:OptimizationContext,V<:DynamicPPL.VarInfo,AD} = Turing.LogDensityFunction{
+    M,V,C,AD
 }
 
 """
@@ -125,9 +125,7 @@ required by Optimization.jl.
 """
 function (f::OptimLogDensity)(z::AbstractVector)
     varinfo = DynamicPPL.unflatten(f.varinfo, z)
-    return -DynamicPPL.getlogp(
-        last(DynamicPPL.evaluate!!(f.model, varinfo, DynamicPPL.getcontext(f)))
-    )
+    return -DynamicPPL.getlogp(last(DynamicPPL.evaluate!!(f.model, varinfo, f.context)))
 end
 
 (f::OptimLogDensity)(z, _) = f(z)

test/mcmc/abstractmcmc.jl

@@ -26,7 +26,7 @@ function initialize_nuts(model::Turing.Model)
     f = Turing.Inference.setvarinfo(
         f,
         DynamicPPL.link!!(Turing.Inference.getvarinfo(f), model),
-        Turing.Inference.getADType(DynamicPPL.getcontext(LogDensityProblemsAD.parent(f))),
+        Turing.Inference.getADType(LogDensityProblemsAD.parent(f).context),
     )
 
     # Choose parameter dimensionality and initial parameter value