remove LogDensityProblemsAD

ext/TuringDynamicHMCExt.jl

@@ -3,17 +3,10 @@
 ### DynamicHMC backend - https://github.com/tpapp/DynamicHMC.jl
 ###
 
-if isdefined(Base, :get_extension)
-    using DynamicHMC: DynamicHMC
-    using Turing
-    using Turing: AbstractMCMC, Random, LogDensityProblems, DynamicPPL
-    using Turing.Inference: ADTypes, LogDensityProblemsAD, TYPEDFIELDS
-else
-    import ..DynamicHMC
-    using ..Turing
-    using ..Turing: AbstractMCMC, Random, LogDensityProblems, DynamicPPL
-    using ..Turing.Inference: ADTypes, LogDensityProblemsAD, TYPEDFIELDS
-end
+using DynamicHMC: DynamicHMC
+using Turing
+using Turing: AbstractMCMC, Random, LogDensityProblems, DynamicPPL
+using Turing.Inference: ADTypes, TYPEDFIELDS
 
 """
     DynamicNUTS
@@ -69,10 +62,11 @@
     end
 
     # Define log-density function.
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(
-            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
-        ),
+    ℓ = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
+        adtype=spl.alg.adtype,
     )
 
     # Perform initial step.

ext/TuringOptimExt.jl

@@ -1,14 +1,8 @@
 module TuringOptimExt
 
-if isdefined(Base, :get_extension)
-    using Turing: Turing
-    import Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
-    using Optim: Optim
-else
-    import ..Turing
-    import ..Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
-    import ..Optim
-end
+using Turing: Turing
+import Turing: DynamicPPL, NamedArrays, Accessors, Optimisation
+using Optim: Optim
 
 ####################
 # Optim.jl methods #
@@ -42,7 +36,7 @@
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -65,7 +59,7 @@
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
 function Optim.optimize(
@@ -112,7 +106,7 @@
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -135,7 +129,7 @@
 )
     ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
     f = Optimisation.OptimLogDensity(model, ctx)
-    init_vals = DynamicPPL.getparams(f)
+    init_vals = DynamicPPL.getparams(f.ldf)
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
 function Optim.optimize(
@@ -162,17 +156,20 @@
 function _optimize(
     model::DynamicPPL.Model,
     f::Optimisation.OptimLogDensity,
-    init_vals::AbstractArray=DynamicPPL.getparams(f),
+    init_vals::AbstractArray=DynamicPPL.getparams(f.ldf),
     optimizer::Optim.AbstractOptimizer=Optim.LBFGS(),
     options::Optim.Options=Optim.Options(),
     args...;
     kwargs...,
 )
     # Convert the initial values, since it is assumed that users provide them
     # in the constrained space.
-    f = Accessors.@set f.varinfo = DynamicPPL.unflatten(f.varinfo, init_vals)
-    f = Accessors.@set f.varinfo = DynamicPPL.link(f.varinfo, model)
-    init_vals = DynamicPPL.getparams(f)
+    # TODO(penelopeysm): As with in src/optimisation/Optimisation.jl, unclear
+    # whether initialisation is really necessary at all
+    vi = DynamicPPL.unflatten(f.ldf.varinfo, init_vals)
+    vi = DynamicPPL.link(vi, f.ldf.model)
+    f = Optimisation.OptimLogDensity(f.ldf.model, vi, f.ldf.context; adtype=f.ldf.adtype)
+    init_vals = DynamicPPL.getparams(f.ldf)
 
     # Optimize!
     M = Optim.optimize(Optim.only_fg!(f), init_vals, optimizer, options, args...; kwargs...)
@@ -186,12 +183,16 @@
     end
 
     # Get the optimum in unconstrained space. `getparams` does the invlinking.
-    f = Accessors.@set f.varinfo = DynamicPPL.unflatten(f.varinfo, M.minimizer)
-    vns_vals_iter = Turing.Inference.getparams(model, f.varinfo)
+    vi = f.ldf.varinfo
+    vi_optimum = DynamicPPL.unflatten(vi, M.minimizer)
+    logdensity_optimum = Optimisation.OptimLogDensity(
+        f.ldf.model, vi_optimum, f.ldf.context
+    )
+    vns_vals_iter = Turing.Inference.getparams(model, vi_optimum)
     varnames = map(Symbol ∘ first, vns_vals_iter)
     vals = map(last, vns_vals_iter)
     vmat = NamedArrays.NamedArray(vals, varnames)
-    return Optimisation.ModeResult(vmat, M, -M.minimum, f)
+    return Optimisation.ModeResult(vmat, M, -M.minimum, logdensity_optimum)
 end
 
 end # module

src/mcmc/Inference.jl

@@ -48,7 +48,6 @@ import AdvancedPS
 import Accessors
 import EllipticalSliceSampling
 import LogDensityProblems
-import LogDensityProblemsAD
 import Random
 import MCMCChains
 import StatsBase: predict
@@ -160,29 +159,6 @@ function externalsampler(
     return ExternalSampler(sampler, adtype, Val(unconstrained))
 end
 
-getADType(spl::Sampler) = getADType(spl.alg)
-getADType(::SampleFromPrior) = Turing.DEFAULT_ADTYPE
-
-getADType(ctx::DynamicPPL.SamplingContext) = getADType(ctx.sampler)
-getADType(ctx::DynamicPPL.AbstractContext) = getADType(DynamicPPL.NodeTrait(ctx), ctx)
-getADType(::DynamicPPL.IsLeaf, ctx::DynamicPPL.AbstractContext) = Turing.DEFAULT_ADTYPE
-function getADType(::DynamicPPL.IsParent, ctx::DynamicPPL.AbstractContext)
-    return getADType(DynamicPPL.childcontext(ctx))
-end
-
-getADType(alg::Hamiltonian) = alg.adtype
-
-function LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)
-    return LogDensityProblemsAD.ADgradient(getADType(ℓ.context), ℓ)
-end
-
-function LogDensityProblems.logdensity(
-    f::Turing.LogDensityFunction{<:AbstractVarInfo,<:Model,<:DynamicPPL.DefaultContext},
-    x::NamedTuple,
-)
-    return DynamicPPL.logjoint(f.model, DynamicPPL.unflatten(f.varinfo, x))
-end
-
 # TODO: make a nicer `set_namedtuple!` and move these functions to DynamicPPL.
 function DynamicPPL.unflatten(vi::TypedVarInfo, θ::NamedTuple)
     set_namedtuple!(deepcopy(vi), θ)

src/mcmc/abstractmcmc.jl

@@ -1,6 +1,6 @@
-struct TuringState{S,F}
+struct TuringState{S,M,V,C}
     state::S
-    logdensity::F
+    ldf::DynamicPPL.LogDensityFunction{M,V,C}
 end
 
 state_to_turing(f::DynamicPPL.LogDensityFunction, state) = TuringState(state, f)
@@ -12,20 +12,10 @@
     return Transition(f.model, varinfo, transition)
 end
 
-state_to_turing(f::LogDensityProblemsAD.ADGradientWrapper, state) = TuringState(state, f)
-function transition_to_turing(f::LogDensityProblemsAD.ADGradientWrapper, transition)
-    return transition_to_turing(parent(f), transition)
-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(DynamicPPL.getmodel(state.logdensity), state.state)
+    θ = getparams(state.ldf.model, state.state)
     # TODO: Do we need to link here first?
-    return DynamicPPL.unflatten(varinfo_from_logdensityfn(state.logdensity), θ)
+    return DynamicPPL.unflatten(state.ldf.varinfo, θ)
 end
 varinfo(state::AbstractVarInfo) = state
 
@@ -40,23 +30,6 @@
 
 getparams(::DynamicPPL.Model, transition::AdvancedMH.Transition) = transition.params
 
-getvarinfo(f::DynamicPPL.LogDensityFunction) = f.varinfo
-function getvarinfo(f::LogDensityProblemsAD.ADGradientWrapper)
-    return getvarinfo(LogDensityProblemsAD.parent(f))
-end
-
-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
-)
-    return LogDensityProblemsAD.ADgradient(
-        adtype, setvarinfo(LogDensityProblemsAD.parent(f), varinfo)
-    )
-end
-
 # TODO: Do we also support `resume`, etc?
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
@@ -69,12 +42,8 @@
     alg = sampler_wrapper.alg
     sampler = alg.sampler
 
-    # Create a log-density function with an implementation of the
-    # gradient so we ensure that we're using the same AD backend as in Turing.
-    f = LogDensityProblemsAD.ADgradient(alg.adtype, DynamicPPL.LogDensityFunction(model))
-
-    # Link the varinfo if needed.
-    varinfo = getvarinfo(f)
+    # Initialise varinfo with initial params and link the varinfo if needed.
+    varinfo = DynamicPPL.VarInfo(model)
     if requires_unconstrained_space(alg)
         if initial_params !== nothing
             # If we have initial parameters, we need to set the varinfo before linking.
@@ -85,9 +54,11 @@
             varinfo = DynamicPPL.link(varinfo, model)
         end
     end
-    f = setvarinfo(f, varinfo, alg.adtype)
 
-    # Then just call `AdvancedHMC.step` with the right arguments.
+    # Construct LogDensityFunction
+    f = DynamicPPL.LogDensityFunction(model, varinfo; adtype=alg.adtype)
+
+    # Then just call `AbstractMCMC.step` with the right arguments.
     if initial_state === nothing
         transition_inner, state_inner = AbstractMCMC.step(
             rng, AbstractMCMC.LogDensityModel(f), sampler; initial_params, kwargs...
@@ -114,7 +85,7 @@
     kwargs...,
 )
     sampler = sampler_wrapper.alg.sampler
-    f = state.logdensity
+    f = state.ldf
 
     # Then just call `AdvancedHMC.step` with the right arguments.
     transition_inner, state_inner = AbstractMCMC.step(

src/mcmc/hmc.jl

@@ -156,19 +156,19 @@
     # Create a Hamiltonian.
     metricT = getmetricT(spl.alg)
     metric = metricT(length(theta))
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(
-            model,
-            vi,
-            # Use the leaf-context from the `model` in case the user has
-            # contextualized the model with something like `PriorContext`
-            # to sample from the prior.
-            DynamicPPL.SamplingContext(rng, spl, DynamicPPL.leafcontext(model.context)),
-        ),
+    ldf = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        # TODO(penelopeysm): Can we just use leafcontext(model.context)? Do we
+        # need to pass in the sampler? (In fact LogDensityFunction defaults to
+        # using leafcontext(model.context) so could we just remove the argument
+        # entirely?)
+        DynamicPPL.SamplingContext(rng, spl, DynamicPPL.leafcontext(model.context));
+        adtype=spl.alg.adtype,
     )
-    logπ = Base.Fix1(LogDensityProblems.logdensity, ℓ)
-    ∂logπ∂θ(x) = LogDensityProblems.logdensity_and_gradient(ℓ, x)
-    hamiltonian = AHMC.Hamiltonian(metric, logπ, ∂logπ∂θ)
+    lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
+    lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
+    hamiltonian = AHMC.Hamiltonian(metric, lp_func, lp_grad_func)
 
     # Compute phase point z.
     z = AHMC.phasepoint(rng, theta, hamiltonian)
@@ -287,16 +287,19 @@
 
 function get_hamiltonian(model, spl, vi, state, n)
     metric = gen_metric(n, spl, state)
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(
-            model,
-            vi,
-            DynamicPPL.SamplingContext(spl, DynamicPPL.leafcontext(model.context)),
-        ),
+    ldf = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        # TODO(penelopeysm): Can we just use leafcontext(model.context)? Do we
+        # need to pass in the sampler? (In fact LogDensityFunction defaults to
+        # using leafcontext(model.context) so could we just remove the argument
+        # entirely?)
+        DynamicPPL.SamplingContext(spl, DynamicPPL.leafcontext(model.context));
+        adtype=spl.alg.adtype,
     )
-    ℓπ = Base.Fix1(LogDensityProblems.logdensity, ℓ)
-    ∂ℓπ∂θ = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ℓ)
-    return AHMC.Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
+    lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
+    lp_grad_func = Base.Fix1(LogDensityProblems.logdensity_and_gradient, ldf)
+    return AHMC.Hamiltonian(metric, lp_func, lp_grad_func)
 end
 
 """

src/mcmc/sghmc.jl

@@ -66,10 +66,11 @@
 
     # Compute initial sample and state.
     sample = Transition(model, vi)
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(
-            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
-        ),
+    ℓ = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
+        adtype=spl.alg.adtype,
     )
     state = SGHMCState(ℓ, vi, zero(vi[spl]))
 
@@ -228,10 +229,11 @@
 
     # Create first sample and state.
     sample = SGLDTransition(model, vi, zero(spl.alg.stepsize(0)))
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(
-            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
-        ),
+    ℓ = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
+        adtype=spl.alg.adtype,
     )
     state = SGLDState(ℓ, vi, 1)