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]>

Author: 4 people

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"

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(

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

src/essential/Essential.jl

@@ -16,6 +16,7 @@ using StatsFuns: logsumexp, softmax
 
 import AdvancedPS
 import LogDensityProblems
+import LogDensityProblemsAD
 
 include("container.jl")
 include("ad.jl")

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

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

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, ℓ)

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.

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"

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