Use FastLDF and removal of NodeTrait

Author: penelopeysm

Project.toml

@@ -64,7 +64,7 @@ Distributions = "0.25.77"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.38"
+DynamicPPL = "0.39"
 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 = "py/ldf"}

src/mcmc/gibbs.jl

@@ -21,10 +21,9 @@ isgibbscomponent(::AdvancedHMC.AbstractHMCSampler) = true
 isgibbscomponent(::AdvancedMH.MetropolisHastings) = true
 isgibbscomponent(spl) = false
 
-function can_be_wrapped(ctx::DynamicPPL.AbstractContext)
-    return DynamicPPL.NodeTrait(ctx) isa DynamicPPL.IsLeaf
-end
-can_be_wrapped(ctx::DynamicPPL.PrefixContext) = can_be_wrapped(ctx.context)
+can_be_wrapped(::DynamicPPL.AbstractContext) = true
+can_be_wrapped(::DynamicPPL.AbstractParentContext) = false
+can_be_wrapped(ctx::DynamicPPL.PrefixContext) = can_be_wrapped(DynamicPPL.childcontext(ctx))
 
 # Basically like a `DynamicPPL.FixedContext` but
 # 1. Hijacks the tilde pipeline to fix variables.
@@ -55,7 +54,7 @@ $(FIELDS)
 """
 struct GibbsContext{
     VNs<:Tuple{Vararg{VarName}},GVI<:Ref{<:AbstractVarInfo},Ctx<:DynamicPPL.AbstractContext
-} <: DynamicPPL.AbstractContext
+} <: DynamicPPL.AbstractParentContext
     """
     the VarNames being sampled
     """
@@ -86,7 +85,6 @@ function GibbsContext(target_varnames, global_varinfo)
     return GibbsContext(target_varnames, global_varinfo, DynamicPPL.DefaultContext())
 end
 
-DynamicPPL.NodeTrait(::GibbsContext) = DynamicPPL.IsParent()
 DynamicPPL.childcontext(context::GibbsContext) = context.context
 function DynamicPPL.setchildcontext(context::GibbsContext, childcontext)
     return GibbsContext(

src/mcmc/hmc.jl

@@ -196,7 +196,7 @@ function Turing.Inference.initialstep(
     # Create a Hamiltonian.
     metricT = getmetricT(spl)
     metric = metricT(length(theta))
-    ldf = DynamicPPL.LogDensityFunction(
+    ldf = DynamicPPL.Experimental.FastLDF(
         model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype
     )
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)
@@ -278,7 +278,7 @@ end
 
 function get_hamiltonian(model, spl, vi, state, n)
     metric = gen_metric(n, spl, state)
-    ldf = DynamicPPL.LogDensityFunction(
+    ldf = DynamicPPL.Experimental.FastLDF(
         model, DynamicPPL.getlogjoint_internal, vi; adtype=spl.adtype
     )
     lp_func = Base.Fix1(LogDensityProblems.logdensity, ldf)

src/mcmc/is.jl

@@ -49,7 +49,6 @@ end
 struct ISContext{R<:AbstractRNG} <: DynamicPPL.AbstractContext
     rng::R
 end
-DynamicPPL.NodeTrait(::ISContext) = DynamicPPL.IsLeaf()
 
 function DynamicPPL.tilde_assume!!(
     ctx::ISContext, dist::Distribution, vn::VarName, vi::AbstractVarInfo

src/mcmc/mh.jl

@@ -410,7 +410,6 @@ end
 struct MHContext{R<:AbstractRNG} <: DynamicPPL.AbstractContext
     rng::R
 end
-DynamicPPL.NodeTrait(::MHContext) = DynamicPPL.IsLeaf()
 
 function DynamicPPL.tilde_assume!!(
     context::MHContext, right::Distribution, vn::VarName, vi::AbstractVarInfo

src/mcmc/particle_mcmc.jl

@@ -7,7 +7,6 @@
 struct ParticleMCMCContext{R<:AbstractRNG} <: DynamicPPL.AbstractContext
     rng::R
 end
-DynamicPPL.NodeTrait(::ParticleMCMCContext) = DynamicPPL.IsLeaf()
 
 struct TracedModel{V<:AbstractVarInfo,M<:Model,E<:Tuple} <: AdvancedPS.AbstractGenericModel
     model::M

src/mcmc/prior.jl

@@ -5,21 +5,30 @@ Algorithm for sampling from the prior.
 """
 struct Prior <: AbstractSampler end
 
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG, model::DynamicPPL.Model, sampler::Prior; kwargs...
+)
+    accs = DynamicPPL.AccumulatorTuple((
+        DynamicPPL.ValuesAsInModelAccumulator(true),
+        DynamicPPL.LogPriorAccumulator(),
+        DynamicPPL.LogLikelihoodAccumulator(),
+    ))
+    sampling_model = DynamicPPL.setleafcontext(
+        model, DynamicPPL.InitContext(rng, DynamicPPL.InitFromPrior())
+    )
+    vi = DynamicPPL.OnlyAccsVarInfo(accs)
+    _, vi = DynamicPPL.evaluate!!(sampling_model, vi)
+    return Transition(sampling_model, vi, nothing; reevaluate=false), (sampling_model, vi)
+end
+
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DynamicPPL.Model,
     sampler::Prior,
-    state=nothing;
+    state::Tuple{DynamicPPL.Model,DynamicPPL.Experimental.OnlyAccsVarInfo};
     kwargs...,
 )
-    vi = DynamicPPL.setaccs!!(
-        DynamicPPL.VarInfo(),
-        (
-            DynamicPPL.ValuesAsInModelAccumulator(true),
-            DynamicPPL.LogPriorAccumulator(),
-            DynamicPPL.LogLikelihoodAccumulator(),
-        ),
-    )
-    _, vi = DynamicPPL.init!!(model, vi, DynamicPPL.InitFromPrior())
-    return Transition(model, vi, nothing; reevaluate=false), nothing
+    model, vi = state
+    _, vi = DynamicPPL.evaluate!!(model, vi)
+    return Transition(model, vi, nothing; reevaluate=false), (model, vi)
 end

test/Project.toml

@@ -53,7 +53,7 @@ Combinatorics = "1"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
-DynamicPPL = "0.38"
+DynamicPPL = "0.39"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10, 1"
 HypothesisTests = "0.11"
@@ -77,3 +77,6 @@ StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
 TimerOutputs = "0.5"
 julia = "1.10"
+
+[sources]
+DynamicPPL = {url = "https://github.com/TuringLang/DynamicPPL.jl", rev = "py/ldf"}

test/ad.jl

@@ -94,7 +94,7 @@ encountered.
 
 """
 struct ADTypeCheckContext{ADType,ChildContext<:DynamicPPL.AbstractContext} <:
-       DynamicPPL.AbstractContext
+       DynamicPPL.AbstractParentContext
     child::ChildContext
 
     function ADTypeCheckContext(adbackend, child)
@@ -108,7 +108,6 @@ end
 
 adtype(_::ADTypeCheckContext{ADType}) where {ADType} = ADType
 
-DynamicPPL.NodeTrait(::ADTypeCheckContext) = DynamicPPL.IsParent()
 DynamicPPL.childcontext(c::ADTypeCheckContext) = c.child
 function DynamicPPL.setchildcontext(c::ADTypeCheckContext, child)
     return ADTypeCheckContext(adtype(c), child)

test/mcmc/gibbs.jl

@@ -159,15 +159,15 @@ end
     # It is modified by the capture_targets_and_algs function.
     targets_and_algs = Any[]
 
-    function capture_targets_and_algs(sampler, context)
-        if DynamicPPL.NodeTrait(context) == DynamicPPL.IsLeaf()
-            return nothing
-        end
+    function capture_targets_and_algs(sampler, context::DynamicPPL.AbstractParentContext)
         if context isa Inference.GibbsContext
             push!(targets_and_algs, (context.target_varnames, sampler))
         end
         return capture_targets_and_algs(sampler, DynamicPPL.childcontext(context))
     end
+    function capture_targets_and_algs(sampler, ::DynamicPPL.AbstractContext)
+        return nothing  # Leaf context.
+    end
 
     # The methods that capture testing information for us.
     function AbstractMCMC.step(