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

This moves the evaluator surface into a self-contained ADProblems module and replaces the old finite-difference test helper with test_autograd backed by AutoFiniteDifferences.

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

Author: yebai

AGENTS.md

@@ -16,7 +16,7 @@ AbstractPPL.jl is a Julia interface package for probabilistic programming. It is
 
   - Preserve the model invariants documented in `src/abstractprobprog.jl`: `condition`/`decondition` and `fix`/`unfix` are intended to round-trip when supported.
   - `rand(model)` and `predict(model, params)` have default RNG/type forwarding behaviour covered by tests; changes here should stay consistent with `AbstractMCMC` expectations.
-  - The evaluator API in `src/evaluator.jl` is structural. `prepare(..., prototype::NamedTuple)` fixes field structure, `capabilities` defaults conservatively to `DerivativeOrder{0}()`, and AD-aware prepared objects are expected to return gradients with the same named structure as inputs.
+  - The ADProblem API in `src/ADProblems.jl` is structural. `prepare(..., prototype::NamedTuple)` fixes field structure, `capabilities` defaults conservatively to `DerivativeOrder{0}()`, and AD-aware prepared objects are expected to return gradients with the same named structure as inputs.
   - `VarName` and optics are the main complexity in this repo. Preserve equality, hashing, pretty-printing, composition/decomposition, and type-stability behaviour.
   - Dynamic indices (`begin`, `end`, expressions containing them) are intentionally deferred until `concretize`; do not silently erase that distinction.
   - Unconcretized dynamic indices must not be serialised. If serialization changes, keep `varname_to_string` / `string_to_varname` round-tripping for supported index types.

docs/src/pplapi.md

@@ -19,7 +19,7 @@ evaluate!!
 AbstractModelTrace
 ```
 
-## Evaluator interface
+## ADProblem interface
 
 ```@docs
 DerivativeOrder

ext/AbstractPPLDifferentiationInterfaceExt.jl

@@ -25,7 +25,7 @@ end
 
 # This extension handles the generic `AbstractADType` vector path directly so
 # DifferentiationInterface backends can opt in without a fallback method in
-# `src/evaluator.jl` forcing a precompile-time method overwrite.
+# `src/ADProblems.jl` forcing a precompile-time method overwrite.
 function AbstractPPL.prepare(
     adtype::ADTypes.AbstractADType, problem, x::AbstractVector{<:AbstractFloat}
 )

ext/AbstractPPLFiniteDifferencesExt.jl

@@ -14,10 +14,6 @@ end
 
 AbstractPPL.capabilities(::Type{<:FDPrepared}) = DerivativeOrder{1}()
 
-function AbstractPPL.test_grad(f, x::AbstractVector{<:AbstractFloat})
-    return FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), f, x)[1]
-end
-
 function AbstractPPL.dimension(::FDPrepared{<:Any,<:Any,<:NamedTuple})
     throw(
         ArgumentError(

src/evaluator.jl src/ADProblems.jlsrc/evaluator.jl renamed to src/ADProblems.jl

@@ -1,5 +1,9 @@
+module ADProblems
+
 using ADTypes: ADTypes
 
+export DerivativeOrder, capabilities, prepare, value_and_gradient, test_autograd, dimension
+
 """
     DerivativeOrder{K}
 
@@ -63,17 +67,48 @@ function value_and_gradient(prepared, x::AbstractVector{<:AbstractFloat})
 end
 
 """
-    test_grad(f, x::AbstractVector{<:AbstractFloat})
+    test_autograd(prepared, x::AbstractVector; atol=1e-5, rtol=1e-5, finite_difference_kwargs...)
 
-Return a finite-difference reference gradient for a scalar-valued callable `f`
-evaluated at the vector input `x`.
+Compare `value_and_gradient(prepared, x)` against a finite-difference reference
+computed via `value_and_gradient(prepare(AutoFiniteDifferences(...), problem, x), x)`.
+Throws an informative error on mismatch. Returns `nothing`.
 
-If the FiniteDifferences extension is not loaded, this warns and returns `nothing`.
+Backends that want this helper should define `prepare_for_test_autograd(prepared, x)`
+to return a pair `(problem, prototype)` suitable for `prepare(AutoFiniteDifferences(...), ...)`.
+Additional keyword arguments are forwarded to `ADTypes.AutoFiniteDifferences`.
 """
-function test_grad end
+function test_autograd end
+
+function prepare_for_test_autograd end
+
+function prepare_for_test_autograd(prepared, x)
+    throw(
+        ArgumentError(
+            "`test_autograd` needs a finite-difference preparation path for $(typeof(prepared)). Define `prepare_for_test_autograd(prepared, x)` to return `(problem, prototype)`.",
+        ),
+    )
+end
 
-function test_grad(f, x)
-    @warn "Finite-difference reference gradients require `using FiniteDifferences`; skipping test_grad."
+function test_autograd(
+    prepared, x::AbstractVector; atol=1e-5, rtol=1e-5, finite_difference_kwargs...
+)
+    val_ad, grad_ad = value_and_gradient(prepared, x)
+    problem, prototype = prepare_for_test_autograd(prepared, x)
+    fd_prepared = prepare(
+        ADTypes.AutoFiniteDifferences(; finite_difference_kwargs...), problem, prototype
+    )
+    val_fd, grad_fd = value_and_gradient(fd_prepared, x)
+
+    isapprox(val_ad, val_fd) || throw(
+        ArgumentError(
+            "Value mismatch against finite differences: got $val_ad, expected $val_fd."
+        ),
+    )
+    isapprox(grad_ad, grad_fd; atol=atol, rtol=rtol) || throw(
+        ArgumentError(
+            "Gradient mismatch against finite differences with atol=$atol and rtol=$rtol.",
+        ),
+    )
     return nothing
 end
 
@@ -83,3 +118,5 @@ end
 Return the number of scalar entries in the vector input expected by a prepared evaluator.
 """
 function dimension end
+
+end # module

src/AbstractPPL.jl

@@ -7,13 +7,14 @@ export AbstractProbabilisticProgram,
 # Abstract traces
 export AbstractModelTrace
 
-# Evaluator interface
-export DerivativeOrder, capabilities, prepare, value_and_gradient, test_grad, dimension
-
+# ADProblem interface
 include("abstractmodeltrace.jl")
 include("abstractprobprog.jl")
 include("evaluate.jl")
-include("evaluator.jl")
+include("ADProblems.jl")
+using .ADProblems:
+    DerivativeOrder, capabilities, prepare, value_and_gradient, test_autograd, dimension
+export DerivativeOrder, capabilities, prepare, value_and_gradient, test_autograd, dimension
 include("utils.jl")
 
 include("varname/optic.jl")

test/ADProblems.jl

@@ -0,0 +1,127 @@
+using AbstractPPL
+using ADTypes: ADTypes
+using Test
+
+struct DummyProblem end
+
+struct DummyPrepared
+    prototype_keys::Tuple
+end
+
+function AbstractPPL.prepare(problem::DummyProblem, values::NamedTuple)
+    return DummyPrepared(keys(values))
+end
+
+function (p::DummyPrepared)(values::NamedTuple)
+    keys(values) == p.prototype_keys ||
+        error("expected fields $(p.prototype_keys), got $(keys(values))")
+    return sum(x -> x isa AbstractArray ? sum(x) : x, values)
+end
+
+struct DummyADPrepared
+    dim::Int
+end
+
+function AbstractPPL.prepare(
+    ::ADTypes.AbstractADType, problem::DummyProblem, x::AbstractVector{<:AbstractFloat}
+)
+    return DummyADPrepared(length(x))
+end
+
+function (p::DummyADPrepared)(x::AbstractVector{<:AbstractFloat})
+    length(x) == p.dim || error("expected vector of length $(p.dim)")
+    return sum(x)
+end
+
+AbstractPPL.capabilities(::Type{DummyADPrepared}) = DerivativeOrder{1}()
+
+function AbstractPPL.value_and_gradient(
+    p::DummyADPrepared, x::AbstractVector{<:AbstractFloat}
+)
+    return (sum(x), ones(length(x)))
+end
+
+struct DummyVectorPrepared
+    dim::Int
+end
+
+AbstractPPL.dimension(p::DummyVectorPrepared) = p.dim
+
+function (p::DummyVectorPrepared)(x::AbstractVector)
+    length(x) == p.dim || error("expected vector of length $(p.dim)")
+    return sum(x)
+end
+
+@testset "ADProblem interface" begin
+    @testset "DerivativeOrder" begin
+        err = try
+            DerivativeOrder{3}()
+            nothing
+        catch err
+            err
+        end
+        @test err isa ArgumentError
+        @test occursin("must be 0, 1, or 2", sprint(showerror, err))
+        @test_throws ArgumentError DerivativeOrder{-1}()
+        @test DerivativeOrder{0}() < DerivativeOrder{1}()
+        @test DerivativeOrder{1}() >= DerivativeOrder{1}()
+        @test DerivativeOrder{1}() < DerivativeOrder{2}()
+        @test !(DerivativeOrder{2}() < DerivativeOrder{1}())
+    end
+
+    @testset "capabilities default" begin
+        @test capabilities(Int) == DerivativeOrder{0}()
+        @test capabilities(42) == DerivativeOrder{0}()
+        @test capabilities(DummyPrepared((:x,))) == DerivativeOrder{0}()
+        @test capabilities(DummyPrepared((:x,))) < DerivativeOrder{1}()
+    end
+
+    @testset "prepare (structural)" begin
+        problem = DummyProblem()
+        values = (x=0.0, y=[1.0, 2.0])
+        prepared = prepare(problem, values)
+        @test prepared isa DummyPrepared
+        @test prepared.prototype_keys == (:x, :y)
+
+        lp = prepared((x=0.5, y=[1.5, 2.5]))
+        @test lp ≈ 0.5 + 1.5 + 2.5
+
+        @test_throws Exception prepared((a=1.0, b=2.0))
+    end
+
+    @testset "prepare (AD-aware)" begin
+        problem = DummyProblem()
+        x0 = zeros(3)
+        adtype = ADTypes.AutoForwardDiff()
+        prepared = prepare(adtype, problem, x0)
+        @test prepared isa DummyADPrepared
+        @test capabilities(prepared) == DerivativeOrder{1}()
+
+        x = [0.5, 1.5, 2.5]
+        @test prepared(x) ≈ 0.5 + 1.5 + 2.5
+
+        val, grad = value_and_gradient(prepared, x)
+        @test val ≈ 0.5 + 1.5 + 2.5
+        @test grad ≈ [1.0, 1.0, 1.0]
+    end
+
+    @testset "dimension and vector adapter" begin
+        prepared = DummyVectorPrepared(3)
+        @test dimension(prepared) == 3
+        @test prepared(ones(3)) ≈ 3.0
+        @test_throws Exception prepared(ones(5))
+    end
+
+    @testset "flatten / unflatten edge cases" begin
+        empty = NamedTuple()
+        @test AbstractPPL.Utils.flatten_to!!(nothing, empty) == Float64[]
+        @test AbstractPPL.Utils.unflatten_to!!(empty, Float64[]) == empty
+
+        view_values = (x=@view([1.0, 2.0, 3.0][2:3]),)
+        flat = AbstractPPL.Utils.flatten_to!!(nothing, view_values)
+        rebuilt = AbstractPPL.Utils.unflatten_to!!(view_values, flat)
+        @test collect(rebuilt.x) == [2.0, 3.0]
+        @test axes(rebuilt.x) == axes(view_values.x)
+        @test parent(rebuilt.x) == [2.0, 3.0]
+    end
+end

test/ext/differentiation_interface/differentiation_interface.jl

@@ -23,6 +23,12 @@ function (::QuadraticPrepared)(x::AbstractVector{<:AbstractFloat})
     return sum(xi -> xi^2, x)
 end
 
+function AbstractPPL.ADProblems.prepare_for_test_autograd(
+    ::AbstractPPL.DIPrepared, x::AbstractVector
+)
+    return (QuadraticProblem(), x)
+end
+
 # Use a backend without a native AbstractPPL extension so this test exercises
 # AbstractPPLDifferentiationInterfaceExt dispatch directly.
 const fdm = FiniteDifferences.central_fdm(5, 1)

test/ext/enzyme/enzyme.jl

@@ -21,6 +21,12 @@ function (::QuadraticPrepared)(x::AbstractVector{<:AbstractFloat})
     return sum(xi -> xi^2, x)
 end
 
+function AbstractPPL.ADProblems.prepare_for_test_autograd(
+    ::AbstractPPL.EnzymePrepared, x::AbstractVector
+)
+    return (QuadraticProblem(), x)
+end
+
 @testset "AbstractPPLEnzymeExt" begin
     problem = QuadraticProblem()
     x0 = zeros(3)

test/ext/forward_diff/forward_diff.jl

@@ -30,6 +30,12 @@ function (::QuadraticVecPrepared)(x::AbstractVector{<:Real})
     return sum(xi -> xi^2, x)
 end
 
+function AbstractPPL.ADProblems.prepare_for_test_autograd(
+    ::AbstractPPL.ForwardDiffPrepared{<:Any,<:Any,<:Any,<:AbstractVector}, x::AbstractVector
+)
+    return (QuadraticProblem(), x)
+end
+
 @testset "AbstractPPLForwardDiffExt" begin
     @testset "NamedTuple path" begin
         problem = QuadraticProblem()