Use typed varinfo in Prior

HISTORY.md

@@ -1,3 +1,12 @@
+# 0.40.0
+
+TODO
+
+  - DynamicPPL 0.37 stuff
+
+  - pMCMC and Gibbs?
+  - Prior is faster
+
 # 0.39.9
 
 Revert a bug introduced in 0.39.5 in the external sampler interface.

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.39.9"
+version = "0.40.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -64,7 +64,7 @@ Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.36.3"
+DynamicPPL = "0.37"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3, 1"
 Libtask = "0.9.3"
@@ -90,3 +90,6 @@ julia = "1.10.8"
 [extras]
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+
+[sources]
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "main"}

ext/TuringDynamicHMCExt.jl

@@ -58,15 +58,12 @@ function DynamicPPL.initialstep(
     # Ensure that initial sample is in unconstrained space.
     if !DynamicPPL.islinked(vi)
         vi = DynamicPPL.link!!(vi, model)
-        vi = last(DynamicPPL.evaluate!!(model, vi, DynamicPPL.SamplingContext(rng, spl)))
+        vi = last(DynamicPPL.evaluate!!(model, vi))
     end
 
     # Define log-density function.
     ℓ = DynamicPPL.LogDensityFunction(
-        model,
-        vi,
-        DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext());
-        adtype=spl.alg.adtype,
+        model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.alg.adtype
     )
 
     # Perform initial step.
@@ -76,12 +73,9 @@ function DynamicPPL.initialstep(
     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.unflatten(vi, Q.q)
-    vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
-
     # Create first sample and state.
-    sample = Turing.Inference.Transition(model, vi)
+    vi = DynamicPPL.unflatten(vi, Q.q)
+    sample = Turing.Inference.Transition(model, vi, nothing)
     state = DynamicNUTSState(ℓ, vi, Q, steps.H.κ, steps.ϵ)
 
     return sample, state
@@ -100,12 +94,9 @@ function AbstractMCMC.step(
     steps = DynamicHMC.mcmc_steps(rng, spl.alg.sampler, state.metric, ℓ, state.stepsize)
     Q, _ = DynamicHMC.mcmc_next_step(steps, state.cache)
 
-    # Update the variables.
-    vi = DynamicPPL.unflatten(vi, Q.q)
-    vi = DynamicPPL.setlogp!!(vi, Q.ℓq)
-
     # Create next sample and state.
-    sample = Turing.Inference.Transition(model, vi)
+    vi = DynamicPPL.unflatten(vi, Q.q)
+    sample = Turing.Inference.Transition(model, vi, nothing)
     newstate = DynamicNUTSState(ℓ, vi, Q, state.metric, state.stepsize)
 
     return sample, newstate

ext/TuringOptimExt.jl

@@ -34,8 +34,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -57,8 +56,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _mle_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -74,8 +72,8 @@ function Optim.optimize(
 end
 
 function _mle_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    ctx = Optimisation.OptimizationContext(DynamicPPL.LikelihoodContext())
-    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getloglikelihood)
+    return _optimize(f, args...; kwargs...)
 end
 
 """
@@ -104,8 +102,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getlogjoint)
     init_vals = DynamicPPL.getparams(f.ldf)
     optimizer = Optim.LBFGS()
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
@@ -127,8 +124,7 @@ function Optim.optimize(
     options::Optim.Options=Optim.Options();
     kwargs...,
 )
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    f = Optimisation.OptimLogDensity(model, ctx)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getlogjoint)
     init_vals = DynamicPPL.getparams(f.ldf)
     return _map_optimize(model, init_vals, optimizer, options; kwargs...)
 end
@@ -144,9 +140,10 @@ function Optim.optimize(
 end
 
 function _map_optimize(model::DynamicPPL.Model, args...; kwargs...)
-    ctx = Optimisation.OptimizationContext(DynamicPPL.DefaultContext())
-    return _optimize(Optimisation.OptimLogDensity(model, ctx), args...; kwargs...)
+    f = Optimisation.OptimLogDensity(model, DynamicPPL.getlogjoint)
+    return _optimize(f, args...; kwargs...)
 end
+
 """
     _optimize(f::OptimLogDensity, optimizer=Optim.LBFGS(), args...; kwargs...)
 
@@ -166,7 +163,9 @@ function _optimize(
     # 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)
+    f = Optimisation.OptimLogDensity(
+        f.ldf.model, f.ldf.getlogdensity, vi; adtype=f.ldf.adtype
+    )
     init_vals = DynamicPPL.getparams(f.ldf)
 
     # Optimize!
@@ -184,7 +183,7 @@ function _optimize(
     vi = f.ldf.varinfo
     vi_optimum = DynamicPPL.unflatten(vi, M.minimizer)
     logdensity_optimum = Optimisation.OptimLogDensity(
-        f.ldf.model, vi_optimum, f.ldf.context
+        f.ldf.model, f.ldf.getlogdensity, vi_optimum; adtype=f.ldf.adtype
     )
     vals_dict = Turing.Inference.getparams(f.ldf.model, vi_optimum)
     iters = map(DynamicPPL.varname_and_value_leaves, keys(vals_dict), values(vals_dict))

src/Turing.jl

@@ -71,7 +71,6 @@ using DynamicPPL:
     unfix,
     prefix,
     conditioned,
-    @submodel,
     to_submodel,
     LogDensityFunction,
     @addlogprob!
@@ -81,7 +80,6 @@ using OrderedCollections: OrderedDict
 # Turing essentials - modelling macros and inference algorithms
 export
     # DEPRECATED
-    @submodel,
     generated_quantities,
     # Modelling - AbstractPPL and DynamicPPL
     @model,