Refine AD extension wrappers and buffer reuse.

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

Author: yebai

ext/AbstractPPLDifferentiationInterfaceExt.jl

@@ -13,9 +13,7 @@ end
 AbstractPPL.capabilities(::Type{<:DIPrepared}) = DerivativeOrder{1}()
 AbstractPPL.dimension(p::DIPrepared) = AbstractPPL.dimension(p.evaluator)
 
-function (p::DIPrepared)(x)
-    return p.evaluator(x)
-end
+(p::DIPrepared)(x) = p.evaluator(x)
 
 function AbstractPPL.prepare(
     adtype::ADTypes.AbstractADType, problem, x::AbstractVector{<:AbstractFloat}

ext/AbstractPPLEnzymeExt.jl

@@ -4,36 +4,40 @@ using AbstractPPL: AbstractPPL, DerivativeOrder
 using ADTypes: AutoEnzyme
 using Enzyme: Enzyme
 
-struct EnzymePrepared{E}
+struct EnzymePrepared{E,G}
     evaluator::E
+    gradient::G
 end
 
 AbstractPPL.capabilities(::Type{<:EnzymePrepared}) = DerivativeOrder{1}()
 AbstractPPL.dimension(p::EnzymePrepared) = AbstractPPL.dimension(p.evaluator)
 
-function (p::EnzymePrepared)(x)
-    return p.evaluator(x)
-end
+(p::EnzymePrepared)(x) = p.evaluator(x)
 
 function AbstractPPL.prepare(::AutoEnzyme, problem, x::AbstractVector{<:AbstractFloat})
     evaluator = AbstractPPL.ADProblems.VectorEvaluator(
         AbstractPPL.prepare(problem, x), length(x)
     )
-    return EnzymePrepared(evaluator)
+    return EnzymePrepared(evaluator, similar(x))
 end
 
 @inline function AbstractPPL.value_and_gradient(
     p::EnzymePrepared, x::AbstractVector{<:AbstractFloat}
 )
-    dx = zero(x)
+    dx = p.gradient
+    length(dx) == length(x) || throw(
+        DimensionMismatch(
+            "Expected a vector of length $(length(dx)), but got length $(length(x))."
+        ),
+    )
+    fill!(dx, zero(eltype(dx)))
     result = Enzyme.autodiff(
         Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal),
         Enzyme.Const(p.evaluator),
         Enzyme.Active,
         Enzyme.Duplicated(x, dx),
     )
-    val = result[2]  # The primal value is returned in the second tuple entry.
-    return (val, dx)
+    return (result[2], copy(dx))
 end
 
 end # module

ext/AbstractPPLFiniteDifferencesExt.jl

@@ -23,9 +23,7 @@ function (p::FDPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator})(values::N
     return p.evaluator(values)
 end
 
-function (p::FDPrepared)(x)
-    return p.evaluator(x)
-end
+(p::FDPrepared)(x) = p.evaluator(x)
 
 function AbstractPPL.prepare(adtype::AutoFiniteDifferences, problem, values::NamedTuple)
     evaluator = AbstractPPL.ADProblems.NamedTupleEvaluator(

ext/AbstractPPLForwardDiffExt.jl

@@ -30,9 +30,7 @@ function (p::ForwardDiffPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator})(
     return p.evaluator(values)
 end
 
-function (p::ForwardDiffPrepared)(x)
-    return p.evaluator(x)
-end
+(p::ForwardDiffPrepared)(x) = p.evaluator(x)
 
 function AbstractPPL.prepare(::AutoForwardDiff, problem, values::NamedTuple)
     evaluator = AbstractPPL.ADProblems.NamedTupleEvaluator(

ext/AbstractPPLMooncakeExt.jl

@@ -12,9 +12,7 @@ end
 AbstractPPL.capabilities(::Type{<:MooncakePrepared}) = DerivativeOrder{1}()
 AbstractPPL.dimension(p::MooncakePrepared) = AbstractPPL.dimension(p.evaluator)
 
-function (p::MooncakePrepared)(x)
-    return p.evaluator(x)
-end
+(p::MooncakePrepared)(x) = p.evaluator(x)
 
 function _prepare(adtype, problem, values::NamedTuple)
     evaluator = AbstractPPL.ADProblems.NamedTupleEvaluator(
@@ -68,11 +66,6 @@ end
     p::MooncakePrepared{<:AbstractPPL.ADProblems.VectorEvaluator},
     x::AbstractVector{<:AbstractFloat},
 )
-    AbstractPPL.dimension(p.evaluator) == length(x) || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(AbstractPPL.dimension(p.evaluator)), but got length $(length(x)).",
-        ),
-    )
     val, (_, grad) = Mooncake.value_and_gradient!!(p.cache, p.evaluator, x)
     return (val, grad)
 end

test/ext/mooncake/mooncake.jl

@@ -114,7 +114,7 @@ const config_forward = ADTypes.AutoMooncakeForward(; config=Mooncake.Config())
                     sprint(showerror, err),
                 )
                 @test_throws MethodError prepared([3, 1, 2])
-                @test_throws DimensionMismatch AbstractPPL.value_and_gradient(
+                @test_throws Mooncake.PreparedCacheSpecError AbstractPPL.value_and_gradient(
                     prepared, [3.0, 1.0, 2.0, 3.0]
                 )
             end