remove LogDensityProblemsAD (#2490)

* Remove LogDensityProblemsAD, part 1

* update Optimisation code to not use LogDensityProblemsAD

* Fix field name change

* Don't put chunksize=0

* Remove LogDensityProblemsAD dep

* Improve OptimLogDensity docstring

* Remove unneeded model argument to _optimize

* Fix more tests

* Remove essential/ad from the list of CI groups

* Fix HMC function

Author: penelopeysm

.github/workflows/Tests.yml

@@ -24,8 +24,6 @@ jobs:
         test:
           # Run some of the slower test files individually. The last one catches everything
           # not included in the others.
-          - name: "essential/ad"
-            args: "essential/ad.jl"
           - name: "mcmc/gibbs"
             args: "mcmc/gibbs.jl"
           - name: "mcmc/hmc"
@@ -37,7 +35,7 @@ jobs:
           - name: "mcmc/ess"
             args: "mcmc/ess.jl"
           - name: "everything else"
-            args: "--skip essential/ad.jl mcmc/gibbs.jl mcmc/hmc.jl mcmc/abstractmcmc.jl mcmc/Inference.jl mcmc/ess.jl"
+            args: "--skip mcmc/gibbs.jl mcmc/hmc.jl mcmc/abstractmcmc.jl mcmc/Inference.jl mcmc/ess.jl"
         runner:
           # Default
           - version: '1'

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.11.3"
 manifest_format = "2.0"
-project_hash = "83ec9face19bc568fc30cc287161517dc49f6c5c"
+project_hash = "afdf28a30966aaa4af542a30879dd92074661565"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "fb97701c117c8162e84dfcf80215caa904aef44f"

Project.toml

@@ -23,7 +23,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
-LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 NamedArrays = "86f7a689-2022-50b4-a561-43c23ac3c673"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -69,7 +68,6 @@ ForwardDiff = "0.10.3"
 Libtask = "0.8.8"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
-LogDensityProblemsAD = "1.7.0"
 MCMCChains = "5, 6"
 NamedArrays = "0.9, 0.10"
 Optim = "1"

ext/TuringDynamicHMCExt.jl

@@ -3,17 +3,10 @@ module TuringDynamicHMCExt
 ### 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 @@ function DynamicPPL.initialstep(
     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 @@ function Optim.optimize(
 )
     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 @@ function Optim.optimize(
 )
     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(
@@ -81,7 +75,7 @@ end
 
 function _mle_optimize(model::DynamicPPL.Model, args...; kwargs...)
     ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    return _optimize(model, Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
 end
 
 """
@@ -112,7 +106,7 @@ function Optim.optimize(
 )
     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 @@ function Optim.optimize(
 )
     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(
@@ -151,28 +145,29 @@ end
 
 function _map_optimize(model::DynamicPPL.Model, args...; kwargs...)
     ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    return _optimize(model, Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
 end
-
 """
-    _optimize(model::Model, f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
+    _optimize(f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
 
 Estimate a mode, i.e., compute a MLE or MAP estimate.
 """
 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 +181,16 @@ function _optimize(
     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(f.ldf.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 @@ function transition_to_turing(f::DynamicPPL.LogDensityFunction, transition)
     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 @@ getstats(transition::AdvancedHMC.Transition) = transition.stat
 
 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 @@ function AbstractMCMC.step(
     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 @@ function AbstractMCMC.step(
             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 @@ function AbstractMCMC.step(
     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/gibbs.jl

@@ -438,7 +438,7 @@ function setparams_varinfo!!(
     state::TuringState,
     params::AbstractVarInfo,
 )
-    logdensity = DynamicPPL.setmodel(state.logdensity, model, sampler.alg.adtype)
+    logdensity = DynamicPPL.setmodel(state.ldf, model, sampler.alg.adtype)
     new_inner_state = setparams_varinfo!!(
         AbstractMCMC.LogDensityModel(logdensity), sampler, state.state, params
     )