Move ad.jl to DynamicPPL (#2158)

* `getADBackend`s are renamed to `getADType` and moved to `Inference` module as well as `LogDensityProblemsAD.ADgradient(ℓ::LogDensityFunction)` * The ` LogDensityProblemsAD.ADgradient(adtype, ℓ)` specific to RD and FD are moved to DynamicPPL * The idea is that with DynamicPPL, call to `ADgradient` must also gives the `adtype`, in Turing, we just use the `adtype` from the algorithm --------- Co-authored-by: Hong Ge <[email protected]>
TuringLang · Feb 19, 2024 · 616a07f · 616a07f · sunxd3 · Feb 19, 2024
1 parent cf647b1
commit 616a07f
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 71 deletions.
diff --git a/HISTORY.md b/HISTORY.md
@@ -1,3 +1,10 @@
+# Release 0.30.5
+
+- `essential/ad.jl` is removed, `ForwardDiff` and `ReverseDiff` integrations via `LogDensityProblemsAD` are moved to `DynamicPPL` and live in corresponding package extensions.
+- `LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)` (i.e. the single argument method) is moved to `Inference` module. It will create `ADgradient` using the `adtype` information stored in `context` field of `ℓ`.
+- `getADbackend` function is renamed to `getADType`, the interface is preserved, but packages that previously used `getADbackend` should be updated to use `getADType`.
+- `TuringTag` for ForwardDiff is also removed, now `DynamicPPLTag` is defined in `DynamicPPL` package and should serve the same [purpose](https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/).
+
 # Release 0.30.0
 
 - [`ADTypes.jl`](https://github.com/SciML/ADTypes.jl) replaced Turing's global AD backend. Users should now specify the desired `ADType` directly in sampler constructors, e.g., `HMC(0.1, 10; adtype=AutoForwardDiff(; chunksize))`, or `HMC(0.1, 10; adtype=AutoReverseDiff(false))` (`false` indicates not to use compiled tape).

diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.30.4"
+version = "0.30.5"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -58,7 +58,7 @@ Distributions = "0.23.3, 0.24, 0.25"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.8, 0.9"
 DynamicHMC = "3.4"
-DynamicPPL = "0.24"
+DynamicPPL = "0.24.7"
 EllipticalSliceSampling = "0.5, 1, 2"
 ForwardDiff = "0.10.3"
 Libtask = "0.7, 0.8"

diff --git a/src/Turing.jl b/src/Turing.jl
@@ -32,14 +32,6 @@ function setprogress!(progress::Bool)
     return progress
 end
 
-# Standard tag: Improves stacktraces
-# Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
-struct TuringTag end
-
-# Allow Turing tag in gradient etc. calls of the log density function
-ForwardDiff.checktag(::Type{ForwardDiff.Tag{TuringTag, V}}, ::LogDensityFunction, ::AbstractArray{V}) where {V} = true
-ForwardDiff.checktag(::Type{ForwardDiff.Tag{TuringTag, V}}, ::Base.Fix1{typeof(LogDensityProblems.logdensity),<:LogDensityFunction}, ::AbstractArray{V}) where {V} = true
-
 # Random probability measures.
 include("stdlib/distributions.jl")
 include("stdlib/RandomMeasures.jl")

diff --git a/src/essential/Essential.jl b/src/essential/Essential.jl
@@ -14,11 +14,8 @@ using StatsFuns: logsumexp, softmax
 using ADTypes: ADTypes, AutoForwardDiff, AutoTracker, AutoReverseDiff, AutoZygote
 
 import AdvancedPS
-import LogDensityProblems
-import LogDensityProblemsAD
 
 include("container.jl")
-include("ad.jl")
 
 export  @model,
         @varname,

diff --git a/src/essential/ad.jl b/src/essential/ad.jl
diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl
@@ -29,7 +29,6 @@ import AdvancedHMC; const AHMC = AdvancedHMC
 import AdvancedMH; const AMH = AdvancedMH
 import AdvancedPS
 import BangBang
-import ..Essential: getADbackend
 import EllipticalSliceSampling
 import LogDensityProblems
 import LogDensityProblemsAD
@@ -78,7 +77,6 @@ abstract type ParticleInference <: InferenceAlgorithm end
 abstract type Hamiltonian <: InferenceAlgorithm end
 abstract type StaticHamiltonian <: Hamiltonian end
 abstract type AdaptiveHamiltonian <: Hamiltonian end
-getADbackend(alg::Hamiltonian) = alg.adtype
 
 """
     ExternalSampler{S<:AbstractSampler}
@@ -98,6 +96,20 @@ Wrap a sampler so it can be used as an inference algorithm.
 """
 externalsampler(sampler::AbstractSampler) = ExternalSampler(sampler)
 
+getADType(spl::Sampler) = getADType(spl.alg)
+getADType(::SampleFromPrior) = AutoForwardDiff(; chunksize=0)
+
+getADType(ctx::DynamicPPL.SamplingContext) = getADType(ctx.sampler)
+getADType(ctx::DynamicPPL.AbstractContext) = getADType(DynamicPPL.NodeTrait(ctx), ctx)
+getADType(::DynamicPPL.IsLeaf, ctx::DynamicPPL.AbstractContext) = AutoForwardDiff(; chunksize=0)
+getADType(::DynamicPPL.IsParent, ctx::DynamicPPL.AbstractContext) = getADType(DynamicPPL.childcontext(ctx))
+
+getADType(alg::Hamiltonian) = alg.adtype
+
+function LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(getADType(ℓ.context), ℓ)
+end
+
 function LogDensityProblems.logdensity(
     f::Turing.LogDensityFunction{<:AbstractVarInfo,<:Model,<:DynamicPPL.DefaultContext},
     x::NamedTuple

diff --git a/test/essential/ad.jl b/test/essential/ad.jl
@@ -165,17 +165,6 @@
 
     end
 
-    @testset "tag" begin
-        for chunksize in (0, 1, 10)
-            ad = Turing.AutoForwardDiff(; chunksize=chunksize)
-            @test ad === Turing.AutoForwardDiff(; chunksize=chunksize)
-            @test Turing.Essential.standardtag(ad)
-            for standardtag in (false, 0, 1)
-                @test !Turing.Essential.standardtag(Turing.AutoForwardDiff(; chunksize=chunksize, tag=standardtag))
-            end
-        end
-    end
-
     @testset "ReverseDiff compiled without linking" begin
         f = DynamicPPL.LogDensityFunction(gdemo_default)
         θ = DynamicPPL.getparams(f)