Internal _make_ad_gradient

Author: willtebbutt

Project.toml

@@ -30,6 +30,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
@@ -38,6 +39,7 @@ DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
+DynamicPPLMooncakeExt = ["Mooncake"]
 DynamicPPLReverseDiffExt = ["ReverseDiff"]
 DynamicPPLZygoteRulesExt = ["ZygoteRules"]
 
@@ -60,10 +62,11 @@ LogDensityProblems = "2"
 LogDensityProblemsAD = "1.7.0"
 MCMCChains = "6"
 MacroTools = "0.5.6"
+Mooncake = "0.4.52"
 OrderedCollections = "1"
 Random = "1.6"
-Requires = "1"
 ReverseDiff = "1"
+Requires = "1"
 Test = "1.6"
 ZygoteRules = "0.2"
 julia = "1.10"

ext/DynamicPPLMooncakeExt.jl

@@ -0,0 +1,8 @@
+module DynamicPPLMooncakeExt
+
+import LogDensityProblemsAD: ADgradient
+using DynamicPPL: ADTypes, _make_ad_gradient, LogDensityFunction
+
+ADgradient(ad::ADTypes.AutoMooncake, f::LogDensityFunction) = _make_ad_gradient(ad, f)
+
+end # module

ext/DynamicPPLReverseDiffExt.jl

@@ -1,26 +1,8 @@
 module DynamicPPLReverseDiffExt
 
-if isdefined(Base, :get_extension)
-    using DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems, LogDensityProblemsAD
-    using ReverseDiff
-else
-    using ..DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems, LogDensityProblemsAD
-    using ..ReverseDiff
-end
+import LogDensityProblemsAD: ADgradient
+using DynamicPPL: ADTypes, _make_ad_gradient, LogDensityFunction
 
-function LogDensityProblemsAD.ADgradient(
-    ad::ADTypes.AutoReverseDiff{Tcompile}, ℓ::DynamicPPL.LogDensityFunction
-) where {Tcompile}
-    return LogDensityProblemsAD.ADgradient(
-        Val(:ReverseDiff),
-        ℓ;
-        compile=Val(Tcompile),
-        # `getparams` can return `Vector{Real}`, in which case, `ReverseDiff` will initialize the gradients to Integer 0
-        # because at https://github.com/JuliaDiff/ReverseDiff.jl/blob/c982cde5494fc166965a9d04691f390d9e3073fd/src/tracked.jl#L473
-        # `zero(D)` will return 0 when D is Real.
-        # here we use `identity` to possibly concretize the type to `Vector{Float64}` in the case of `Vector{Real}`.
-        x=map(identity, DynamicPPL.getparams(ℓ)),
-    )
-end
+ADgradient(ad::ADTypes.AutoReverseDiff, f::LogDensityFunction) = _make_ad_gradient(ad, f)
 
 end # module

src/logdensityfunction.jl

@@ -149,7 +149,7 @@ LogDensityProblems.dimension(f::LogDensityFunction) = length(getparams(f))
 # parameters, or DifferentiationInterface will not have sufficient information to e.g.
 # compile a rule for Mooncake (because it won't know the type of the input), or pre-allocate
 # a tape when using ReverseDiff.jl.
-function LogDensityProblemsAD.ADgradient(ad::ADTypes.AbstractADType, ℓ::LogDensityFunction)
+function _make_ad_gradient(ad::ADTypes.AbstractADType, ℓ::LogDensityFunction)
     x = map(identity, getparams(ℓ)) # ensure we concretise the elements of the params
     return LogDensityProblemsAD.ADgradient(ad, ℓ; x)
 end