Fix LogDensityFunction argument order

Author: mhauru

ext/TuringDynamicHMCExt.jl

@@ -70,7 +70,9 @@ function DynamicPPL.initialstep(
 
     # Define log-density function.
     ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
+        Turing.LogDensityFunction(
+            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
+        ),
     )
 
     # Perform initial step.

src/mcmc/ess.jl

@@ -59,7 +59,9 @@ function AbstractMCMC.step(
         rng,
         EllipticalSliceSampling.ESSModel(
             ESSPrior(model, spl, vi),
-            Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()),
+            Turing.LogDensityFunction(
+                model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
+            ),
         ),
         EllipticalSliceSampling.ESS(),
         oldstate,

src/mcmc/hmc.jl

@@ -158,8 +158,8 @@ function DynamicPPL.initialstep(
     metric = metricT(length(theta))
     ℓ = LogDensityProblemsAD.ADgradient(
         Turing.LogDensityFunction(
-            vi,
             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.
@@ -289,8 +289,8 @@ function get_hamiltonian(model, spl, vi, state, n)
     metric = gen_metric(n, spl, state)
     ℓ = LogDensityProblemsAD.ADgradient(
         Turing.LogDensityFunction(
-            vi,
             model,
+            vi,
             DynamicPPL.SamplingContext(spl, DynamicPPL.leafcontext(model.context)),
         ),
     )

src/mcmc/mh.jl

@@ -310,8 +310,8 @@ function propose!!(
         Base.Fix1(
             LogDensityProblems.logdensity,
             Turing.LogDensityFunction(
-                vi,
                 model,
+                vi,
                 DynamicPPL.SamplingContext(rng, spl, DynamicPPL.leafcontext(model.context)),
             ),
         ),
@@ -345,8 +345,8 @@ function propose!!(
         Base.Fix1(
             LogDensityProblems.logdensity,
             Turing.LogDensityFunction(
-                vi,
                 model,
+                vi,
                 DynamicPPL.SamplingContext(rng, spl, DynamicPPL.leafcontext(model.context)),
             ),
         ),

src/mcmc/sghmc.jl

@@ -67,7 +67,9 @@ function DynamicPPL.initialstep(
     # Compute initial sample and state.
     sample = Transition(model, vi)
     ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
+        Turing.LogDensityFunction(
+            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
+        ),
     )
     state = SGHMCState(ℓ, vi, zero(vi[spl]))
 
@@ -227,7 +229,9 @@ function DynamicPPL.initialstep(
     # Create first sample and state.
     sample = SGLDTransition(model, vi, zero(spl.alg.stepsize(0)))
     ℓ = LogDensityProblemsAD.ADgradient(
-        Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
+        Turing.LogDensityFunction(
+            model, vi, DynamicPPL.SamplingContext(spl, DynamicPPL.DefaultContext())
+        ),
     )
     state = SGLDState(ℓ, vi, 1)
 

src/optimisation/Optimisation.jl

@@ -110,7 +110,7 @@ Create a callable `OptimLogDensity` struct that evaluates a model using the give
 """
 function OptimLogDensity(model::DynamicPPL.Model, context::OptimizationContext)
     init = DynamicPPL.VarInfo(model)
-    return Turing.LogDensityFunction(init, model, context)
+    return Turing.LogDensityFunction(model, init, context)
 end
 
 """

test/essential/ad.jl

@@ -43,7 +43,9 @@ function test_model_ad(model, f, syms::Vector{Symbol})
         ℓ = LogDensityProblemsAD.ADgradient(
             Turing.AutoForwardDiff(; chunksize=chunksize, tag=standardtag),
             Turing.LogDensityFunction(
-                vi, model, SampleFromPrior(), DynamicPPL.DefaultContext()
+                model,
+                vi,
+                DynamicPPL.SamplingContext(SampleFromPrior(), DynamicPPL.DefaultContext()),
             ),
         )
         l, ∇E = LogDensityProblems.logdensity_and_gradient(ℓ, z)
@@ -84,7 +86,9 @@ end
         grad_FWAD = sort(g(_x))
 
         ℓ = Turing.LogDensityFunction(
-            vi, ad_test_f, SampleFromPrior(), DynamicPPL.DefaultContext()
+            ad_test_f,
+            vi,
+            DynamicPPL.SamplingContext(SampleFromPrior(), DynamicPPL.DefaultContext()),
         )
         x = map(x -> Float64(x), vi[:])
 

test/skipped/unit_test_helper.jl

@@ -11,7 +11,9 @@ function test_grad(turing_model, grad_f; trans=Dict())
         ℓ = LogDensityProblemsAD.ADgradient(
             Turing.AutoTracker(),
             Turing.LogDensityFunction(
-                vi, model_f, SampleFromPrior(), DynamicPPL.DefaultContext()
+                model_f,
+                vi,
+                DynamicPPL.SamplingContext(SampleFromPrior(), DynamicPPL.DefaultContext()),
             ),
         )
         for _ in 1:10000