Rename evaluator API to ADProblems and fold autograd testing into the core interface.

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

Author: yebai

docs/src/pplapi.md

@@ -26,5 +26,6 @@ DerivativeOrder
 capabilities
 prepare
 value_and_gradient
+test_autograd
 dimension
 ```

ext/AbstractPPLDifferentiationInterfaceExt.jl

@@ -8,18 +8,12 @@ struct DIPrepared{E,B,C}
     evaluator::E
     backend::B
     prep::C
-    dim::Int
 end
 
 AbstractPPL.capabilities(::Type{<:DIPrepared}) = DerivativeOrder{1}()
-AbstractPPL.dimension(p::DIPrepared) = p.dim
+AbstractPPL.dimension(p::DIPrepared) = AbstractPPL.dimension(p.evaluator)
 
-function (p::DIPrepared)(x::AbstractVector{<:AbstractFloat})
-    length(x) == p.dim || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(p.dim), but got length $(length(x))."
-        ),
-    )
+function (p::DIPrepared)(x)
     return p.evaluator(x)
 end
 
@@ -29,19 +23,16 @@ end
 function AbstractPPL.prepare(
     adtype::ADTypes.AbstractADType, problem, x::AbstractVector{<:AbstractFloat}
 )
-    evaluator = AbstractPPL.prepare(problem, x)
+    evaluator = AbstractPPL.ADProblems.VectorEvaluator(
+        AbstractPPL.prepare(problem, x), length(x)
+    )
     prep = DI.prepare_gradient(evaluator, adtype, x)
-    return DIPrepared(evaluator, adtype, prep, length(x))
+    return DIPrepared(evaluator, adtype, prep)
 end
 
 @inline function AbstractPPL.value_and_gradient(
     p::DIPrepared, x::AbstractVector{<:AbstractFloat}
 )
-    length(x) == p.dim || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(p.dim), but got length $(length(x))."
-        ),
-    )
     return DI.value_and_gradient(p.evaluator, p.prep, p.backend, x)
 end
 

ext/AbstractPPLEnzymeExt.jl

@@ -6,34 +6,25 @@ using Enzyme: Enzyme
 
 struct EnzymePrepared{E}
     evaluator::E
-    dim::Int
 end
 
 AbstractPPL.capabilities(::Type{<:EnzymePrepared}) = DerivativeOrder{1}()
-AbstractPPL.dimension(p::EnzymePrepared) = p.dim
+AbstractPPL.dimension(p::EnzymePrepared) = AbstractPPL.dimension(p.evaluator)
 
-function (p::EnzymePrepared)(x::AbstractVector{<:AbstractFloat})
-    length(x) == p.dim || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(p.dim), but got length $(length(x))."
-        ),
-    )
+function (p::EnzymePrepared)(x)
     return p.evaluator(x)
 end
 
 function AbstractPPL.prepare(::AutoEnzyme, problem, x::AbstractVector{<:AbstractFloat})
-    evaluator = AbstractPPL.prepare(problem, x)
-    return EnzymePrepared(evaluator, length(x))
+    evaluator = AbstractPPL.ADProblems.VectorEvaluator(
+        AbstractPPL.prepare(problem, x), length(x)
+    )
+    return EnzymePrepared(evaluator)
 end
 
 @inline function AbstractPPL.value_and_gradient(
     p::EnzymePrepared, x::AbstractVector{<:AbstractFloat}
 )
-    length(x) == p.dim || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(p.dim), but got length $(length(x))."
-        ),
-    )
     dx = zero(x)
     result = Enzyme.autodiff(
         Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal),

ext/AbstractPPLFiniteDifferencesExt.jl

@@ -5,87 +5,63 @@ using AbstractPPL.Utils: flatten_to!!, unflatten_to!!
 using ADTypes: AutoFiniteDifferences
 using FiniteDifferences: FiniteDifferences
 
-struct FDPrepared{F,M,P}
-    evaluator
+struct FDPrepared{E,F,M}
+    evaluator::E
     f::F
     fdm::M
-    inputspec::P
 end
 
 AbstractPPL.capabilities(::Type{<:FDPrepared}) = DerivativeOrder{1}()
+AbstractPPL.dimension(p::FDPrepared) = AbstractPPL.dimension(p.evaluator)
 
-function AbstractPPL.dimension(::FDPrepared{<:Any,<:Any,<:NamedTuple})
-    throw(
-        ArgumentError(
-            "`dimension` is only available for evaluators prepared with a vector of floating-point numbers.",
-        ),
-    )
-end
-function AbstractPPL.dimension(p::FDPrepared{<:Any,<:Any,<:AbstractVector})
-    return length(p.inputspec)
-end
-
-function (p::FDPrepared{<:Any,<:Any,<:NamedTuple})(values::NamedTuple)
-    typeof(values) === typeof(p.inputspec) || throw(
+function (p::FDPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator})(values::NamedTuple)
+    typeof(values) === typeof(p.evaluator.inputspec) || throw(
         ArgumentError(
             "Expected the same NamedTuple structure that was used to prepare this evaluator.",
         ),
     )
     return p.evaluator(values)
 end
 
-function (p::FDPrepared{<:Any,<:Any,<:NamedTuple})(x::AbstractVector)
-    return p.f(x)
-end
-
-function (p::FDPrepared{<:Any,<:Any,<:AbstractVector})(x::AbstractVector{<:AbstractFloat})
-    length(x) == length(p.inputspec) || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(length(p.inputspec)), but got length $(length(x)).",
-        ),
-    )
-    return p.evaluator(x)
-end
-
 function (p::FDPrepared)(x)
-    throw(MethodError(p, (x,)))
+    return p.evaluator(x)
 end
 
 function AbstractPPL.prepare(adtype::AutoFiniteDifferences, problem, values::NamedTuple)
-    evaluator = AbstractPPL.prepare(problem, values)
+    evaluator = AbstractPPL.ADProblems.NamedTupleEvaluator(
+        AbstractPPL.prepare(problem, values), values
+    )
     f = x -> evaluator(unflatten_to!!(values, x))
-    return FDPrepared(evaluator, f, adtype.fdm, values)
+    return FDPrepared(evaluator, f, adtype.fdm)
 end
 
 function AbstractPPL.prepare(
     adtype::AutoFiniteDifferences, problem, x::AbstractVector{<:AbstractFloat}
 )
-    evaluator = AbstractPPL.prepare(problem, x)
-    return FDPrepared(evaluator, evaluator, adtype.fdm, x)
+    evaluator = AbstractPPL.ADProblems.VectorEvaluator(
+        AbstractPPL.prepare(problem, x), length(x)
+    )
+    return FDPrepared(evaluator, evaluator, adtype.fdm)
 end
 
 function AbstractPPL.value_and_gradient(
-    p::FDPrepared{<:Any,<:Any,<:NamedTuple}, values::NamedTuple
+    p::FDPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator}, values::NamedTuple
 )
-    typeof(values) === typeof(p.inputspec) || throw(
+    typeof(values) === typeof(p.evaluator.inputspec) || throw(
         ArgumentError(
             "Expected the same NamedTuple structure that was used to prepare this evaluator.",
         ),
     )
     x = flatten_to!!(nothing, values)
     val = p.evaluator(values)
     grad = FiniteDifferences.grad(p.fdm, p.f, x)[1]
-    return (val, unflatten_to!!(p.inputspec, grad))
+    return (val, unflatten_to!!(p.evaluator.inputspec, grad))
 end
 
 function AbstractPPL.value_and_gradient(
-    p::FDPrepared{<:Any,<:Any,<:AbstractVector}, x::AbstractVector{<:AbstractFloat}
+    p::FDPrepared{<:AbstractPPL.ADProblems.VectorEvaluator},
+    x::AbstractVector{<:AbstractFloat},
 )
-    length(x) == length(p.inputspec) || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(length(p.inputspec)), but got length $(length(x)).",
-        ),
-    )
     val = p.evaluator(x)
     grad = FiniteDifferences.grad(p.fdm, p.f, x)[1]
     return (val, grad)

ext/AbstractPPLForwardDiffExt.jl

@@ -5,79 +5,62 @@ using AbstractPPL.Utils: flatten_to!!, unflatten_to!!
 using ADTypes: AutoForwardDiff
 using ForwardDiff: ForwardDiff
 
-struct ForwardDiffPrepared{F,C,R,P}
-    evaluator
+struct ForwardDiffPrepared{E,F,C,R}
+    evaluator::E
     f::F
     config::C
     result::R
-    inputspec::P
 end
 
 AbstractPPL.capabilities(::Type{<:ForwardDiffPrepared}) = DerivativeOrder{1}()
+AbstractPPL.dimension(p::ForwardDiffPrepared) = AbstractPPL.dimension(p.evaluator)
 
-function AbstractPPL.dimension(::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:NamedTuple})
-    throw(
-        ArgumentError(
-            "`dimension` is only available for evaluators prepared with a vector of floating-point numbers.",
-        ),
-    )
-end
-function AbstractPPL.dimension(p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:AbstractVector})
-    return length(p.inputspec)
+function (p::ForwardDiffPrepared)(x::AbstractVector{<:Integer})
+    throw(MethodError(p, (x,)))
 end
 
-function (p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:NamedTuple})(values::NamedTuple)
-    typeof(values) === typeof(p.inputspec) || throw(
+function (p::ForwardDiffPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator})(
+    values::NamedTuple
+)
+    typeof(values) === typeof(p.evaluator.inputspec) || throw(
         ArgumentError(
             "Expected the same NamedTuple structure that was used to prepare this evaluator.",
         ),
     )
     return p.evaluator(values)
 end
 
-function (p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:AbstractVector})(
-    x::AbstractVector{<:Integer}
-)
-    throw(MethodError(p, (x,)))
-end
-
-function (p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:AbstractVector})(x::AbstractVector)
-    length(x) == length(p.inputspec) || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(length(p.inputspec)), but got length $(length(x)).",
-        ),
-    )
-    return p.evaluator(x)
-end
-
 function (p::ForwardDiffPrepared)(x)
-    throw(MethodError(p, (x,)))
+    return p.evaluator(x)
 end
 
 function AbstractPPL.prepare(::AutoForwardDiff, problem, values::NamedTuple)
-    evaluator = AbstractPPL.prepare(problem, values)
+    evaluator = AbstractPPL.ADProblems.NamedTupleEvaluator(
+        AbstractPPL.prepare(problem, values), values
+    )
     x = flatten_to!!(nothing, values)
     f = let evaluator = evaluator, values = values
         x -> evaluator(unflatten_to!!(values, x))
     end
     result = ForwardDiff.DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),))
     cfg = ForwardDiff.GradientConfig(f, x)
-    return ForwardDiffPrepared(evaluator, f, cfg, result, values)
+    return ForwardDiffPrepared(evaluator, f, cfg, result)
 end
 
 function AbstractPPL.prepare(::AutoForwardDiff, problem, x::AbstractVector{<:AbstractFloat})
-    evaluator = AbstractPPL.prepare(problem, x)
-    f = evaluator
-    cfg = ForwardDiff.GradientConfig(f, x)
+    evaluator = AbstractPPL.ADProblems.VectorEvaluator(
+        AbstractPPL.prepare(problem, x), length(x)
+    )
+    cfg = ForwardDiff.GradientConfig(evaluator, x)
     grad_buf = similar(x)
     result = ForwardDiff.DiffResults.MutableDiffResult(zero(eltype(x)), (grad_buf,))
-    return ForwardDiffPrepared(evaluator, f, cfg, result, x)
+    return ForwardDiffPrepared(evaluator, evaluator, cfg, result)
 end
 
 @inline function AbstractPPL.value_and_gradient(
-    p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:NamedTuple}, values::NamedTuple
+    p::ForwardDiffPrepared{<:AbstractPPL.ADProblems.NamedTupleEvaluator}, values::NamedTuple
 )
-    typeof(values) === typeof(p.inputspec) || throw(
+    typeof(values) === typeof(p.evaluator.inputspec) || throw(
         ArgumentError(
             "Expected the same NamedTuple structure that was used to prepare this evaluator.",
         ),
@@ -86,18 +69,13 @@ end
     ForwardDiff.gradient!(p.result, p.f, x, p.config)
     val = ForwardDiff.DiffResults.value(p.result)
     grad = copy(ForwardDiff.DiffResults.gradient(p.result))
-    return (val, unflatten_to!!(p.inputspec, grad))
+    return (val, unflatten_to!!(p.evaluator.inputspec, grad))
 end
 
 @inline function AbstractPPL.value_and_gradient(
-    p::ForwardDiffPrepared{<:Any,<:Any,<:Any,<:AbstractVector},
+    p::ForwardDiffPrepared{<:AbstractPPL.ADProblems.VectorEvaluator},
     x::AbstractVector{<:AbstractFloat},
 )
-    length(x) == length(p.inputspec) || throw(
-        DimensionMismatch(
-            "Expected a vector of length $(length(p.inputspec)), but got length $(length(x)).",
-        ),
-    )
     ForwardDiff.gradient!(p.result, p.f, x, p.config)
     val = ForwardDiff.DiffResults.value(p.result)
     grad = copy(ForwardDiff.DiffResults.gradient(p.result))