Support for DynamicPPL v0.35 (#2488)

* Progress towards compat with DPPL v0.35

* More fixing of DPPL v0.35 stuff

* Fix LogDensityFunction argument order

* More minor bugfixes

* [TEMP] Commit Manifest pointing to DynamicPPL#release-0.35

* 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

* more test fixes (#2491)

* 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

* More test fixes

* Remove Manifest

* More fixes for DynamicPPL 0.35 (#2494)

* Remove test/dynamicppl/compiler.jl

* Remove old regression tests

* Remove vdemo2

* Fix last test

* Add HISTORY.md entry about DPPL 0.35

* Allow ESS to sample variables with different symbols

* Update a TODO note

---------

Co-authored-by: Penelope Yong <[email protected]>

Author: mhauru

.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'

HISTORY.md

@@ -2,8 +2,21 @@
 
 ## Breaking changes
 
+### Gibbs constructors
+
 0.37 removes the old Gibbs constructors deprecated in 0.36.
 
+### DynamicPPL 0.35
+
+Turing.jl v0.37 uses DynamicPPL v0.35, which brings with it several breaking changes:
+
+  - The right hand side of `.~` must from now on be a univariate distribution.
+  - Indexing `VarInfo` objects by samplers has been removed completely.
+  - The order in which nested submodel prefixes are applied has been reversed.
+  - The arguments for the constructor of `LogDensityFunction` have changed. `LogDensityFunction` also now satisfies the `LogDensityProblems` interface, without needing a wrapper object.
+
+For more details about all of the above, see the changelog of DynamicPPL [here](https://github.com/TuringLang/DynamicPPL.jl/releases/tag/v0.35.0).
+
 # Release 0.36.0
 
 ## Breaking changes

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"
@@ -63,13 +62,12 @@ Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.34.1"
+DynamicPPL = "0.35"
 EllipticalSliceSampling = "0.5, 1, 2"
 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
@@ -25,22 +18,15 @@ To use it, make sure you have DynamicHMC package (version >= 2) loaded:
 using DynamicHMC
 ```
 """
-struct DynamicNUTS{AD,space,T<:DynamicHMC.NUTS} <: Turing.Inference.Hamiltonian
+struct DynamicNUTS{AD,T<:DynamicHMC.NUTS} <: Turing.Inference.Hamiltonian
     sampler::T
     adtype::AD
 end
 
-function DynamicNUTS(
-    spl::DynamicHMC.NUTS=DynamicHMC.NUTS(),
-    space::Tuple=();
-    adtype::ADTypes.AbstractADType=Turing.DEFAULT_ADTYPE,
-)
-    return DynamicNUTS{typeof(adtype),space,typeof(spl)}(spl, adtype)
-end
+DynamicNUTS() = DynamicNUTS(DynamicHMC.NUTS())
+DynamicNUTS(spl) = DynamicNUTS(spl, Turing.DEFAULT_ADTYPE)
 Turing.externalsampler(spl::DynamicHMC.NUTS) = DynamicNUTS(spl)
 
-DynamicPPL.getspace(::DynamicNUTS{<:Any,space}) where {space} = space
-
 """
     DynamicNUTSState
 
@@ -70,25 +56,28 @@ function DynamicPPL.initialstep(
     kwargs...,
 )
     # Ensure that initial sample is in unconstrained space.
-    if !DynamicPPL.islinked(vi, spl)
-        vi = DynamicPPL.link!!(vi, spl, model)
+    if !DynamicPPL.islinked(vi)
+        vi = DynamicPPL.link!!(vi, model)
         vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
     end
 
     # Define log-density function.
-    ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
+    ℓ = DynamicPPL.LogDensityFunction(
+        model,
+        vi,
+        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
+        adtype=spl.alg.adtype,
     )
 
     # Perform initial step.
     results = DynamicHMC.mcmc_keep_warmup(
-        rng, ℓ, 0; initialization=(q=vi[spl],), reporter=DynamicHMC.NoProgressReport()
+        rng, ℓ, 0; initialization=(q=vi[:],), reporter=DynamicHMC.NoProgressReport()
     )
     steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
     Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)
 
     # Update the variables.
-    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.unflatten(vi, Q.q)
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create first sample and state.
@@ -112,7 +101,7 @@ function AbstractMCMC.step(
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
     # Update the variables.
-    vi = DynamicPPL.setindex!!(vi, Q.q, spl)
+    vi = DynamicPPL.unflatten(vi, Q.q)
     vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
 
     # Create next sample and state.

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/Turing.jl

@@ -9,7 +9,7 @@ using Compat: pkgversion
 
 using AdvancedVI: AdvancedVI
 using DynamicPPL: DynamicPPL, LogDensityFunction
-import DynamicPPL: getspace, NoDist, NamedDist
+import DynamicPPL: NoDist, NamedDist
 using LogDensityProblems: LogDensityProblems
 using NamedArrays: NamedArrays
 using Accessors: Accessors

src/essential/container.jl

@@ -39,7 +39,7 @@ function AdvancedPS.advance!(
 end
 
 function AdvancedPS.delete_retained!(trace::TracedModel)
-    DynamicPPL.set_retained_vns_del_by_spl!(trace.varinfo, trace.sampler)
+    DynamicPPL.set_retained_vns_del!(trace.varinfo)
     return trace
 end