Move the content of ad.jl from Turing.jl to here

Project.toml

@@ -1,6 +1,5 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-
 version = "0.24.4"
 
 [deps]
@@ -15,6 +14,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -24,18 +24,23 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [extensions]
+DynamicPPLADTypesExt = ["ADTypes"]
 DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
+DynamicPPLForwardDiffExt = ["ADTypes", "ForwardDiff"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLZygoteRulesExt = ["ZygoteRules"]
 
 [compat]
+ADTypes = "0.2"
 AbstractMCMC = "5"
 AbstractPPL = "0.7"
 BangBang = "0.3"
@@ -46,16 +51,17 @@ ConstructionBase = "1.5.4"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.6"
+LinearAlgebra = "1.6"
 LogDensityProblems = "2"
+LogDensityProblemsAD = "1.7.0"
 MCMCChains = "6"
 MacroTools = "0.5.6"
 OrderedCollections = "1"
+Random = "1.6"
 Requires = "1"
 Setfield = "1"
-ZygoteRules = "0.2"
-LinearAlgebra = "1.6"
-Random = "1.6"
 Test = "1.6"
+ZygoteRules = "0.2"
 julia = "1.6"
 
 [extras]

ext/DynamicPPLADTypesExt.jl

@@ -0,0 +1,33 @@
+module DynamicPPLADTypesExt
+
+if isdefined(Base, :get_extension)
+    using DynamicPPL: DynamicPPL, LogDensityProblemsAD
+    using ADTypes
+else
+    using ..DynamicPPL: DynamicPPL, LogDensityProblemsAD
+    using ..ADTypes
+end
+
+getADType(spl::DynamicPPL.Sampler) = getADType(spl.alg)
+getADType(::DynamicPPL.SampleFromPrior) = ADTypes.AutoForwardDiff(; chunksize=0)
+getADType(ctx::DynamicPPL.SamplingContext) = getADType(ctx.sampler)
+getADType(ctx::DynamicPPL.AbstractContext) = getADType(DynamicPPL.NodeTrait(ctx), ctx)
+
+function getADType(::DynamicPPL.IsLeaf, ctx::DynamicPPL.AbstractContext)
+    return ADTypes.AutoForwardDiff(; chunksize=0)
+end
+function getADType(::DynamicPPL.IsParent, ctx::DynamicPPL.AbstractContext)
+    return getADType(DynamicPPL.childcontext(ctx))
+end
+
+function LogDensityProblemsAD.ADgradient(ℓ::DynamicPPL.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(getADType(ℓ.context), ℓ)
+end
+
+function LogDensityProblemsAD.ADgradient(ad::AutoReverseDiff, ℓ::DynamicPPL.LogDensityFunction)
+    return LogDensityProblemsAD.ADgradient(
+        Val(:ReverseDiff), ℓ; compile=Val(ad.compile), x=DynamicPPL.getparams(ℓ)
+    )
+end
+
+end # module

ext/DynamicPPLForwardDiffExt.jl

@@ -0,0 +1,56 @@
+module DynamicPPLForwardDiffExt
+
+if isdefined(Base, :get_extension)
+    using DynamicPPL: DynamicPPL, LogDensityProblemsAD
+    using ADTypes
+    using ForwardDiff
+else
+    using ..DynamicPPL: DynamicPPL, LogDensityProblemsAD
+    using ..ADTypes
+    using ..ForwardDiff
+end
+
+getchunksize(::AutoForwardDiff{chunk}) where {chunk} = chunk
+
+standardtag(::AutoForwardDiff{<:Any,Nothing}) = true
+standardtag(::AutoForwardDiff) = false
+
+function LogDensityProblemsAD.ADgradient(
+    ad::AutoForwardDiff, ℓ::DynamicPPL.LogDensityFunction
+)
+    θ = DynamicPPL.getparams(ℓ)
+    f = Base.Fix1(LogDensityProblems.logdensity, ℓ)
+
+    # Define configuration for ForwardDiff.
+    tag = if standardtag(ad)
+        ForwardDiff.Tag(DynamicPPL.DynamicPPLTag(), eltype(θ))
+    else
+        ForwardDiff.Tag(f, eltype(θ))
+    end
+    chunk_size = getchunksize(ad)
+    chunk = if chunk_size == 0
+        ForwardDiff.Chunk(θ)
+    else
+        ForwardDiff.Chunk(length(θ), chunk_size)
+    end
+
+    return LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓ; chunk, tag, x=θ)
+end
+
+# Allow Turing tag in gradient etc. calls of the log density function
+function ForwardDiff.checktag(
+    ::Type{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag(),V}},
+    ::LogDensityFunction,
+    ::AbstractArray{V},
+) where {V}
+    return true
+end
+function ForwardDiff.checktag(
+    ::Type{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag(),V}},
+    ::Base.Fix1{typeof(LogDensityProblems.logdensity),<:LogDensityFunction},
+    ::AbstractArray{V},
+) where {V}
+    return true
+end
+
+end # module

src/DynamicPPL.jl

@@ -183,12 +183,20 @@ end
 
 @static if !isdefined(Base, :get_extension)
     function __init__()
+        @require ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" include(
+            "../ext/DynamicPPLADTypesExt.jl"
+        )
         @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" include(
             "../ext/DynamicPPLChainRulesCoreExt.jl"
         )
         @require EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" include(
             "../ext/DynamicPPLEnzymeCoreExt.jl"
         )
+        @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" begin
+            @require ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" include(
+                "../ext/DynamicPPLForwardDiffExt.jl"
+            )
+        end
         @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(
             "../ext/DynamicPPLMCMCChainsExt.jl"
         )
@@ -198,4 +206,8 @@ end
     end
 end
 
+# Standard tag: Improves stacktraces
+# Ref: https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
+struct DynamicPPLTag end
+
 end # module

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
@@ -23,6 +24,7 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+ADTypes = "0.2"
 AbstractMCMC = "5"
 AbstractPPL = "0.7"
 Bijectors = "0.13"