fix test_grad

Author: yebai

ext/AbstractPPLFiniteDifferencesExt.jl

@@ -14,6 +14,10 @@ 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/AbstractPPL.jl

@@ -8,7 +8,7 @@ export AbstractProbabilisticProgram,
 export AbstractModelTrace
 
 # Evaluator interface
-export DerivativeOrder, capabilities, prepare, value_and_gradient, dimension
+export DerivativeOrder, capabilities, prepare, value_and_gradient, test_grad, dimension
 
 include("abstractmodeltrace.jl")
 include("abstractprobprog.jl")

src/evaluator.jl

@@ -62,6 +62,21 @@ function value_and_gradient(prepared, x::AbstractVector{<:AbstractFloat})
     )
 end
 
+"""
+    test_grad(f, x::AbstractVector{<:AbstractFloat})
+
+Return a finite-difference reference gradient for a scalar-valued callable `f`
+evaluated at the vector input `x`.
+
+If the FiniteDifferences extension is not loaded, this warns and returns `nothing`.
+"""
+function test_grad end
+
+function test_grad(f, x)
+    @warn "Finite-difference reference gradients require `using FiniteDifferences`; skipping test_grad."
+    return nothing
+end
+
 """
     dimension(prepared)::Int
 

test/test_utils.jl

@@ -5,29 +5,15 @@ Include this file inside `@testset` blocks in `test/ext/*/` tests after loading
 AbstractPPL and Test.
 """
 
-function _fd_gradient(f, x::AbstractVector)
-    T = float(eltype(x))
-    h = cbrt(eps(T))
-    grad = Vector{T}(undef, length(x))
-    for i in eachindex(x)
-        xp = copy(x)
-        xp[i] += h
-        xm = copy(x)
-        xm[i] -= h
-        grad[i] = (f(xp) - f(xm)) / (2h)
-    end
-    return grad
-end
-
 """
     test_autograd(prepared, x::AbstractVector; atol=1e-5, rtol=1e-5)
 
-Compare `value_and_gradient(prepared, x)` against a central finite-difference
+Compare `value_and_gradient(prepared, x)` against a finite-difference
 reference. Calls `@test` internally; use inside a `@testset` block.
 """
 function test_autograd(prepared, x::AbstractVector; atol=1e-5, rtol=1e-5)
     val_ad, grad_ad = AbstractPPL.value_and_gradient(prepared, x)
-    grad_fd = _fd_gradient(prepared, x)
+    grad_fd = AbstractPPL.test_grad(prepared, x)
     @test val_ad ≈ prepared(x)
-    @test grad_ad ≈ grad_fd atol = atol rtol = rtol
+    return isnothing(grad_fd) || @test grad_ad ≈ grad_fd atol = atol rtol = rtol
 end