Use native AD APIs for ForwardDiff, Enzyme, and Mooncake; make DI optional

- Move DifferentiationInterface to [weakdeps]; add DynamicPPLDifferentiationInterfaceExt
  as fallback for backends without native implementations
- Add native ForwardDiff gradient via GradientConfig (DynamicPPLForwardDiffExt)
- Add native Enzyme gradient via autodiff(ReverseWithPrimal, ...) (new DynamicPPLEnzymeExt)
- Keep native Mooncake reverse/forward gradient (DynamicPPLMooncakeExt)
- Add Enzyme to test env; drop DI from test env

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: yebai

Project.toml

@@ -12,7 +12,6 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -29,6 +28,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -38,11 +39,13 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
+DynamicPPLDifferentiationInterfaceExt = ["DifferentiationInterface"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
+DynamicPPLEnzymeExt = ["Enzyme"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"]
-DynamicPPLMooncakeExt = ["Mooncake", "DifferentiationInterface"]
+DynamicPPLMooncakeExt = ["Mooncake"]
 DynamicPPLReverseDiffExt = ["ReverseDiff"]
 
 [compat]
@@ -55,9 +58,9 @@ Bijectors = "0.15.17"
 Chairmarks = "1.3.1"
 Compat = "4"
 ConstructionBase = "1.5.4"
-DifferentiationInterface = "0.6.41, 0.7"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
+Enzyme = "0.13"
 EnzymeCore = "0.6 - 0.8"
 FillArrays = "1.16.0"
 ForwardDiff = "0.10.12, 1"

ext/DynamicPPLDifferentiationInterfaceExt.jl

@@ -0,0 +1,64 @@
+module DynamicPPLDifferentiationInterfaceExt
+
+import DifferentiationInterface as DI
+using DynamicPPL:
+    DynamicPPL,
+    AccumulatorTuple,
+    LogDensityAt,
+    Model,
+    VarNamedTuple,
+    AbstractTransformStrategy,
+    _use_closure,
+    logdensity_at
+using ADTypes: ADTypes
+
+function DynamicPPL._prepare_gradient(
+    adtype::ADTypes.AbstractADType,
+    x::AbstractVector{<:Real},
+    model::Model,
+    getlogdensity::Any,
+    varname_ranges::VarNamedTuple,
+    transform_strategy::AbstractTransformStrategy,
+    accs::AccumulatorTuple,
+)
+    args = (model, getlogdensity, varname_ranges, transform_strategy, accs)
+    return if _use_closure(adtype)
+        DI.prepare_gradient(LogDensityAt(args...), adtype, x)
+    else
+        DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...)
+    end
+end
+
+function DynamicPPL._value_and_gradient(
+    adtype::ADTypes.AbstractADType,
+    prep,
+    params::AbstractVector{<:Real},
+    model::Model,
+    getlogdensity::Any,
+    varname_ranges::VarNamedTuple,
+    transform_strategy::AbstractTransformStrategy,
+    accs::AccumulatorTuple,
+)
+    return if _use_closure(adtype)
+        DI.value_and_gradient(
+            LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs),
+            prep,
+            adtype,
+            params,
+        )
+    else
+        DI.value_and_gradient(
+            logdensity_at,
+            prep,
+            adtype,
+            params,
+            DI.Constant(model),
+            DI.Constant(getlogdensity),
+            DI.Constant(varname_ranges),
+            DI.Constant(transform_strategy),
+            DI.Constant(accs),
+        )
+    end
+end
+
+end # module

ext/DynamicPPLEnzymeExt.jl

@@ -0,0 +1,45 @@
+module DynamicPPLEnzymeExt
+
+using DynamicPPL: ADTypes, DynamicPPL
+using Enzyme: Enzyme
+
+function DynamicPPL._prepare_gradient(
+    ::ADTypes.AutoEnzyme,
+    x::AbstractVector{<:Real},
+    model::DynamicPPL.Model,
+    getlogdensity::Any,
+    varname_ranges::DynamicPPL.VarNamedTuple,
+    transform_strategy::DynamicPPL.AbstractTransformStrategy,
+    accs::DynamicPPL.AccumulatorTuple,
+)
+    return (; dx=similar(x))
+end
+
+function DynamicPPL._value_and_gradient(
+    ::ADTypes.AutoEnzyme,
+    prep,
+    params::AbstractVector{<:Real},
+    model::DynamicPPL.Model,
+    getlogdensity::Any,
+    varname_ranges::DynamicPPL.VarNamedTuple,
+    transform_strategy::DynamicPPL.AbstractTransformStrategy,
+    accs::DynamicPPL.AccumulatorTuple,
+)
+    f = DynamicPPL.LogDensityAt(
+        model, getlogdensity, varname_ranges, transform_strategy, accs
+    )
+    dx = prep.dx
+    fill!(dx, zero(eltype(dx)))
+    # Const(f): LogDensityAt is not being differentiated; without Const, Enzyme errors
+    # because it cannot prove the function argument is readonly.
+    # autodiff(ReverseWithPrimal, ...) returns ((), val); dx is mutated in-place.
+    _, val = Enzyme.autodiff(
+        Enzyme.ReverseWithPrimal,
+        Enzyme.Const(f),
+        Enzyme.Active,
+        Enzyme.Duplicated(params, dx),
+    )
+    return val, copy(dx)
+end
+
+end # module

ext/DynamicPPLForwardDiffExt.jl

@@ -12,10 +12,6 @@ function DynamicPPL.tweak_adtype(
 ) where {chunk_size}
     # 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
@@ -32,4 +28,47 @@ function DynamicPPL.tweak_adtype(
     return ADTypes.AutoForwardDiff(; chunksize=ForwardDiff.chunksize(chunk), tag=tag)
 end
 
+function DynamicPPL._prepare_gradient(
+    adtype::ADTypes.AutoForwardDiff{chunk_size},
+    x::AbstractVector{<:Real},
+    model::DynamicPPL.Model,
+    getlogdensity::Any,
+    varname_ranges::DynamicPPL.VarNamedTuple,
+    transform_strategy::DynamicPPL.AbstractTransformStrategy,
+    accs::DynamicPPL.AccumulatorTuple,
+) where {chunk_size}
+    f = DynamicPPL.LogDensityAt(
+        model, getlogdensity, varname_ranges, transform_strategy, accs
+    )
+    chunk = if chunk_size == 0 || chunk_size === nothing
+        ForwardDiff.Chunk(x)
+    else
+        ForwardDiff.Chunk(length(x), chunk_size)
+    end
+    cfg = ForwardDiff.GradientConfig(f, x, chunk, adtype.tag)
+    grad = similar(x)
+    return (; cfg, grad)
+end
+
+function DynamicPPL._value_and_gradient(
+    ::ADTypes.AutoForwardDiff,
+    prep,
+    params::AbstractVector{<:Real},
+    model::DynamicPPL.Model,
+    getlogdensity::Any,
+    varname_ranges::DynamicPPL.VarNamedTuple,
+    transform_strategy::DynamicPPL.AbstractTransformStrategy,
+    accs::DynamicPPL.AccumulatorTuple,
+)
+    f = DynamicPPL.LogDensityAt(
+        model, getlogdensity, varname_ranges, transform_strategy, accs
+    )
+    # Val{false}() skips tag checking, since our DynamicPPLTag is reused across calls
+    # with different LogDensityAt instances.
+    ForwardDiff.gradient!(prep.grad, f, params, prep.cfg, Val{false}())
+    # gradient!(::AbstractArray, ...) doesn't return the value, so evaluate separately.
+    value = f(params)
+    return value, copy(prep.grad)
+end
+
 end # module

ext/DynamicPPLMooncakeExt.jl

@@ -3,9 +3,13 @@ module DynamicPPLMooncakeExt
 using DynamicPPL: DynamicPPL, is_transformed
 using Mooncake:
     Mooncake,
+    Dual,
     NoTangent,
     prepare_derivative_cache,
     prepare_gradient_cache,
+    primal,
+    tangent,
+    value_and_derivative!!,
     value_and_gradient!!
 
 # These are purely optimisations (although quite significant ones sometimes, especially for
@@ -61,7 +65,8 @@ function DynamicPPL._prepare_gradient(
     accs::DynamicPPL.AccumulatorTuple,
 )
     f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs)
-    return prepare_derivative_cache(f, x; config=_cache_config(adtype))
+    cache = prepare_derivative_cache(f, x; config=_cache_config(adtype))
+    return (; cache, dx=similar(x), grad=similar(x))
 end
 
 function DynamicPPL._value_and_gradient(
@@ -90,7 +95,16 @@ function DynamicPPL._value_and_gradient(
     accs::DynamicPPL.AccumulatorTuple,
 )
     f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs)
-    value, grad = value_and_gradient!!(prep, f, params)
+    (; cache, dx, grad) = prep
+    value = zero(eltype(grad))
+    fill!(dx, zero(eltype(dx)))
+    @inbounds for i in eachindex(grad, dx)
+        dx[i] = one(eltype(dx))
+        result = value_and_derivative!!(cache, Dual(f, NoTangent()), Dual(params, dx))
+        value = primal(result)
+        grad[i] = tangent(result)
+        dx[i] = zero(eltype(dx))
+    end
     return value, copy(grad)
 end
 

src/logdensityfunction.jl

@@ -23,7 +23,6 @@ using ADTypes: ADTypes
 using BangBang: BangBang
 using AbstractPPL: AbstractPPL, VarName
 using LogDensityProblems: LogDensityProblems
-import DifferentiationInterface as DI
 using Random: Random
 
 """
@@ -404,54 +403,9 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real})
     )
 end
 
-function _prepare_gradient(
-    adtype::ADTypes.AbstractADType,
-    x::AbstractVector{<:Real},
-    model::Model,
-    getlogdensity::Any,
-    varname_ranges::VarNamedTuple,
-    transform_strategy::AbstractTransformStrategy,
-    accs::AccumulatorTuple,
-)
-    args = (model, getlogdensity, varname_ranges, transform_strategy, accs)
-    return if _use_closure(adtype)
-        DI.prepare_gradient(LogDensityAt(args...), adtype, x)
-    else
-        DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...)
-    end
-end
-
-function _value_and_gradient(
-    adtype::ADTypes.AbstractADType,
-    prep,
-    params::AbstractVector{<:Real},
-    model::Model,
-    getlogdensity::Any,
-    varname_ranges::VarNamedTuple,
-    transform_strategy::AbstractTransformStrategy,
-    accs::AccumulatorTuple,
-)
-    return if _use_closure(adtype)
-        DI.value_and_gradient(
-            LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs),
-            prep,
-            adtype,
-            params,
-        )
-    else
-        DI.value_and_gradient(
-            logdensity_at,
-            prep,
-            adtype,
-            params,
-            DI.Constant(model),
-            DI.Constant(getlogdensity),
-            DI.Constant(varname_ranges),
-            DI.Constant(transform_strategy),
-            DI.Constant(accs),
-        )
-    end
-end
+# Extensible hooks: backends provide methods via package extensions.
+function _prepare_gradient end
+function _value_and_gradient end
 
 function LogDensityProblems.logdensity(
     ldf::LogDensityFunction, params::AbstractVector{<:Real}

src/test_utils/ad.jl

@@ -2,7 +2,6 @@ module AD
 
 using ADTypes: AbstractADType, AutoForwardDiff
 using Chairmarks: @be
-import DifferentiationInterface as DI
 using DocStringExtensions
 using DynamicPPL:
     DynamicPPL,

test/Project.toml

@@ -9,11 +9,12 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -43,19 +44,18 @@ BangBang = "0.4"
 Bijectors = "0.15.17"
 Chairmarks = "1"
 Combinatorics = "1"
-DifferentiationInterface = "0.6.41, 0.7"
 DimensionalData = "0.30"
 Distributions = "0.25"
 Documenter = "1"
+Enzyme = "0.13"
 ForwardDiff = "0.10.12, 1"
 InvertedIndices = "1"
 LogDensityProblems = "2"
 MCMCChains = "7.2.1"
 MacroTools = "0.5.6"
 MarginalLogDensities = "0.4"
-Mooncake = "0.4, 0.5"
-OrderedCollections = "1"
 OffsetArrays = "1"
+OrderedCollections = "1"
 ReverseDiff = "1"
 SpecialFunctions = "2.6.1"
 StableRNGs = "1"

test/logdensityfunction.jl

@@ -13,9 +13,10 @@ using LogDensityProblems: LogDensityProblems
 using Random: Xoshiro
 using StableRNGs: StableRNG
 
+using Enzyme: Enzyme
 using ForwardDiff: ForwardDiff
-using ReverseDiff: ReverseDiff
 using Mooncake: Mooncake
+using ReverseDiff: ReverseDiff
 
 @testset "LogDensityFunction: constructors" begin
     dist = Beta(2, 2)
@@ -177,12 +178,12 @@ end
         struct ErrorAccumulatorException <: Exception end
         struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end
         DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR
-        DynamicPPL.accumulate_assume!!(
-            ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any
-        ) = throw(ErrorAccumulatorException())
-        DynamicPPL.accumulate_observe!!(
-            ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any
-        ) = throw(ErrorAccumulatorException())
+        DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw(
+            ErrorAccumulatorException()
+        )
+        DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw(
+            ErrorAccumulatorException()
+        )
         DynamicPPL.reset(ea::ErrorAccumulator) = ea
         Base.copy(ea::ErrorAccumulator) = ea
         # Construct an LDF