Re-add ForwardDiffExt

Author: penelopeysm

Project.toml

@@ -29,6 +29,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
@@ -37,6 +38,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [extensions]
 DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
+DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLJETExt = ["JET"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLMooncakeExt = ["Mooncake"]

ext/DynamicPPLForwardDiffExt.jl

@@ -1,54 +1,40 @@
 module DynamicPPLForwardDiffExt
 
-if isdefined(Base, :get_extension)
-    using DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems, LogDensityProblemsAD
-    using ForwardDiff
-else
-    using ..DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems, LogDensityProblemsAD
-    using ..ForwardDiff
-end
-
-getchunksize(::ADTypes.AutoForwardDiff{chunk}) where {chunk} = chunk
-
-standardtag(::ADTypes.AutoForwardDiff{<:Any,Nothing}) = true
-standardtag(::ADTypes.AutoForwardDiff) = false
-
-function LogDensityProblemsAD.ADgradient(
-    ad::ADTypes.AutoForwardDiff, ℓ::DynamicPPL.LogDensityFunction
-)
-    θ = DynamicPPL.getparams(ℓ)
-    f = Base.Fix1(LogDensityProblems.logdensity, ℓ)
-
-    # Define configuration for ForwardDiff.
-    tag = if standardtag(ad)
-        ForwardDiff.Tag(DynamicPPL.DynamicPPLTag(), eltype(θ))
+using DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems
+using ForwardDiff
+
+# check if the AD type already has a tag
+use_dynamicppl_tag(::ADTypes.AutoForwardDiff{<:Any,Nothing}) = true
+use_dynamicppl_tag(::ADTypes.AutoForwardDiff) = false
+
+function DynamicPPL.tweak_adtype(
+    ad::ADTypes.AutoForwardDiff{chunk_size},
+    ::DynamicPPL.Model,
+    vi::DynamicPPL.AbstractVarInfo,
+    ::DynamicPPL.AbstractContext,
+) where {chunk_size}
+    params = vi[:]
+
+    # Use DynamicPPL tag to improve stack traces
+    # https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
+    # NOTE: DifferentiationInterface disables tag checking if the
+    # tag inside the AutoForwardDiff type is not nothing. See
+    # https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/1df562180bdcc3e91c885aa5f4162a0be2ced850/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl#L338-L350.
+    # So we don't currently need to override ForwardDiff.checktag as well.
+    tag = if use_dynamicppl_tag(ad)
+        ForwardDiff.Tag(DynamicPPL.DynamicPPLTag(), eltype(params))
     else
-        ForwardDiff.Tag(f, eltype(θ))
+        ad.tag
     end
-    chunk_size = getchunksize(ad)
+
+    # Optimise chunk size according to size of model
     chunk = if chunk_size == 0 || chunk_size === nothing
-        ForwardDiff.Chunk(θ)
+        ForwardDiff.Chunk(params)
     else
-        ForwardDiff.Chunk(length(θ), chunk_size)
+        ForwardDiff.Chunk(length(params), 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}},
-    ::DynamicPPL.LogDensityFunction,
-    ::AbstractArray{W},
-) where {V,W}
-    return true
-end
-function ForwardDiff.checktag(
-    ::Type{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag,V}},
-    ::Base.Fix1{typeof(LogDensityProblems.logdensity),<:DynamicPPL.LogDensityFunction},
-    ::AbstractArray{W},
-) where {V,W}
-    return true
+    return ADTypes.AutoForwardDiff(; chunksize=ForwardDiff.chunksize(chunk), tag=tag)
 end
 
 end # module

src/logdensityfunction.jl

@@ -116,7 +116,9 @@ struct LogDensityFunction{
         if adtype === nothing
             prep = nothing
         else
-            # Check support
+            # Make backend-specific tweaks to the adtype
+            adtype = tweak_adtype(adtype, model, varinfo, context)
+            # Check whether it is supported
             is_supported(adtype) ||
                 @warn "The AD backend $adtype is not officially supported by DynamicPPL. Gradient calculations may still work, but compatibility is not guaranteed."
             # Get a set of dummy params to use for prep
@@ -227,6 +229,25 @@ LogDensityProblems.dimension(f::LogDensityFunction) = length(getparams(f))
 
 ### Utils
 
+"""
+    tweak_adtype(
+        adtype::ADTypes.AbstractADType,
+        model::Model,
+        varinfo::AbstractVarInfo,
+        context::AbstractContext
+    )
+
+Return an 'optimised' form of the adtype. This is useful for doing
+backend-specific optimisation of the adtype (e.g., for ForwardDiff, calculating
+the chunk size: see the method override in `ext/DynamicPPLForwardDiffExt.jl`).
+The model is passed as a parameter in case the optimisation depends on the
+model.
+
+By default, this just returns the input unchanged.
+"""
+tweak_adtype(adtype::ADTypes.AbstractADType, ::Model, ::AbstractVarInfo, ::AbstractContext) =
+    adtype
+
 """
     use_closure(adtype::ADTypes.AbstractADType)