CompatHelper: bump compat for LogDensityProblems to 2, (keep existing…

… compat) (#1917) * CompatHelper: bump compat for LogDensityProblems to 2, (keep existing compat) * Use LogDensityProblemsAD * Fix tests for Gibbs (#1920) * Update ad.jl * Update ad.jl Co-authored-by: CompatHelper Julia <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: Hong Ge <[email protected]>
TuringLang · Dec 10, 2022 · 61b06f6 · 61b06f6 · devmotion · Dec 10, 2022
1 parent 51fddaf
commit 61b06f6
Show file tree

Hide file tree

Showing 14 changed files with 35 additions and 27 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.23.0"
+version = "0.23.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -20,6 +20,7 @@ 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"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
@@ -49,8 +50,9 @@ DocStringExtensions = "0.8, 0.9"
 DynamicPPL = "0.21"
 EllipticalSliceSampling = "0.5, 1"
 ForwardDiff = "0.10.3"
-LogDensityProblems = "0.12, 1"
 Libtask = "0.7, 0.8"
+LogDensityProblems = "2"
+LogDensityProblemsAD = "1"
 MCMCChains = "5"
 NamedArrays = "0.9"
 Reexport = "0.2, 1"

diff --git a/src/contrib/inference/dynamichmc.jl b/src/contrib/inference/dynamichmc.jl
@@ -44,7 +44,7 @@ function gibbs_state(
     varinfo::AbstractVarInfo,
 )
     # Update the log density function and its cached evaluation.
-    ℓ = LogDensityProblems.ADgradient(Turing.LogDensityFunction(varinfo, model, spl, DynamicPPL.DefaultContext()))
+    ℓ = LogDensityProblemsAD.ADgradient(Turing.LogDensityFunction(varinfo, model, spl, DynamicPPL.DefaultContext()))
     Q = DynamicHMC.evaluate_ℓ(ℓ, varinfo[spl])
     return DynamicNUTSState(ℓ, varinfo, Q, state.metric, state.stepsize)
 end
@@ -65,7 +65,7 @@ function DynamicPPL.initialstep(
     end
 
     # Define log-density function.
-    ℓ = LogDensityProblems.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
+    ℓ = LogDensityProblemsAD.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
 
     # Perform initial step.
     results = DynamicHMC.mcmc_keep_warmup(

diff --git a/src/contrib/inference/sghmc.jl b/src/contrib/inference/sghmc.jl
@@ -62,7 +62,7 @@ function DynamicPPL.initialstep(
 
     # Compute initial sample and state.
     sample = Transition(vi)
-    ℓ = LogDensityProblems.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
+    ℓ = LogDensityProblemsAD.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
     state = SGHMCState(ℓ, vi, zero(vi[spl]))
 
     return sample, state
@@ -215,7 +215,7 @@ function DynamicPPL.initialstep(
 
     # Create first sample and state.
     sample = SGLDTransition(vi, zero(spl.alg.stepsize(0)))
-    ℓ = LogDensityProblems.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
+    ℓ = LogDensityProblemsAD.ADgradient(Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext()))
     state = SGLDState(ℓ, vi, 1)
 
     return sample, state

diff --git a/src/essential/Essential.jl b/src/essential/Essential.jl
@@ -16,6 +16,7 @@ using StatsFuns: logsumexp, softmax
 
 import AdvancedPS
 import LogDensityProblems
+import LogDensityProblemsAD
 
 include("container.jl")
 include("ad.jl")

diff --git a/src/essential/ad.jl b/src/essential/ad.jl
@@ -78,11 +78,11 @@ Find the autodifferentiation backend of the algorithm `alg`.
 getADbackend(spl::Sampler) = getADbackend(spl.alg)
 getADbackend(::SampleFromPrior) = ADBackend()()
 
-function LogDensityProblems.ADgradient(ℓ::Turing.LogDensityFunction)
-    return LogDensityProblems.ADgradient(getADbackend(ℓ.sampler), ℓ)
+function LogDensityProblemsAD.ADgradient(ℓ::Turing.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(getADbackend(ℓ.sampler), ℓ)
 end
 
-function LogDensityProblems.ADgradient(ad::ForwardDiffAD, ℓ::Turing.LogDensityFunction)
+function LogDensityProblemsAD.ADgradient(ad::ForwardDiffAD, ℓ::Turing.LogDensityFunction)
     θ = ℓ.varinfo[ℓ.sampler]
     f = Base.Fix1(LogDensityProblems.logdensity, ℓ)
 
@@ -99,21 +99,21 @@ function LogDensityProblems.ADgradient(ad::ForwardDiffAD, ℓ::Turing.LogDensity
         ForwardDiff.GradientConfig(f, θ, ForwardDiff.Chunk(length(θ), chunk_size), tag)
     end
 
-    return LogDensityProblems.ADgradient(Val(:ForwardDiff), ℓ; gradientconfig=config)
+    return LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; gradientconfig=config)
 end
 
-function LogDensityProblems.ADgradient(::TrackerAD, ℓ::Turing.LogDensityFunction)
-    return LogDensityProblems.ADgradient(Val(:Tracker), ℓ)
+function LogDensityProblemsAD.ADgradient(::TrackerAD, ℓ::Turing.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(Val(:Tracker), ℓ)
 end
 
-function LogDensityProblems.ADgradient(::ZygoteAD, ℓ::Turing.LogDensityFunction)
-    return LogDensityProblems.ADgradient(Val(:Zygote), ℓ)
+function LogDensityProblemsAD.ADgradient(::ZygoteAD, ℓ::Turing.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(Val(:Zygote), ℓ)
 end
 
 for cache in (:true, :false)
     @eval begin
-        function LogDensityProblems.ADgradient(::ReverseDiffAD{$cache}, ℓ::Turing.LogDensityFunction)
-            return LogDensityProblems.ADgradient(Val(:ReverseDiff), ℓ; compile=Val($cache))
+        function LogDensityProblemsAD.ADgradient(::ReverseDiffAD{$cache}, ℓ::Turing.LogDensityFunction)
+            return LogDensityProblemsAD.ADgradient(Val(:ReverseDiff), ℓ; compile=Val($cache))
         end
     end
 end

diff --git a/src/inference/Inference.jl b/src/inference/Inference.jl
@@ -31,6 +31,7 @@ import BangBang
 import ..Essential: getADbackend
 import EllipticalSliceSampling
 import LogDensityProblems
+import LogDensityProblemsAD
 import Random
 import MCMCChains
 import StatsBase: predict

diff --git a/src/inference/hmc.jl b/src/inference/hmc.jl
@@ -158,7 +158,7 @@ function DynamicPPL.initialstep(
     # Create a Hamiltonian.
     metricT = getmetricT(spl.alg)
     metric = metricT(length(theta))
-    ℓ = LogDensityProblems.ADgradient(
+    ℓ = LogDensityProblemsAD.ADgradient(
         Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
     )
     logπ = Base.Fix1(LogDensityProblems.logdensity, ℓ)
@@ -264,7 +264,7 @@ end
 
 function get_hamiltonian(model, spl, vi, state, n)
     metric = gen_metric(n, spl, state)
-    ℓ = LogDensityProblems.ADgradient(
+    ℓ = LogDensityProblemsAD.ADgradient(
         Turing.LogDensityFunction(vi, model, spl, DynamicPPL.DefaultContext())
     )
     ℓπ = Base.Fix1(LogDensityProblems.logdensity, ℓ)

diff --git a/src/modes/ModeEstimation.jl b/src/modes/ModeEstimation.jl
@@ -12,6 +12,7 @@ using DynamicPPL: Model, AbstractContext, VarInfo, VarName,
     get_and_set_val!, istrans
 
 import LogDensityProblems
+import LogDensityProblemsAD
 
 export  constrained_space,  
         MAP,
@@ -111,7 +112,7 @@ function (f::OptimLogDensity)(F, G, z)
     if G !== nothing
         # Calculate negative log joint and its gradient.
         # TODO: Make OptimLogDensity already an LogDensityProblems.ADgradient? Allow to specify AD?
-        ℓ = LogDensityProblems.ADgradient(f)
+        ℓ = LogDensityProblemsAD.ADgradient(f)
         neglogp, ∇neglogp = LogDensityProblems.logdensity_and_gradient(ℓ, z)
 
         # Save the gradient to the pre-allocated array.

diff --git a/test/Project.toml b/test/Project.toml
@@ -12,6 +12,7 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 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"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
@@ -42,7 +43,8 @@ DynamicHMC = "2.1.6, 3.0"
 DynamicPPL = "0.21"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32"
-LogDensityProblems = "0.12, 1"
+LogDensityProblems = "2"
+LogDensityProblemsAD = "1"
 MCMCChains = "5"
 NamedArrays = "0.9.4"
 Optim = "0.22, 1.0"

diff --git a/test/essential/ad.jl b/test/essential/ad.jl
@@ -30,14 +30,14 @@
         ℓ = Turing.LogDensityFunction(vi, ad_test_f, SampleFromPrior(), DynamicPPL.DefaultContext())
         x = map(x->Float64(x), vi[SampleFromPrior()])
 
-        trackerℓ = LogDensityProblems.ADgradient(TrackerAD(), ℓ)
-        @test trackerℓ isa LogDensityProblems.TrackerGradientLogDensity
+        trackerℓ = LogDensityProblemsAD.ADgradient(TrackerAD(), ℓ)
+        @test trackerℓ isa LogDensityProblemsAD.TrackerGradientLogDensity
         @test trackerℓ.ℓ === ℓ
         ∇E1 = LogDensityProblems.logdensity_and_gradient(trackerℓ, x)[2]
         @test sort(∇E1) ≈ grad_FWAD atol=1e-9
 
-        zygoteℓ = LogDensityProblems.ADgradient(ZygoteAD(), ℓ)
-        @test zygoteℓ isa LogDensityProblems.ZygoteGradientLogDensity
+        zygoteℓ = LogDensityProblemsAD.ADgradient(ZygoteAD(), ℓ)
+        @test zygoteℓ isa LogDensityProblemsAD.ZygoteGradientLogDensity
         @test zygoteℓ.ℓ === ℓ
         ∇E2 = LogDensityProblems.logdensity_and_gradient(zygoteℓ, x)[2]
         @test sort(∇E2) ≈ grad_FWAD atol=1e-9

diff --git a/test/inference/Inference.jl b/test/inference/Inference.jl
@@ -89,7 +89,7 @@
         chn3 = sample(gdemo_default, alg3, 5_000; save_state=true)
         check_gdemo(chn3)
 
-        chn3_contd = sample(gdemo_default, alg3, 2_000; resume_from=chn3)
+        chn3_contd = sample(gdemo_default, alg3, 5_000; resume_from=chn3)
         check_gdemo(chn3_contd)
     end
     @testset "Contexts" begin

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -42,6 +42,7 @@ using Turing.Essential: TuringDenseMvNormal, TuringDiagMvNormal
 using Turing.Variational: TruncatedADAGrad, DecayedADAGrad, AdvancedVI
 
 import LogDensityProblems
+import LogDensityProblemsAD
 
 setprogress!(false)
 

diff --git a/test/skipped/unit_test_helper.jl b/test/skipped/unit_test_helper.jl
@@ -8,7 +8,7 @@ function test_grad(turing_model, grad_f; trans=Dict())
     end
     d = length(vi.vals)
     @testset "Gradient using random inputs" begin
-        ℓ = LogDensityProblems.ADgradient(
+        ℓ = LogDensityProblemsAD.ADgradient(
             TrackerAD(),
             Turing.LogDensityFunction(vi, model_f, SampleFromPrior(), DynamicPPL.DefaultContext()),
         )

diff --git a/test/test_utils/ad_utils.jl b/test/test_utils/ad_utils.jl
@@ -93,7 +93,7 @@ function test_model_ad(model, f, syms::Vector{Symbol})
     # Compare with `logdensity_and_gradient`.
     z = vi[SampleFromPrior()]
     for chunksize in (0, 1, 10), standardtag in (true, false, 0, 3)
-        ℓ = LogDensityProblems.ADgradient(
+        ℓ = LogDensityProblemsAD.ADgradient(
             ForwardDiffAD{chunksize, standardtag}(),
             Turing.LogDensityFunction(vi, model, SampleFromPrior(), DynamicPPL.DefaultContext()),
         )