From 6baeca3a3814cf383072d3317ea2f9cffd277101 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sat, 2 May 2026 10:41:03 +0100 Subject: [PATCH 01/16] Add DifferentiationInterface extension on top of evaluator interface Layer the DifferentiationInterface AD extension onto the simplified evaluator interface introduced on the `evaluators` branch. This is the catch-all path for any `<:AbstractADType` backend without a native AbstractPPL extension. - ext/AbstractPPLDifferentiationInterfaceExt.jl: DI-backed `prepare`, `value_and_gradient!!`, and `value_and_jacobian!!` for vector-input evaluators. The evaluator is passed as a `DI.Constant` so DynamicPPL-style problem state stays constant across calls. Compiled `AutoReverseDiff{true}` takes the one-argument tape path (the `DI.Constant` route would invalidate the compiled tape across calls). Empty inputs short-circuit before `DI.prepare_*` since many backends fail on length-zero arrays. - Project.toml: DifferentiationInterface added as a weakdep with extension trigger and compat bound. - test/autograd_tests.jl: shared problem definitions and `run_autograd_tests` entry point with separate gradient / jacobian / empty-input helpers. - test/ext/differentiationinterface/: isolated test env using a local `DummyADType <: AbstractADType` to exercise the catch-all dispatch without pulling in a real AD package. - .github/workflows/CI.yml: extend the ext matrix with `ext/differentiationinterface` alongside the existing `ext/logdensityproblems`. - test/run_extras.jl: register the new label. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/CI.yml | 7 +- Project.toml | 3 + ext/AbstractPPLDifferentiationInterfaceExt.jl | 92 +++++++++++++++++ test/autograd_tests.jl | 98 +++++++++++++++++++ .../ext/differentiationinterface/Project.toml | 11 +++ test/ext/differentiationinterface/main.jl | 34 +++++++ test/run_extras.jl | 5 +- 7 files changed, 246 insertions(+), 4 deletions(-) create mode 100644 ext/AbstractPPLDifferentiationInterfaceExt.jl create mode 100644 test/autograd_tests.jl create mode 100644 test/ext/differentiationinterface/Project.toml create mode 100644 test/ext/differentiationinterface/main.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 60609e16..17ae718f 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -54,11 +54,14 @@ jobs: fail_ci_if_error: true ext: - name: Ext (logdensityproblems, ${{ matrix.version }}) + name: Ext (${{ matrix.label }}, ${{ matrix.version }}) runs-on: ubuntu-latest strategy: fail-fast: false matrix: + label: + - ext/differentiationinterface + - ext/logdensityproblems version: - '1' - 'min' @@ -71,4 +74,4 @@ jobs: - uses: julia-actions/julia-buildpkg@v1 - run: julia --project=. test/run_extras.jl env: - LABEL: ext/logdensityproblems + LABEL: ${{ matrix.label }} diff --git a/Project.toml b/Project.toml index 3678edc6..8b007eef 100644 --- a/Project.toml +++ b/Project.toml @@ -19,15 +19,18 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" [extensions] +AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] [compat] ADTypes = "1" +DifferentiationInterface = "0.6, 0.7" LogDensityProblems = "2" AbstractMCMC = "2, 3, 4, 5" Accessors = "0.1" diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl new file mode 100644 index 00000000..73052f26 --- /dev/null +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -0,0 +1,92 @@ +module AbstractPPLDifferentiationInterfaceExt + +using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: + Prepared, VectorEvaluator, _assert_jacobian_output, _assert_supported_output +using ADTypes: ADTypes, AbstractADType, AutoReverseDiff +using DifferentiationInterface: DifferentiationInterface as DI + +# Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so +# that in DynamicPPL the model and other evaluator state stay constant. +@inline _call_evaluator(x, evaluator) = evaluator(x) + +struct DICache{F,GP,JP} + target::F + gradient_prep::GP + jacobian_prep::JP + use_context::Bool +end + +# Compiled ReverseDiff only reuses a compiled tape on the one-argument path; +# the DI.Constant route would invalidate the tape across calls. +function _prepare_gradient(adtype::AutoReverseDiff{true}, x, evaluator) + target = Base.Fix2(_call_evaluator, evaluator) + gradient_prep = DI.prepare_gradient(target, adtype, x) + return target, gradient_prep, false +end + +function _prepare_gradient(adtype::AbstractADType, x, evaluator) + target = _call_evaluator + gradient_prep = DI.prepare_gradient(target, adtype, x, DI.Constant(evaluator)) + return target, gradient_prep, true +end + +function AbstractPPL.prepare( + adtype::AbstractADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true +) + raw = AbstractPPL.prepare(problem, x) + evaluator = VectorEvaluator{check_dims}(raw, length(x)) + # Empty inputs bypass `DI.prepare_*` (many backends fail on length-zero arrays); + # `value_and_gradient!!` / `value_and_jacobian!!` short-circuit on `length(x) == 0`. + length(x) == 0 && + return Prepared(adtype, evaluator, DICache(_call_evaluator, nothing, nothing, true)) + y = evaluator(x) + _assert_supported_output(y) + if y isa Number + target, gradient_prep, use_context = _prepare_gradient(adtype, x, evaluator) + return Prepared( + adtype, evaluator, DICache(target, gradient_prep, nothing, use_context) + ) + else + _assert_jacobian_output(y) + jacobian_prep = DI.prepare_jacobian( + _call_evaluator, adtype, x, DI.Constant(evaluator) + ) + return Prepared( + adtype, evaluator, DICache(_call_evaluator, nothing, jacobian_prep, true) + ) + end +end + +@inline function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T} +) where {T<:Real} + length(x) == 0 && return (p.evaluator(x), T[]) + p.cache.gradient_prep === nothing && + throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function.")) + val, grad = if p.cache.use_context + DI.value_and_gradient( + p.cache.target, p.cache.gradient_prep, p.adtype, x, DI.Constant(p.evaluator) + ) + else + DI.value_and_gradient(p.cache.target, p.cache.gradient_prep, p.adtype, x) + end + # Some DI backends may return a non-`Vector` gradient; normalise. + return (val, grad isa Vector ? grad : collect(grad)) +end + +@inline function AbstractPPL.value_and_jacobian!!( + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{<:Real} +) + if length(x) == 0 + val = p.evaluator(x) + return (val, similar(x, length(val), 0)) + end + p.cache.jacobian_prep === nothing && + throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) + return DI.value_and_jacobian( + p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator) + ) +end + +end # module diff --git a/test/autograd_tests.jl b/test/autograd_tests.jl new file mode 100644 index 00000000..88b6845c --- /dev/null +++ b/test/autograd_tests.jl @@ -0,0 +1,98 @@ +# Shared problem definitions and test helpers for AD backend integration tests. +# Include this file after `using AbstractPPL, Test` and any backend-specific setup. + +struct QuadraticProblem end +struct QuadraticVecPrepared end + +function AbstractPPL.prepare(::QuadraticProblem, x::AbstractVector{<:Real}) + return QuadraticVecPrepared() +end + +(::QuadraticVecPrepared)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) + +struct VectorValuedProblem end +struct VectorValuedPrepared end + +function AbstractPPL.prepare(::VectorValuedProblem, x::AbstractVector{<:Real}) + return VectorValuedPrepared() +end + +# y = [x[1]*x[2], x[2]+x[3]] -> J = [x[2] x[1] 0; 0 1 1] +(::VectorValuedPrepared)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] + +""" + run_shared_gradient_tests(adtype, x0, x; atol=0, rtol=1e-10) + +Test the vector-input gradient path for `adtype` on `QuadraticProblem`. +`x0` is the prototype (zeros), `x = [3.0, 1.0, 2.0]` is the test point. +""" +function run_shared_gradient_tests(adtype, x0, x; atol=0, rtol=1e-10) + @testset "gradient path" begin + problem = QuadraticProblem() + prepared = AbstractPPL.prepare(adtype, problem, x0) + + @test prepared(x) ≈ 14.0 + + val, grad = AbstractPPL.value_and_gradient!!(prepared, x) + @test val ≈ 14.0 atol = atol rtol = rtol + @test grad ≈ [6.0, 2.0, 4.0] atol = atol rtol = rtol + + @test_throws DimensionMismatch prepared([3.0, 1.0, 2.0, 99.0]) + @test_throws r"floating-point" prepared([3, 1, 2]) + end +end + +""" + run_shared_jacobian_tests(adtype, x0, xj; atol=0, rtol=1e-10) + +Test the jacobian path for `adtype` on `VectorValuedProblem`. +`x0` is the prototype (zeros(3)), `xj` is the test point. +""" +function run_shared_jacobian_tests(adtype, x0, xj; atol=0, rtol=1e-10) + @testset "jacobian path" begin + problem = VectorValuedProblem() + prepared = AbstractPPL.prepare(adtype, problem, x0) + + @test prepared(xj) ≈ [6.0, 7.0] + + val, jac = AbstractPPL.value_and_jacobian!!(prepared, xj) + @test val ≈ [6.0, 7.0] atol = atol rtol = rtol + @test jac ≈ [3.0 2.0 0.0; 0.0 1.0 1.0] atol = atol rtol = rtol + + @test_throws r"scalar-valued" AbstractPPL.value_and_gradient!!(prepared, xj) + end +end + +""" + run_shared_empty_input_tests(adtype) + +Test the empty-input short-circuit for `adtype` on both scalar- and +vector-valued evaluators. +""" +function run_shared_empty_input_tests(adtype) + @testset "empty input" begin + x_empty = Float64[] + prepared = AbstractPPL.prepare(adtype, x -> 7.5, x_empty) + val, grad = AbstractPPL.value_and_gradient!!(prepared, x_empty) + @test val == 7.5 + @test grad == Float64[] + + prepared_jac = AbstractPPL.prepare(adtype, x -> [2.0, 3.0], x_empty) + valj, jac = AbstractPPL.value_and_jacobian!!(prepared_jac, x_empty) + @test valj == [2.0, 3.0] + @test size(jac) == (2, 0) + end +end + +""" + run_autograd_tests(adtype; kwargs...) + +Run the gradient, jacobian, and empty-input shared tests on `adtype`. `kwargs` +(`atol`, `rtol`, …) are forwarded to the shared helpers. +""" +function run_autograd_tests(adtype; kwargs...) + run_shared_gradient_tests(adtype, zeros(3), [3.0, 1.0, 2.0]; kwargs...) + run_shared_jacobian_tests(adtype, zeros(3), [2.0, 3.0, 4.0]; kwargs...) + run_shared_empty_input_tests(adtype) + return nothing +end diff --git a/test/ext/differentiationinterface/Project.toml b/test/ext/differentiationinterface/Project.toml new file mode 100644 index 00000000..90b65758 --- /dev/null +++ b/test/ext/differentiationinterface/Project.toml @@ -0,0 +1,11 @@ +[deps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[compat] +ADTypes = "1" +DifferentiationInterface = "0.6, 0.7" +julia = "1.10" diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl new file mode 100644 index 00000000..a7bb6217 --- /dev/null +++ b/test/ext/differentiationinterface/main.jl @@ -0,0 +1,34 @@ +using Pkg +Pkg.activate(@__DIR__) +Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) +Pkg.instantiate() + +using AbstractPPL +using ADTypes: ADTypes +using DifferentiationInterface: DifferentiationInterface as DI +using Test + +include(joinpath(@__DIR__, "..", "..", "autograd_tests.jl")) + +# Stub backend without a native AbstractPPL extension; this exercises the +# `AbstractPPLDifferentiationInterfaceExt` catch-all dispatch on +# `<:AbstractADType` rather than depending on a real AD package. +struct DummyADType <: ADTypes.AbstractADType end +const adtype = DummyADType() + +DI.prepare_gradient(f, ::DummyADType, x, ::DI.Constant) = Val(:gradient) +function DI.value_and_gradient(f, prep, ::DummyADType, x, ctx::DI.Constant) + return (f(x, ctx.data), 2 .* x) +end +DI.prepare_jacobian(f, ::DummyADType, x, ::DI.Constant) = Val(:jacobian) +function DI.value_and_jacobian(f, prep, ::DummyADType, x, ctx::DI.Constant) + jac = [ + x[2] x[1] zero(eltype(x)) + zero(eltype(x)) one(eltype(x)) one(eltype(x)) + ] + return (f(x, ctx.data), jac) +end + +@testset "AbstractPPLDifferentiationInterfaceExt" begin + run_autograd_tests(adtype; atol=1e-6, rtol=1e-6) +end diff --git a/test/run_extras.jl b/test/run_extras.jl index 0e71bbcd..6b43493d 100644 --- a/test/run_extras.jl +++ b/test/run_extras.jl @@ -1,9 +1,10 @@ # Run a named extension test in its own isolated Julia environment. # # Usage (from the repo root): -# LABEL=ext/logdensityproblems julia test/run_extras.jl +# LABEL=ext/differentiationinterface julia test/run_extras.jl +# LABEL=ext/logdensityproblems julia test/run_extras.jl -const VALID_LABELS = ("ext/logdensityproblems",) +const VALID_LABELS = ("ext/differentiationinterface", "ext/logdensityproblems") label = get(ENV, "LABEL", nothing) label in VALID_LABELS || From 9c371b3703538369169079a72243b139b4daab40 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 4 May 2026 18:13:54 +0100 Subject: [PATCH 02/16] Inline AD output-shape assertions into the DI extension The upstream `Simplify Evaluators module after review` commit dropped `_assert_supported_output` / `_assert_jacobian_output` from `Evaluators.jl`, expecting the AD-extension PRs that actually call them to bring them back. Move both helpers into the DI extension, the sole caller after rebase. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLDifferentiationInterfaceExt.jl | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 73052f26..57e25e0b 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -1,11 +1,28 @@ module AbstractPPLDifferentiationInterfaceExt using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: - Prepared, VectorEvaluator, _assert_jacobian_output, _assert_supported_output +using AbstractPPL.Evaluators: Prepared, VectorEvaluator using ADTypes: ADTypes, AbstractADType, AutoReverseDiff using DifferentiationInterface: DifferentiationInterface as DI +function _assert_supported_output(y) + (y isa Number || y isa AbstractVector) || throw( + ArgumentError( + "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", + ), + ) + return nothing +end + +function _assert_jacobian_output(y) + y isa AbstractVector || throw( + ArgumentError( + "`value_and_jacobian!!` requires the prepared function to return an AbstractVector; got $(typeof(y)).", + ), + ) + return nothing +end + # Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so # that in DynamicPPL the model and other evaluator state stay constant. @inline _call_evaluator(x, evaluator) = evaluator(x) From 6bc564a5bb62244d3c25431795e4eaae10117853 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 4 May 2026 18:16:20 +0100 Subject: [PATCH 03/16] Fold AD output-shape checks into prepare's dispatch Replace `_assert_supported_output` / `_assert_jacobian_output` with an `if y isa Number / elseif y isa AbstractVector / else throw` cascade. The jacobian assertion was unreachable after the scalar branch, and the supported-output helper had a single call site; inlining the remaining check as the `else` arm is clearer than two named helpers. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLDifferentiationInterfaceExt.jl | 28 +++++-------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 57e25e0b..b45a9d63 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -5,24 +5,6 @@ using AbstractPPL.Evaluators: Prepared, VectorEvaluator using ADTypes: ADTypes, AbstractADType, AutoReverseDiff using DifferentiationInterface: DifferentiationInterface as DI -function _assert_supported_output(y) - (y isa Number || y isa AbstractVector) || throw( - ArgumentError( - "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", - ), - ) - return nothing -end - -function _assert_jacobian_output(y) - y isa AbstractVector || throw( - ArgumentError( - "`value_and_jacobian!!` requires the prepared function to return an AbstractVector; got $(typeof(y)).", - ), - ) - return nothing -end - # Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so # that in DynamicPPL the model and other evaluator state stay constant. @inline _call_evaluator(x, evaluator) = evaluator(x) @@ -58,20 +40,24 @@ function AbstractPPL.prepare( length(x) == 0 && return Prepared(adtype, evaluator, DICache(_call_evaluator, nothing, nothing, true)) y = evaluator(x) - _assert_supported_output(y) if y isa Number target, gradient_prep, use_context = _prepare_gradient(adtype, x, evaluator) return Prepared( adtype, evaluator, DICache(target, gradient_prep, nothing, use_context) ) - else - _assert_jacobian_output(y) + elseif y isa AbstractVector jacobian_prep = DI.prepare_jacobian( _call_evaluator, adtype, x, DI.Constant(evaluator) ) return Prepared( adtype, evaluator, DICache(_call_evaluator, nothing, jacobian_prep, true) ) + else + throw( + ArgumentError( + "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", + ), + ) end end From c4a6e3c774db1d9ba23efcd5dd4adc16d66d7987 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 4 May 2026 20:20:32 +0100 Subject: [PATCH 04/16] Move AD test cases into package test resources Expose shared AD backend fixtures through TestResources so extension test environments can reuse the same cases without including files from test/. Co-Authored-By: Claude Opus 4.7 --- src/AbstractPPL.jl | 1 + src/test_resources.jl | 139 ++++++++++++++++++++++ test/autograd_tests.jl | 98 --------------- test/ext/differentiationinterface/main.jl | 38 +++++- 4 files changed, 175 insertions(+), 101 deletions(-) create mode 100644 src/test_resources.jl delete mode 100644 test/autograd_tests.jl diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 48a2cb44..567cedb2 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -12,6 +12,7 @@ include("abstractprobprog.jl") include("evaluate.jl") include("evaluators/Evaluators.jl") using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!! +include("test_resources.jl") @static if VERSION >= v"1.11.0" eval(Meta.parse("public prepare, value_and_gradient!!, value_and_jacobian!!")) end diff --git a/src/test_resources.jl b/src/test_resources.jl new file mode 100644 index 00000000..6e26dd31 --- /dev/null +++ b/src/test_resources.jl @@ -0,0 +1,139 @@ +module TestResources + +import ..AbstractPPL: prepare + +export TestCase, generate_testcases + +struct QuadraticProblem end +struct QuadraticPrepared end +prepare(::QuadraticProblem, ::AbstractVector{<:Real}) = QuadraticPrepared() +(::QuadraticPrepared)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) + +struct VectorValuedProblem end +struct VectorValuedPrepared end +function prepare(::VectorValuedProblem, ::AbstractVector{<:Real}) + return VectorValuedPrepared() +end +(::VectorValuedPrepared)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] + +struct TestCase{F,XP,X,V,G,J,O,E} + name::String + f::F + x_proto::XP + x::X + value::V + gradient::G + jacobian::J + operation::O + exception::E +end + +function TestCase(name, f, x_proto, x, value, gradient, jacobian) + return TestCase(name, f, x_proto, x, value, gradient, jacobian, :value, nothing) +end + +function generate_testcases(::Val{:vector}) + return ( + TestCase( + "quadratic (scalar output)", + QuadraticProblem(), + zeros(3), + [3.0, 1.0, 2.0], + 14.0, + [6.0, 2.0, 4.0], + nothing, + ), + TestCase( + "vector-valued (vector output)", + VectorValuedProblem(), + zeros(3), + [2.0, 3.0, 4.0], + [6.0, 7.0], + nothing, + [3.0 2.0 0.0; 0.0 1.0 1.0], + ), + TestCase( + "empty input, scalar output", + x -> 7.5, + Float64[], + Float64[], + 7.5, + Float64[], + nothing, + ), + TestCase( + "empty input, vector output", + x -> [2.0, 3.0], + Float64[], + Float64[], + [2.0, 3.0], + nothing, + zeros(2, 0), + ), + ) +end + +function generate_testcases(::Val{:namedtuple}) + return ( + TestCase( + "scalar output over (x::Real, y::Vector)", + vs -> vs.x^2 + sum(abs2, vs.y), + (x=0.0, y=zeros(2)), + (x=3.0, y=[1.0, 2.0]), + 14.0, + (x=6.0, y=[2.0, 4.0]), + nothing, + ), + TestCase( + "wrong NamedTuple structure", + vs -> vs.x^2 + sum(abs2, vs.y), + (x=0.0, y=zeros(2)), + (x=3.0, z=[1.0, 2.0]), + nothing, + nothing, + nothing, + :gradient, + r"same NamedTuple structure", + ), + ) +end + +function generate_testcases(::Val{:edge}) + return ( + TestCase( + "wrong vector length", + QuadraticProblem(), + zeros(3), + [3.0, 1.0, 2.0, 99.0], + nothing, + nothing, + nothing, + :call, + DimensionMismatch, + ), + TestCase( + "non-floating-point vector", + QuadraticProblem(), + zeros(3), + [3, 1, 2], + nothing, + nothing, + nothing, + :call, + r"floating-point", + ), + TestCase( + "gradient of vector-valued output", + VectorValuedProblem(), + zeros(3), + [2.0, 3.0, 4.0], + nothing, + nothing, + nothing, + :gradient, + Exception, + ), + ) +end + +end # module diff --git a/test/autograd_tests.jl b/test/autograd_tests.jl deleted file mode 100644 index 88b6845c..00000000 --- a/test/autograd_tests.jl +++ /dev/null @@ -1,98 +0,0 @@ -# Shared problem definitions and test helpers for AD backend integration tests. -# Include this file after `using AbstractPPL, Test` and any backend-specific setup. - -struct QuadraticProblem end -struct QuadraticVecPrepared end - -function AbstractPPL.prepare(::QuadraticProblem, x::AbstractVector{<:Real}) - return QuadraticVecPrepared() -end - -(::QuadraticVecPrepared)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) - -struct VectorValuedProblem end -struct VectorValuedPrepared end - -function AbstractPPL.prepare(::VectorValuedProblem, x::AbstractVector{<:Real}) - return VectorValuedPrepared() -end - -# y = [x[1]*x[2], x[2]+x[3]] -> J = [x[2] x[1] 0; 0 1 1] -(::VectorValuedPrepared)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] - -""" - run_shared_gradient_tests(adtype, x0, x; atol=0, rtol=1e-10) - -Test the vector-input gradient path for `adtype` on `QuadraticProblem`. -`x0` is the prototype (zeros), `x = [3.0, 1.0, 2.0]` is the test point. -""" -function run_shared_gradient_tests(adtype, x0, x; atol=0, rtol=1e-10) - @testset "gradient path" begin - problem = QuadraticProblem() - prepared = AbstractPPL.prepare(adtype, problem, x0) - - @test prepared(x) ≈ 14.0 - - val, grad = AbstractPPL.value_and_gradient!!(prepared, x) - @test val ≈ 14.0 atol = atol rtol = rtol - @test grad ≈ [6.0, 2.0, 4.0] atol = atol rtol = rtol - - @test_throws DimensionMismatch prepared([3.0, 1.0, 2.0, 99.0]) - @test_throws r"floating-point" prepared([3, 1, 2]) - end -end - -""" - run_shared_jacobian_tests(adtype, x0, xj; atol=0, rtol=1e-10) - -Test the jacobian path for `adtype` on `VectorValuedProblem`. -`x0` is the prototype (zeros(3)), `xj` is the test point. -""" -function run_shared_jacobian_tests(adtype, x0, xj; atol=0, rtol=1e-10) - @testset "jacobian path" begin - problem = VectorValuedProblem() - prepared = AbstractPPL.prepare(adtype, problem, x0) - - @test prepared(xj) ≈ [6.0, 7.0] - - val, jac = AbstractPPL.value_and_jacobian!!(prepared, xj) - @test val ≈ [6.0, 7.0] atol = atol rtol = rtol - @test jac ≈ [3.0 2.0 0.0; 0.0 1.0 1.0] atol = atol rtol = rtol - - @test_throws r"scalar-valued" AbstractPPL.value_and_gradient!!(prepared, xj) - end -end - -""" - run_shared_empty_input_tests(adtype) - -Test the empty-input short-circuit for `adtype` on both scalar- and -vector-valued evaluators. -""" -function run_shared_empty_input_tests(adtype) - @testset "empty input" begin - x_empty = Float64[] - prepared = AbstractPPL.prepare(adtype, x -> 7.5, x_empty) - val, grad = AbstractPPL.value_and_gradient!!(prepared, x_empty) - @test val == 7.5 - @test grad == Float64[] - - prepared_jac = AbstractPPL.prepare(adtype, x -> [2.0, 3.0], x_empty) - valj, jac = AbstractPPL.value_and_jacobian!!(prepared_jac, x_empty) - @test valj == [2.0, 3.0] - @test size(jac) == (2, 0) - end -end - -""" - run_autograd_tests(adtype; kwargs...) - -Run the gradient, jacobian, and empty-input shared tests on `adtype`. `kwargs` -(`atol`, `rtol`, …) are forwarded to the shared helpers. -""" -function run_autograd_tests(adtype; kwargs...) - run_shared_gradient_tests(adtype, zeros(3), [3.0, 1.0, 2.0]; kwargs...) - run_shared_jacobian_tests(adtype, zeros(3), [2.0, 3.0, 4.0]; kwargs...) - run_shared_empty_input_tests(adtype) - return nothing -end diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index a7bb6217..f7b1abbb 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -4,12 +4,11 @@ Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() using AbstractPPL +using AbstractPPL.TestResources: generate_testcases using ADTypes: ADTypes using DifferentiationInterface: DifferentiationInterface as DI using Test -include(joinpath(@__DIR__, "..", "..", "autograd_tests.jl")) - # Stub backend without a native AbstractPPL extension; this exercises the # `AbstractPPLDifferentiationInterfaceExt` catch-all dispatch on # `<:AbstractADType` rather than depending on a real AD package. @@ -29,6 +28,39 @@ function DI.value_and_jacobian(f, prep, ::DummyADType, x, ctx::DI.Constant) return (f(x, ctx.data), jac) end +const ATOL = 1e-6 +const RTOL = 1e-6 + @testset "AbstractPPLDifferentiationInterfaceExt" begin - run_autograd_tests(adtype; atol=1e-6, rtol=1e-6) + @testset "vector input" begin + @testset "$(case.name)" for case in generate_testcases(Val(:vector)) + prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) + @test prepared(case.x) ≈ case.value atol = ATOL rtol = RTOL + if case.gradient !== nothing + val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) + @test val ≈ case.value atol = ATOL rtol = RTOL + @test grad ≈ case.gradient atol = ATOL rtol = RTOL + end + if case.jacobian !== nothing + val, jac = AbstractPPL.value_and_jacobian!!(prepared, case.x) + @test val ≈ case.value atol = ATOL rtol = RTOL + @test jac ≈ case.jacobian atol = ATOL rtol = RTOL + end + end + + @testset "edge cases" begin + @testset "$(case.name)" for case in generate_testcases(Val(:edge)) + prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) + if case.operation === :call + @test_throws case.exception prepared(case.x) + elseif case.operation === :gradient + @test_throws case.exception AbstractPPL.value_and_gradient!!( + prepared, case.x + ) + else + error("Unknown edge-test operation: $(case.operation)") + end + end + end + end end From 76d2ea897002a5c3a2babec9b90e4bb7006626d7 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 4 May 2026 20:42:42 +0100 Subject: [PATCH 05/16] Use ForwardDiff in DI extension tests Exercise the DifferentiationInterface catch-all with a real backend instead of a stub so the extension test covers the integration path used by downstream AD packages. Co-Authored-By: Claude Opus 4.7 --- src/test_resources.jl | 1 - .../ext/differentiationinterface/Project.toml | 2 + test/ext/differentiationinterface/main.jl | 76 ++++++++----------- 3 files changed, 33 insertions(+), 46 deletions(-) diff --git a/src/test_resources.jl b/src/test_resources.jl index 6e26dd31..143f9e25 100644 --- a/src/test_resources.jl +++ b/src/test_resources.jl @@ -1,7 +1,6 @@ module TestResources import ..AbstractPPL: prepare - export TestCase, generate_testcases struct QuadraticProblem end diff --git a/test/ext/differentiationinterface/Project.toml b/test/ext/differentiationinterface/Project.toml index 90b65758..2a48e31a 100644 --- a/test/ext/differentiationinterface/Project.toml +++ b/test/ext/differentiationinterface/Project.toml @@ -2,10 +2,12 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] ADTypes = "1" DifferentiationInterface = "0.6, 0.7" +ForwardDiff = "1" julia = "1.10" diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index f7b1abbb..5fb0d1da 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -5,60 +5,46 @@ Pkg.instantiate() using AbstractPPL using AbstractPPL.TestResources: generate_testcases -using ADTypes: ADTypes +using ADTypes: AutoForwardDiff using DifferentiationInterface: DifferentiationInterface as DI +using ForwardDiff using Test -# Stub backend without a native AbstractPPL extension; this exercises the -# `AbstractPPLDifferentiationInterfaceExt` catch-all dispatch on -# `<:AbstractADType` rather than depending on a real AD package. -struct DummyADType <: ADTypes.AbstractADType end -const adtype = DummyADType() - -DI.prepare_gradient(f, ::DummyADType, x, ::DI.Constant) = Val(:gradient) -function DI.value_and_gradient(f, prep, ::DummyADType, x, ctx::DI.Constant) - return (f(x, ctx.data), 2 .* x) -end -DI.prepare_jacobian(f, ::DummyADType, x, ::DI.Constant) = Val(:jacobian) -function DI.value_and_jacobian(f, prep, ::DummyADType, x, ctx::DI.Constant) - jac = [ - x[2] x[1] zero(eltype(x)) - zero(eltype(x)) one(eltype(x)) one(eltype(x)) - ] - return (f(x, ctx.data), jac) -end - +const adtype = AutoForwardDiff() const ATOL = 1e-6 const RTOL = 1e-6 @testset "AbstractPPLDifferentiationInterfaceExt" begin - @testset "vector input" begin - @testset "$(case.name)" for case in generate_testcases(Val(:vector)) - prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) - @test prepared(case.x) ≈ case.value atol = ATOL rtol = RTOL - if case.gradient !== nothing - val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) - @test val ≈ case.value atol = ATOL rtol = RTOL - @test grad ≈ case.gradient atol = ATOL rtol = RTOL - end - if case.jacobian !== nothing - val, jac = AbstractPPL.value_and_jacobian!!(prepared, case.x) - @test val ≈ case.value atol = ATOL rtol = RTOL - @test jac ≈ case.jacobian atol = ATOL rtol = RTOL + # Use a real DI backend to exercise AbstractPPL's catch-all ADType dispatch. + @testset "ForwardDiff" begin + @testset "vector input" begin + @testset "$(case.name)" for case in generate_testcases(Val(:vector)) + prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) + @test prepared(case.x) ≈ case.value atol = ATOL rtol = RTOL + if case.gradient !== nothing + val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) + @test val ≈ case.value atol = ATOL rtol = RTOL + @test grad ≈ case.gradient atol = ATOL rtol = RTOL + end + if case.jacobian !== nothing + val, jac = AbstractPPL.value_and_jacobian!!(prepared, case.x) + @test val ≈ case.value atol = ATOL rtol = RTOL + @test jac ≈ case.jacobian atol = ATOL rtol = RTOL + end end - end - @testset "edge cases" begin - @testset "$(case.name)" for case in generate_testcases(Val(:edge)) - prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) - if case.operation === :call - @test_throws case.exception prepared(case.x) - elseif case.operation === :gradient - @test_throws case.exception AbstractPPL.value_and_gradient!!( - prepared, case.x - ) - else - error("Unknown edge-test operation: $(case.operation)") + @testset "edge cases" begin + @testset "$(case.name)" for case in generate_testcases(Val(:edge)) + prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) + if case.operation === :call + @test_throws case.exception prepared(case.x) + elseif case.operation === :gradient + @test_throws case.exception AbstractPPL.value_and_gradient!!( + prepared, case.x + ) + else + error("Unknown edge-test operation: $(case.operation)") + end end end end From ff96fc10c61c563ca2715db20b26530060766413 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 5 May 2026 16:51:28 +0100 Subject: [PATCH 06/16] Address PR review: simplify AD prepare and move test fixtures to Test extension Replace `prepare(adtype, problem, x)`'s two-step `prepare` + `VectorEvaluator` wrap with a single delegated call to `AbstractPPL.prepare(problem, x; check_dims)`, eliminating a latent double-wrap when `check_dims=false`. Move `TestResources` (callable problems, test cases, runner) out of `src/` and into a new `AbstractPPLTestExt` package extension, so `Test` is loaded only when downstream code opts in. Split the unified `TestCase` into `ValueCase` / `ErrorCase`, drop the unused `:namedtuple` fixtures, and store the error operation as a callable instead of a `:call`/`:gradient` symbol. Flatten the DI extension test file to two `run_testcases(Val(...); ...)` calls. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 3 + ext/AbstractPPLDifferentiationInterfaceExt.jl | 3 +- ext/AbstractPPLTestExt.jl | 133 +++++++++++++++++ src/AbstractPPL.jl | 5 +- src/test_resources.jl | 138 ------------------ test/ext/differentiationinterface/main.jl | 43 +----- 6 files changed, 144 insertions(+), 181 deletions(-) create mode 100644 ext/AbstractPPLTestExt.jl delete mode 100644 src/test_resources.jl diff --git a/Project.toml b/Project.toml index 8b007eef..529ebd73 100644 --- a/Project.toml +++ b/Project.toml @@ -22,11 +22,13 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] +AbstractPPLTestExt = ["Test"] [compat] ADTypes = "1" @@ -43,4 +45,5 @@ MacroTools = "0.5" OrderedCollections = "1.8.1" Random = "1.6" StatsBase = "0.32, 0.33, 0.34" +Test = "1" julia = "1.10.8" diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index b45a9d63..25dc6332 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -33,8 +33,7 @@ end function AbstractPPL.prepare( adtype::AbstractADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true ) - raw = AbstractPPL.prepare(problem, x) - evaluator = VectorEvaluator{check_dims}(raw, length(x)) + evaluator = AbstractPPL.prepare(problem, x; check_dims)::VectorEvaluator # Empty inputs bypass `DI.prepare_*` (many backends fail on length-zero arrays); # `value_and_gradient!!` / `value_and_jacobian!!` short-circuit on `length(x) == 0`. length(x) == 0 && diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl new file mode 100644 index 00000000..3f6519f5 --- /dev/null +++ b/ext/AbstractPPLTestExt.jl @@ -0,0 +1,133 @@ +module AbstractPPLTestExt + +using AbstractPPL: AbstractPPL, generate_testcases, run_testcases +using Test: @test, @test_throws, @testset + +struct QuadraticProblem end +(::QuadraticProblem)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) + +struct VectorValuedProblem end +(::VectorValuedProblem)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] + +struct ValueCase + name::String + f::Any + x_proto::Any + x::Any + value::Any + gradient::Any + jacobian::Any +end + +struct ErrorCase + name::String + f::Any + x_proto::Any + x::Any + op::Any + exception::Any +end + +function AbstractPPL.generate_testcases(::Val{:vector}) + return ( + ValueCase( + "quadratic (scalar output)", + QuadraticProblem(), + zeros(3), + [3.0, 1.0, 2.0], + 14.0, + [6.0, 2.0, 4.0], + nothing, + ), + ValueCase( + "vector-valued (vector output)", + VectorValuedProblem(), + zeros(3), + [2.0, 3.0, 4.0], + [6.0, 7.0], + nothing, + [3.0 2.0 0.0; 0.0 1.0 1.0], + ), + ValueCase( + "empty input, scalar output", + x -> 7.5, + Float64[], + Float64[], + 7.5, + Float64[], + nothing, + ), + ValueCase( + "empty input, vector output", + x -> [2.0, 3.0], + Float64[], + Float64[], + [2.0, 3.0], + nothing, + zeros(2, 0), + ), + ) +end + +function AbstractPPL.generate_testcases(::Val{:edge}) + return ( + ErrorCase( + "wrong vector length", + QuadraticProblem(), + zeros(3), + [3.0, 1.0, 2.0, 99.0], + (prepared, x) -> prepared(x), + DimensionMismatch, + ), + ErrorCase( + "non-floating-point vector", + QuadraticProblem(), + zeros(3), + [3, 1, 2], + (prepared, x) -> prepared(x), + r"floating-point", + ), + ErrorCase( + "gradient of vector-valued output", + VectorValuedProblem(), + zeros(3), + [2.0, 3.0, 4.0], + (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + Exception, + ), + ) +end + +function AbstractPPL.run_testcases( + ::Val{:vector}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 +) + for case in generate_testcases(Val(:vector)) + @testset "$(case.name)" begin + prepared = prepare_fn(adtype, case.f, case.x_proto) + @test prepared(case.x) ≈ case.value atol = atol rtol = rtol + if case.gradient !== nothing + val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) + @test val ≈ case.value atol = atol rtol = rtol + @test grad ≈ case.gradient atol = atol rtol = rtol + end + if case.jacobian !== nothing + val, jac = AbstractPPL.value_and_jacobian!!(prepared, case.x) + @test val ≈ case.value atol = atol rtol = rtol + @test jac ≈ case.jacobian atol = atol rtol = rtol + end + end + end + return nothing +end + +function AbstractPPL.run_testcases(::Val{:edge}, prepare_fn=AbstractPPL.prepare; adtype) + for case in generate_testcases(Val(:edge)) + @testset "$(case.name)" begin + prepared = prepare_fn(adtype, case.f, case.x_proto) + @test_throws case.exception case.op(prepared, case.x) + end + end + return nothing +end + +end # module diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 567cedb2..dd6bb7e2 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -12,11 +12,14 @@ include("abstractprobprog.jl") include("evaluate.jl") include("evaluators/Evaluators.jl") using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!! -include("test_resources.jl") @static if VERSION >= v"1.11.0" eval(Meta.parse("public prepare, value_and_gradient!!, value_and_jacobian!!")) end +# Stubs implemented by `AbstractPPLTestExt` (loaded with `Test`). +function generate_testcases end +function run_testcases end + include("varname/optic.jl") include("varname/varname.jl") include("varname/subsumes.jl") diff --git a/src/test_resources.jl b/src/test_resources.jl deleted file mode 100644 index 143f9e25..00000000 --- a/src/test_resources.jl +++ /dev/null @@ -1,138 +0,0 @@ -module TestResources - -import ..AbstractPPL: prepare -export TestCase, generate_testcases - -struct QuadraticProblem end -struct QuadraticPrepared end -prepare(::QuadraticProblem, ::AbstractVector{<:Real}) = QuadraticPrepared() -(::QuadraticPrepared)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) - -struct VectorValuedProblem end -struct VectorValuedPrepared end -function prepare(::VectorValuedProblem, ::AbstractVector{<:Real}) - return VectorValuedPrepared() -end -(::VectorValuedPrepared)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] - -struct TestCase{F,XP,X,V,G,J,O,E} - name::String - f::F - x_proto::XP - x::X - value::V - gradient::G - jacobian::J - operation::O - exception::E -end - -function TestCase(name, f, x_proto, x, value, gradient, jacobian) - return TestCase(name, f, x_proto, x, value, gradient, jacobian, :value, nothing) -end - -function generate_testcases(::Val{:vector}) - return ( - TestCase( - "quadratic (scalar output)", - QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0], - 14.0, - [6.0, 2.0, 4.0], - nothing, - ), - TestCase( - "vector-valued (vector output)", - VectorValuedProblem(), - zeros(3), - [2.0, 3.0, 4.0], - [6.0, 7.0], - nothing, - [3.0 2.0 0.0; 0.0 1.0 1.0], - ), - TestCase( - "empty input, scalar output", - x -> 7.5, - Float64[], - Float64[], - 7.5, - Float64[], - nothing, - ), - TestCase( - "empty input, vector output", - x -> [2.0, 3.0], - Float64[], - Float64[], - [2.0, 3.0], - nothing, - zeros(2, 0), - ), - ) -end - -function generate_testcases(::Val{:namedtuple}) - return ( - TestCase( - "scalar output over (x::Real, y::Vector)", - vs -> vs.x^2 + sum(abs2, vs.y), - (x=0.0, y=zeros(2)), - (x=3.0, y=[1.0, 2.0]), - 14.0, - (x=6.0, y=[2.0, 4.0]), - nothing, - ), - TestCase( - "wrong NamedTuple structure", - vs -> vs.x^2 + sum(abs2, vs.y), - (x=0.0, y=zeros(2)), - (x=3.0, z=[1.0, 2.0]), - nothing, - nothing, - nothing, - :gradient, - r"same NamedTuple structure", - ), - ) -end - -function generate_testcases(::Val{:edge}) - return ( - TestCase( - "wrong vector length", - QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0, 99.0], - nothing, - nothing, - nothing, - :call, - DimensionMismatch, - ), - TestCase( - "non-floating-point vector", - QuadraticProblem(), - zeros(3), - [3, 1, 2], - nothing, - nothing, - nothing, - :call, - r"floating-point", - ), - TestCase( - "gradient of vector-valued output", - VectorValuedProblem(), - zeros(3), - [2.0, 3.0, 4.0], - nothing, - nothing, - nothing, - :gradient, - Exception, - ), - ) -end - -end # module diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index 5fb0d1da..0fc14c6d 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -3,50 +3,13 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL -using AbstractPPL.TestResources: generate_testcases +using AbstractPPL: AbstractPPL, run_testcases using ADTypes: AutoForwardDiff using DifferentiationInterface: DifferentiationInterface as DI using ForwardDiff using Test -const adtype = AutoForwardDiff() -const ATOL = 1e-6 -const RTOL = 1e-6 - @testset "AbstractPPLDifferentiationInterfaceExt" begin - # Use a real DI backend to exercise AbstractPPL's catch-all ADType dispatch. - @testset "ForwardDiff" begin - @testset "vector input" begin - @testset "$(case.name)" for case in generate_testcases(Val(:vector)) - prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) - @test prepared(case.x) ≈ case.value atol = ATOL rtol = RTOL - if case.gradient !== nothing - val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) - @test val ≈ case.value atol = ATOL rtol = RTOL - @test grad ≈ case.gradient atol = ATOL rtol = RTOL - end - if case.jacobian !== nothing - val, jac = AbstractPPL.value_and_jacobian!!(prepared, case.x) - @test val ≈ case.value atol = ATOL rtol = RTOL - @test jac ≈ case.jacobian atol = ATOL rtol = RTOL - end - end - - @testset "edge cases" begin - @testset "$(case.name)" for case in generate_testcases(Val(:edge)) - prepared = AbstractPPL.prepare(adtype, case.f, case.x_proto) - if case.operation === :call - @test_throws case.exception prepared(case.x) - elseif case.operation === :gradient - @test_throws case.exception AbstractPPL.value_and_gradient!!( - prepared, case.x - ) - else - error("Unknown edge-test operation: $(case.operation)") - end - end - end - end - end + run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + run_testcases(Val(:edge); adtype=AutoForwardDiff()) end From 936ddb4fd3ee8197cb6a4dc586e35bca9985718b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 5 May 2026 21:38:59 +0100 Subject: [PATCH 07/16] Fix empty-input gradient arity, advertise LDP capabilities for DI - value_and_gradient!! / value_and_jacobian!! now check arity before the length-zero short-circuit, so empty-input vector-valued / scalar-valued functions surface the correct ArgumentError instead of silently returning the wrong shape. Uses Val(0) as the cache sentinel for empty inputs. - New AbstractPPLDifferentiationInterfaceLogDensityProblemsExt advertises LogDensityOrder{1} for any DI-prepared evaluator, registering the method in __init__ to side-step extension precompile-visibility constraints. - Empty-input arity errors and the new capability advertisement are exercised in the DI test environment, which now also loads LDP. - Project.toml [compat] alphabetised. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 5 ++- ext/AbstractPPLDifferentiationInterfaceExt.jl | 41 ++++++++++--------- ...entiationInterfaceLogDensityProblemsExt.jl | 20 +++++++++ ext/AbstractPPLTestExt.jl | 16 ++++++++ .../ext/differentiationinterface/Project.toml | 2 + test/ext/differentiationinterface/main.jl | 27 +++++++++++- 6 files changed, 88 insertions(+), 23 deletions(-) create mode 100644 ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl diff --git a/Project.toml b/Project.toml index 529ebd73..16a8f758 100644 --- a/Project.toml +++ b/Project.toml @@ -26,21 +26,22 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] +AbstractPPLDifferentiationInterfaceLogDensityProblemsExt = ["DifferentiationInterface", "LogDensityProblems"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] AbstractPPLTestExt = ["Test"] [compat] ADTypes = "1" -DifferentiationInterface = "0.6, 0.7" -LogDensityProblems = "2" AbstractMCMC = "2, 3, 4, 5" Accessors = "0.1" BangBang = "0.4" DensityInterface = "0.4" +DifferentiationInterface = "0.6, 0.7" Distributions = "0.25" JSON = "0.19 - 0.21, 1" LinearAlgebra = "<0.0.1, 1" +LogDensityProblems = "2" MacroTools = "0.5" OrderedCollections = "1.8.1" Random = "1.6" diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 25dc6332..0dddd2b3 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -34,38 +34,39 @@ function AbstractPPL.prepare( adtype::AbstractADType, problem, x::AbstractVector{<:Real}; check_dims::Bool=true ) evaluator = AbstractPPL.prepare(problem, x; check_dims)::VectorEvaluator - # Empty inputs bypass `DI.prepare_*` (many backends fail on length-zero arrays); - # `value_and_gradient!!` / `value_and_jacobian!!` short-circuit on `length(x) == 0`. - length(x) == 0 && - return Prepared(adtype, evaluator, DICache(_call_evaluator, nothing, nothing, true)) y = evaluator(x) + y isa Union{Number,AbstractVector} || throw( + ArgumentError( + "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", + ), + ) + if length(x) == 0 + # `Val(0)` marks "no DI prep, but this slot's arity is supported" — + # DI's prep paths hit errors on length-0 input (e.g. ForwardDiff + # `BoundsError`), so we bypass them. The non-`nothing` marker keeps + # the scalar-vs-vector arity check in `value_and_{gradient,jacobian}!!` + # meaningful when both prep slots would otherwise be `nothing`. + gp, jp = y isa Number ? (Val(0), nothing) : (nothing, Val(0)) + return Prepared(adtype, evaluator, DICache(_call_evaluator, gp, jp, true)) + end if y isa Number target, gradient_prep, use_context = _prepare_gradient(adtype, x, evaluator) return Prepared( adtype, evaluator, DICache(target, gradient_prep, nothing, use_context) ) - elseif y isa AbstractVector - jacobian_prep = DI.prepare_jacobian( - _call_evaluator, adtype, x, DI.Constant(evaluator) - ) - return Prepared( - adtype, evaluator, DICache(_call_evaluator, nothing, jacobian_prep, true) - ) - else - throw( - ArgumentError( - "A prepared AD evaluator must return a scalar or AbstractVector; got $(typeof(y)).", - ), - ) end + jacobian_prep = DI.prepare_jacobian(_call_evaluator, adtype, x, DI.Constant(evaluator)) + return Prepared( + adtype, evaluator, DICache(_call_evaluator, nothing, jacobian_prep, true) + ) end @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T} ) where {T<:Real} - length(x) == 0 && return (p.evaluator(x), T[]) p.cache.gradient_prep === nothing && throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function.")) + length(x) == 0 && return (p.evaluator(x), T[]) val, grad = if p.cache.use_context DI.value_and_gradient( p.cache.target, p.cache.gradient_prep, p.adtype, x, DI.Constant(p.evaluator) @@ -80,12 +81,12 @@ end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{<:Real} ) + p.cache.jacobian_prep === nothing && + throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) if length(x) == 0 val = p.evaluator(x) return (val, similar(x, length(val), 0)) end - p.cache.jacobian_prep === nothing && - throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) return DI.value_and_jacobian( p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator) ) diff --git a/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl b/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl new file mode 100644 index 00000000..10fe4402 --- /dev/null +++ b/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl @@ -0,0 +1,20 @@ +module AbstractPPLDifferentiationInterfaceLogDensityProblemsExt + +using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: Prepared, VectorEvaluator +using LogDensityProblems: LogDensityProblems + +# `DICache` lives in the DI extension, which isn't a named dependency of +# this triple extension — so resolve it via `Base.get_extension` at load +# time and register the capability method then. +function __init__() + di_ext = Base.get_extension(AbstractPPL, :AbstractPPLDifferentiationInterfaceExt) + DICache = di_ext.DICache + @eval function LogDensityProblems.capabilities( + ::Type{<:Prepared{<:Any,<:VectorEvaluator,<:$DICache}} + ) + return LogDensityProblems.LogDensityOrder{1}() + end +end + +end # module diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 3f6519f5..e4d519ad 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -95,6 +95,22 @@ function AbstractPPL.generate_testcases(::Val{:edge}) (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), Exception, ), + ErrorCase( + "gradient of vector-valued output, empty input", + x -> [2.0, 3.0], + Float64[], + Float64[], + (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + r"scalar-valued", + ), + ErrorCase( + "jacobian of scalar output, empty input", + x -> 7.5, + Float64[], + Float64[], + (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + r"vector-valued", + ), ) end diff --git a/test/ext/differentiationinterface/Project.toml b/test/ext/differentiationinterface/Project.toml index 2a48e31a..4d8517ae 100644 --- a/test/ext/differentiationinterface/Project.toml +++ b/test/ext/differentiationinterface/Project.toml @@ -3,6 +3,7 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" @@ -10,4 +11,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" ADTypes = "1" DifferentiationInterface = "0.6, 0.7" ForwardDiff = "1" +LogDensityProblems = "2" julia = "1.10" diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index 0fc14c6d..f2ad9507 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -3,13 +3,38 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: AbstractPPL, run_testcases +using AbstractPPL: AbstractPPL, prepare, run_testcases using ADTypes: AutoForwardDiff using DifferentiationInterface: DifferentiationInterface as DI using ForwardDiff +using LogDensityProblems: LogDensityProblems using Test @testset "AbstractPPLDifferentiationInterfaceExt" begin run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:edge); adtype=AutoForwardDiff()) + + @testset "LogDensityProblems capabilities" begin + p_scalar = prepare(AutoForwardDiff(), x -> -0.5 * sum(abs2, x), zeros(3)) + @test LogDensityProblems.capabilities(p_scalar) == + LogDensityProblems.LogDensityOrder{1}() + x = [1.0, 2.0, 3.0] + val, grad = LogDensityProblems.logdensity_and_gradient(p_scalar, x) + @test val ≈ -0.5 * sum(abs2, x) + @test grad ≈ -x + + # All DI-prepared evaluators advertise order 1; mismatched arity + # surfaces as a runtime error from `value_and_gradient!!` rather + # than a capability downgrade. + p_vec = prepare(AutoForwardDiff(), x -> [x[1] * x[2], x[2] + x[3]], zeros(3)) + @test LogDensityProblems.capabilities(p_vec) == + LogDensityProblems.LogDensityOrder{1}() + + p_empty = prepare(AutoForwardDiff(), x -> 7.5, Float64[]) + @test LogDensityProblems.capabilities(p_empty) == + LogDensityProblems.LogDensityOrder{1}() + val, grad = LogDensityProblems.logdensity_and_gradient(p_empty, Float64[]) + @test val == 7.5 + @test grad == Float64[] + end end From c7ed718f80f10099bfc245901ad6fb2d1e436538 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 5 May 2026 21:47:26 +0100 Subject: [PATCH 08/16] Unify LDP capabilities: Prepared = order 1, bare evaluator = order 0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the triple extension and have AbstractPPLLogDensityProblemsExt advertise LogDensityOrder{1} for any Prepared{<:Any,<:VectorEvaluator}. The contract is now uniform: anything from `prepare(adtype, …)` claims gradient capability; a bare VectorEvaluator from `prepare(problem, x)` stays at order 0. Backends that don't implement value_and_gradient!! or that wrap a vector-output function surface a runtime error at call time rather than via a capability downgrade. Documents the rule in docs/src/evaluators.md. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 1 - docs/src/evaluators.md | 21 +++++++++++++++ ...entiationInterfaceLogDensityProblemsExt.jl | 20 -------------- ext/AbstractPPLLogDensityProblemsExt.jl | 26 +++++++++++-------- test/ext/logdensityproblems/main.jl | 25 +++++------------- 5 files changed, 42 insertions(+), 51 deletions(-) delete mode 100644 ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl diff --git a/Project.toml b/Project.toml index 16a8f758..00b9dddd 100644 --- a/Project.toml +++ b/Project.toml @@ -26,7 +26,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] -AbstractPPLDifferentiationInterfaceLogDensityProblemsExt = ["DifferentiationInterface", "LogDensityProblems"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] AbstractPPLTestExt = ["Test"] diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md index 4f1e1512..8240c28f 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -154,6 +154,27 @@ p = prepare(sumsimple, zeros(3)) # `VectorEvaluator{true}(sumsimple, 3)` p([1.0, 2.0, 3.0]) ``` +## LogDensityProblems integration + +Loading `LogDensityProblems` activates an extension that exposes any +vector-input evaluator as an `LogDensityProblems` problem. The capability +advertised follows a single rule: + +- **Without an AD backend** — `prepare(problem, x)` returns a bare + `VectorEvaluator`, which advertises `LogDensityOrder{0}` (primal only). +- **With an AD backend** — `prepare(adtype, problem, x)` returns a + `Prepared`, which always advertises `LogDensityOrder{1}`. + +The order-1 advertisement is unconditional on output arity. A `Prepared` +wrapping a vector-valued function still advertises `LogDensityOrder{1}`; +calling `LogDensityProblems.logdensity_and_gradient` on it raises an +`ArgumentError` from `value_and_gradient!!` (since gradient is only defined +for scalar output). Likewise, a `Prepared` whose backend doesn't implement +`value_and_gradient!!` will surface a `MethodError` at call time. The trade +is a uniform contract — "anything that came out of `prepare(adtype, …)` +claims gradient capability" — at the cost of a runtime failure for +mismatched arity or unimplemented backends. + ## API reference ```@docs diff --git a/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl b/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl deleted file mode 100644 index 10fe4402..00000000 --- a/ext/AbstractPPLDifferentiationInterfaceLogDensityProblemsExt.jl +++ /dev/null @@ -1,20 +0,0 @@ -module AbstractPPLDifferentiationInterfaceLogDensityProblemsExt - -using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator -using LogDensityProblems: LogDensityProblems - -# `DICache` lives in the DI extension, which isn't a named dependency of -# this triple extension — so resolve it via `Base.get_extension` at load -# time and register the capability method then. -function __init__() - di_ext = Base.get_extension(AbstractPPL, :AbstractPPLDifferentiationInterfaceExt) - DICache = di_ext.DICache - @eval function LogDensityProblems.capabilities( - ::Type{<:Prepared{<:Any,<:VectorEvaluator,<:$DICache}} - ) - return LogDensityProblems.LogDensityOrder{1}() - end -end - -end # module diff --git a/ext/AbstractPPLLogDensityProblemsExt.jl b/ext/AbstractPPLLogDensityProblemsExt.jl index 8f00ebb7..96542792 100644 --- a/ext/AbstractPPLLogDensityProblemsExt.jl +++ b/ext/AbstractPPLLogDensityProblemsExt.jl @@ -2,34 +2,38 @@ module AbstractPPLLogDensityProblemsExt using AbstractPPL: AbstractPPL using AbstractPPL.Evaluators: Prepared, VectorEvaluator +using ADTypes: AbstractADType using LogDensityProblems: LogDensityProblems # LDP integration is restricted to vector-input evaluators; `NamedTupleEvaluator` # does not satisfy LDP's vector-input contract. Scalar output is a runtime # contract the user must satisfy. -LogDensityProblems.logdensity(p::Prepared{<:Any,<:VectorEvaluator}, x) = p(x) +LogDensityProblems.logdensity(p::Prepared{<:AbstractADType,<:VectorEvaluator}, x) = p(x) LogDensityProblems.logdensity(e::VectorEvaluator, x) = e(x) -function LogDensityProblems.dimension(p::Prepared{<:Any,<:VectorEvaluator}) +function LogDensityProblems.dimension(p::Prepared{<:AbstractADType,<:VectorEvaluator}) return LogDensityProblems.dimension(p.evaluator) end LogDensityProblems.dimension(e::VectorEvaluator) = e.dim -# Order 0 by default. AD-backend extensions (DifferentiationInterface, -# ForwardDiff, Mooncake, etc.) must overload this for their cache type to -# advertise `LogDensityOrder{1}` — without that overload, -# `logdensity_and_gradient` will hit the `value_and_gradient!!` stub and fail. -# LDP defines `capabilities(x) = capabilities(typeof(x))`, so the type method -# alone covers both call shapes. -function LogDensityProblems.capabilities(::Type{<:Prepared{<:Any,<:VectorEvaluator}}) - return LogDensityProblems.LogDensityOrder{0}() +# A `Prepared` (i.e. `prepare(adtype, ...)`) advertises gradient capability; +# a bare evaluator (i.e. `prepare(problem, x)`, no adtype) is primal-only. +# Backends that don't implement `value_and_gradient!!` for their `Prepared` +# type will surface a `MethodError` at call time — that's the trade-off of +# this rule, in exchange for a uniform contract. +function LogDensityProblems.capabilities( + ::Type{<:Prepared{<:AbstractADType,<:VectorEvaluator}} +) + return LogDensityProblems.LogDensityOrder{1}() end function LogDensityProblems.capabilities(::Type{<:VectorEvaluator}) return LogDensityProblems.LogDensityOrder{0}() end -function LogDensityProblems.logdensity_and_gradient(p::Prepared{<:Any,<:VectorEvaluator}, x) +function LogDensityProblems.logdensity_and_gradient( + p::Prepared{<:AbstractADType,<:VectorEvaluator}, x +) val, grad = AbstractPPL.value_and_gradient!!(p, x) # `value_and_gradient!!` may alias internal storage; LDP requires a stable result. return (val, copy(grad)) diff --git a/test/ext/logdensityproblems/main.jl b/test/ext/logdensityproblems/main.jl index c0176035..80d88a2a 100644 --- a/test/ext/logdensityproblems/main.jl +++ b/test/ext/logdensityproblems/main.jl @@ -20,34 +20,21 @@ function AbstractPPL.value_and_gradient!!( return (p(x), ones(length(x))) end -# Backend extensions opt into gradient capability by overloading `capabilities` -# (typically on their cache type, e.g. `<:Prepared{<:Any,<:VectorEvaluator,<:MyCache}`). -# Here we dispatch on the AD type for simplicity. -function LogDensityProblems.capabilities(::Type{<:Prepared{TestADType,<:VectorEvaluator}}) - return LogDensityProblems.LogDensityOrder{1}() -end - @testset "AbstractPPLLogDensityProblemsExt" begin @testset "VectorEvaluator" begin ve = VectorEvaluator(sum, 3) @test LogDensityProblems.dimension(ve) == 3 @test LogDensityProblems.logdensity(ve, [1.0, 2.0, 3.0]) == 6.0 - # A bare VectorEvaluator never advertises gradient capability; - # only the wrapping `Prepared` does. + # A bare evaluator (no `Prepared` wrapper) is primal-only. @test LogDensityProblems.capabilities(ve) == LogDensityProblems.LogDensityOrder{0}() end @testset "Prepared capabilities" begin - # Without a backend overload the fallback advertises order 0 only. - p_no_overload = Prepared(AutoForwardDiff(), VectorEvaluator(sum, 3)) - @test LogDensityProblems.capabilities(p_no_overload) == - LogDensityProblems.LogDensityOrder{0}() - - # A backend that overloads capabilities advertises order 1. - p_overloaded = Prepared(TestADType(), VectorEvaluator(sum, 3)) - @test LogDensityProblems.capabilities(p_overloaded) == - LogDensityProblems.LogDensityOrder{1}() - @test LogDensityProblems.capabilities(typeof(p_overloaded)) == + # Any `Prepared` advertises order 1 — backends that don't implement + # `value_and_gradient!!` will fail at call time, not via capabilities. + p = Prepared(AutoForwardDiff(), VectorEvaluator(sum, 3)) + @test LogDensityProblems.capabilities(p) == LogDensityProblems.LogDensityOrder{1}() + @test LogDensityProblems.capabilities(typeof(p)) == LogDensityProblems.LogDensityOrder{1}() # NamedTupleEvaluator-backed Prepared has no LDP methods defined; the From 3f487636c9fcdf612cc44f2bc0bd3811b88c14ca Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 5 May 2026 23:42:28 +0100 Subject: [PATCH 09/16] format --- docs/src/evaluators.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md index 8240c28f..be7186ec 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -160,10 +160,10 @@ Loading `LogDensityProblems` activates an extension that exposes any vector-input evaluator as an `LogDensityProblems` problem. The capability advertised follows a single rule: -- **Without an AD backend** — `prepare(problem, x)` returns a bare - `VectorEvaluator`, which advertises `LogDensityOrder{0}` (primal only). -- **With an AD backend** — `prepare(adtype, problem, x)` returns a - `Prepared`, which always advertises `LogDensityOrder{1}`. + - **Without an AD backend** — `prepare(problem, x)` returns a bare + `VectorEvaluator`, which advertises `LogDensityOrder{0}` (primal only). + - **With an AD backend** — `prepare(adtype, problem, x)` returns a + `Prepared`, which always advertises `LogDensityOrder{1}`. The order-1 advertisement is unconditional on output arity. A `Prepared` wrapping a vector-valued function still advertises `LogDensityOrder{1}`; From f9ac110829c6b4e03ad10706b7eba0c2a65d652a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 6 May 2026 23:16:15 +0100 Subject: [PATCH 10/16] Skip Aqua persistent_tasks on Julia 1.10; drop dead --project from extras Aqua's persistent_tasks check spawns a wrapper subprocess that runs Pkg.precompile() on a package depending on AbstractPPL. On Julia 1.10 this hits "Declaring __precompile__(false) is not allowed in files that are being precompiled" inside the wrapper's extension precompile path. The dedicated Ext CI jobs load and exercise every extension on min Julia and pass, so this is a Julia 1.10 / Aqua interaction, not a defect in our extensions. Re-enable when min is bumped past 1.10. The CI extras step also had a redundant --project=. that was overridden by run_extras.jl's own Pkg.activate; drop it. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/CI.yml | 2 +- test/Aqua.jl | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 17ae718f..88fce19f 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -72,6 +72,6 @@ jobs: version: ${{ matrix.version }} - uses: julia-actions/cache@v3 - uses: julia-actions/julia-buildpkg@v1 - - run: julia --project=. test/run_extras.jl + - run: julia test/run_extras.jl env: LABEL: ${{ matrix.label }} diff --git a/test/Aqua.jl b/test/Aqua.jl index 27d4103f..9fd508a4 100644 --- a/test/Aqua.jl +++ b/test/Aqua.jl @@ -3,6 +3,13 @@ module AquaTests using Aqua: Aqua using AbstractPPL -Aqua.test_all(AbstractPPL) +# `persistent_tasks` spawns a subprocess that runs `Pkg.precompile()` on a +# wrapper package depending on AbstractPPL. On Julia 1.10, this hits +# "Declaring __precompile__(false) is not allowed in files that are being +# precompiled" inside the wrapper's extension precompile path — a known +# brittleness in the Aqua/Julia 1.10 interaction, not in our extensions +# (the dedicated `Ext` CI jobs load and exercise every extension on `min` +# Julia and pass). Re-enable once `min` is bumped past 1.10. +Aqua.test_all(AbstractPPL; persistent_tasks=VERSION >= v"1.11") end From 4861b66461caa98a78a74cb752fd6e545d6dd525 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 6 May 2026 23:16:23 +0100 Subject: [PATCH 11/16] Polish evaluator API: public stubs, docstrings, AD-missing hints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Mark generate_testcases / run_testcases as public (1.11+) so downstream AD-backend packages can reuse the conformance suite without reaching into private API. - Note in the Prepared docstring that the two-arg constructor is for backends that allocate fresh storage per call. - Expand the _ad_output_arity error message to call out matrix/tuple outputs explicitly and recommend flattening. - Add a MethodError hint on value_and_gradient!! / value_and_jacobian!! that fires when no AD backend extension is loaded — also surfaces through LogDensityProblems.logdensity_and_gradient, which delegates to value_and_gradient!!. Suppressed once any backend registers a method, where the standard candidate list is more informative. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/AbstractPPL.jl | 15 +++++++++++---- src/evaluators/Evaluators.jl | 16 +++++++++++++++- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index dd6bb7e2..cf644ce2 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -12,14 +12,21 @@ include("abstractprobprog.jl") include("evaluate.jl") include("evaluators/Evaluators.jl") using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!! -@static if VERSION >= v"1.11.0" - eval(Meta.parse("public prepare, value_and_gradient!!, value_and_jacobian!!")) -end -# Stubs implemented by `AbstractPPLTestExt` (loaded with `Test`). +# Stubs implemented by `AbstractPPLTestExt` (loaded with `Test`). Part of the +# public API for downstream AD-backend packages that want to reuse the shared +# conformance suite in their own test setups. function generate_testcases end function run_testcases end +@static if VERSION >= v"1.11.0" + eval( + Meta.parse( + "public prepare, value_and_gradient!!, value_and_jacobian!!, generate_testcases, run_testcases", + ), + ) +end + include("varname/optic.jl") include("varname/varname.jl") include("varname/subsumes.jl") diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index ff2d5407..3e6a63e0 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -16,7 +16,8 @@ AD-prepared evaluator parameterised by backend type `AD`. `NamedTupleEvaluator`); forwarded on `p(x)`. - `cache` — backend-specific pre-allocated state (ForwardDiff configs, Mooncake caches, DifferentiationInterface preps, etc.). `Nothing` when the backend requires - no cached state. + no cached state. The two-argument constructor is for backends that allocate + fresh storage on every `value_and_gradient!!` call. Extension packages implement `value_and_gradient!!` (and optionally `value_and_jacobian!!`) by specialising on the `cache` type: @@ -258,6 +259,19 @@ function __init__() "\nCalling `prepare` with an AD backend requires loading the corresponding extension (e.g., `using DifferentiationInterface`).", ) end + # `value_and_gradient!!` / `value_and_jacobian!!` are stubs until an AD + # backend extension adds methods. Suppress the hint once any backend is + # loaded — the standard `MethodError` candidate list is then more useful + # than a generic "load an extension" message. Also fires when reached via + # `LogDensityProblems.logdensity_and_gradient`, which delegates here. + Base.Experimental.register_error_hint(MethodError) do io, exc, args, kwargs + exc.f === value_and_gradient!! || exc.f === value_and_jacobian!! || return nothing + isempty(methods(exc.f)) || return nothing + print( + io, + "\nNo AD backend extension is loaded. Load `DifferentiationInterface` (with a backend like `ForwardDiff`) or `Mooncake` to enable gradient/jacobian computation.", + ) + end end end # module From 78489478b65debd8f92c11a02b02e25fb08ebf56 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Thu, 7 May 2026 18:00:12 +0100 Subject: [PATCH 12/16] Update JULIA.md --- JULIA.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/JULIA.md b/JULIA.md index 4067ee32..15e8b55d 100644 --- a/JULIA.md +++ b/JULIA.md @@ -1,6 +1,6 @@ # JULIA.md -Shared day-to-day Julia practices. DynamicPPL-specific review notes live in `AGENTS.md`; newcomer context lives in `docs/src/onboarding.md`. +Shared day-to-day Julia practices. AbstractPPL-specific review notes live in `AGENTS.md`. ## Engineering From 5797c8339cd54e4d79417f9955fd3f972bb6148d Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 7 May 2026 18:21:40 +0100 Subject: [PATCH 13/16] Address PR review: drop LDP ext, fix DI gradient/jacobian edges - Drop AbstractPPLLogDensityProblemsExt and its test launcher; downstream packages can define LDP methods on `Prepared` directly when needed. - Add `_check_vector_length` to enforce dim mismatch on `value_and_gradient!!`/`value_and_jacobian!!` before the `Val(0)` short-circuit and DI call (matches `VectorEvaluator{true}` behaviour). - Unify gradient/jacobian DI prep into `_prepare_di(prep, ...)`; the `AutoReverseDiff{true}` carve-out now also fixes Jacobian tape recording. - Tidy DI tests: remove LDP block, drop redundant constants and edge cases, tighten gradient-of-vector-valued case to `ArgumentError`. - Document prep-time `evaluator(x)` call and reserved `Val{:vector}` / `Val{:edge}` testcase keys. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/CI.yml | 1 - AGENTS.md | 2 - Project.toml | 3 - docs/src/evaluators.md | 21 ------- ext/AbstractPPLDifferentiationInterfaceExt.jl | 46 ++++++++------- ext/AbstractPPLLogDensityProblemsExt.jl | 42 -------------- ext/AbstractPPLTestExt.jl | 18 +++++- src/AbstractPPL.jl | 22 ++++++- src/evaluators/Evaluators.jl | 21 ++++--- .../ext/differentiationinterface/Project.toml | 2 - test/ext/differentiationinterface/main.jl | 27 +-------- test/ext/logdensityproblems/Project.toml | 11 ---- test/ext/logdensityproblems/main.jl | 57 ------------------- test/run_extras.jl | 3 +- 14 files changed, 78 insertions(+), 198 deletions(-) delete mode 100644 ext/AbstractPPLLogDensityProblemsExt.jl delete mode 100644 test/ext/logdensityproblems/Project.toml delete mode 100644 test/ext/logdensityproblems/main.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 88fce19f..1e8f957e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -61,7 +61,6 @@ jobs: matrix: label: - ext/differentiationinterface - - ext/logdensityproblems version: - '1' - 'min' diff --git a/AGENTS.md b/AGENTS.md index 3f69c1c2..139496fe 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -24,14 +24,12 @@ AbstractPPL.jl provides shared interfaces and utilities for probabilistic progra - When touching `VarName` or optics code, test symbolic, indexed, nested, and serialization round-trip cases. - Respect evaluator contracts: `VectorEvaluator` is for flat vectors; `NamedTupleEvaluator` is for stable named structures; `!!` derivative APIs may return cache-aliased arrays. - Use `check_dims=false` only for trusted AD hot paths. Public evaluator calls should validate user input. - - Keep `LogDensityProblems` integration vector-only unless its contract changes. ## Tests - Core tests: `GROUP=Tests julia --project=test test/runtests.jl` - Doctests: `GROUP=Doctests julia --project=test test/runtests.jl` - Full package tests: `julia --project=. -e 'using Pkg; Pkg.test()'` - - LogDensityProblems extension: `LABEL=ext/logdensityproblems julia --project=. test/run_extras.jl` - Docs: `julia --project=docs docs/make.jl` Run the smallest relevant test first, then broaden when changing public interfaces, extensions, or downstream-facing behaviour. Do not weaken tests just to make CI pass. diff --git a/Project.toml b/Project.toml index 00b9dddd..a07f4a20 100644 --- a/Project.toml +++ b/Project.toml @@ -21,13 +21,11 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] -AbstractPPLLogDensityProblemsExt = ["LogDensityProblems"] AbstractPPLTestExt = ["Test"] [compat] @@ -40,7 +38,6 @@ DifferentiationInterface = "0.6, 0.7" Distributions = "0.25" JSON = "0.19 - 0.21, 1" LinearAlgebra = "<0.0.1, 1" -LogDensityProblems = "2" MacroTools = "0.5" OrderedCollections = "1.8.1" Random = "1.6" diff --git a/docs/src/evaluators.md b/docs/src/evaluators.md index be7186ec..4f1e1512 100644 --- a/docs/src/evaluators.md +++ b/docs/src/evaluators.md @@ -154,27 +154,6 @@ p = prepare(sumsimple, zeros(3)) # `VectorEvaluator{true}(sumsimple, 3)` p([1.0, 2.0, 3.0]) ``` -## LogDensityProblems integration - -Loading `LogDensityProblems` activates an extension that exposes any -vector-input evaluator as an `LogDensityProblems` problem. The capability -advertised follows a single rule: - - - **Without an AD backend** — `prepare(problem, x)` returns a bare - `VectorEvaluator`, which advertises `LogDensityOrder{0}` (primal only). - - **With an AD backend** — `prepare(adtype, problem, x)` returns a - `Prepared`, which always advertises `LogDensityOrder{1}`. - -The order-1 advertisement is unconditional on output arity. A `Prepared` -wrapping a vector-valued function still advertises `LogDensityOrder{1}`; -calling `LogDensityProblems.logdensity_and_gradient` on it raises an -`ArgumentError` from `value_and_gradient!!` (since gradient is only defined -for scalar output). Likewise, a `Prepared` whose backend doesn't implement -`value_and_gradient!!` will surface a `MethodError` at call time. The trade -is a uniform contract — "anything that came out of `prepare(adtype, …)` -claims gradient capability" — at the cost of a runtime failure for -mismatched arity or unimplemented backends. - ## API reference ```@docs diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index 0dddd2b3..edcb8a42 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -1,8 +1,8 @@ module AbstractPPLDifferentiationInterfaceExt using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator -using ADTypes: ADTypes, AbstractADType, AutoReverseDiff +using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator +using ADTypes: AbstractADType, AutoReverseDiff using DifferentiationInterface: DifferentiationInterface as DI # Differentiate only `x`; the evaluator is passed as a `DI.Constant` context so @@ -17,17 +17,15 @@ struct DICache{F,GP,JP} end # Compiled ReverseDiff only reuses a compiled tape on the one-argument path; -# the DI.Constant route would invalidate the tape across calls. -function _prepare_gradient(adtype::AutoReverseDiff{true}, x, evaluator) +# `DI.Constant` deactivates tape recording, so close the evaluator into the +# target and call DI without contexts. +function _prepare_di(prep::F, adtype::AutoReverseDiff{true}, x, evaluator) where {F} target = Base.Fix2(_call_evaluator, evaluator) - gradient_prep = DI.prepare_gradient(target, adtype, x) - return target, gradient_prep, false + return target, prep(target, adtype, x), false end -function _prepare_gradient(adtype::AbstractADType, x, evaluator) - target = _call_evaluator - gradient_prep = DI.prepare_gradient(target, adtype, x, DI.Constant(evaluator)) - return target, gradient_prep, true +function _prepare_di(prep::F, adtype::AbstractADType, x, evaluator) where {F} + return _call_evaluator, prep(_call_evaluator, adtype, x, DI.Constant(evaluator)), true end function AbstractPPL.prepare( @@ -50,15 +48,17 @@ function AbstractPPL.prepare( return Prepared(adtype, evaluator, DICache(_call_evaluator, gp, jp, true)) end if y isa Number - target, gradient_prep, use_context = _prepare_gradient(adtype, x, evaluator) + target, gradient_prep, use_context = _prepare_di( + DI.prepare_gradient, adtype, x, evaluator + ) return Prepared( adtype, evaluator, DICache(target, gradient_prep, nothing, use_context) ) end - jacobian_prep = DI.prepare_jacobian(_call_evaluator, adtype, x, DI.Constant(evaluator)) - return Prepared( - adtype, evaluator, DICache(_call_evaluator, nothing, jacobian_prep, true) + target, jacobian_prep, use_context = _prepare_di( + DI.prepare_jacobian, adtype, x, evaluator ) + return Prepared(adtype, evaluator, DICache(target, nothing, jacobian_prep, use_context)) end @inline function AbstractPPL.value_and_gradient!!( @@ -66,16 +66,17 @@ end ) where {T<:Real} p.cache.gradient_prep === nothing && throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function.")) + Evaluators._check_vector_length(p.evaluator.dim, x) + # Bypass DI on length-0 input — DI prep paths fail (e.g. ForwardDiff + # `BoundsError`); typed `T[]` matches the caller's element type. length(x) == 0 && return (p.evaluator(x), T[]) - val, grad = if p.cache.use_context + return if p.cache.use_context DI.value_and_gradient( p.cache.target, p.cache.gradient_prep, p.adtype, x, DI.Constant(p.evaluator) ) else DI.value_and_gradient(p.cache.target, p.cache.gradient_prep, p.adtype, x) end - # Some DI backends may return a non-`Vector` gradient; normalise. - return (val, grad isa Vector ? grad : collect(grad)) end @inline function AbstractPPL.value_and_jacobian!!( @@ -83,13 +84,18 @@ end ) p.cache.jacobian_prep === nothing && throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) + Evaluators._check_vector_length(p.evaluator.dim, x) if length(x) == 0 val = p.evaluator(x) return (val, similar(x, length(val), 0)) end - return DI.value_and_jacobian( - p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator) - ) + return if p.cache.use_context + DI.value_and_jacobian( + p.cache.target, p.cache.jacobian_prep, p.adtype, x, DI.Constant(p.evaluator) + ) + else + DI.value_and_jacobian(p.cache.target, p.cache.jacobian_prep, p.adtype, x) + end end end # module diff --git a/ext/AbstractPPLLogDensityProblemsExt.jl b/ext/AbstractPPLLogDensityProblemsExt.jl deleted file mode 100644 index 96542792..00000000 --- a/ext/AbstractPPLLogDensityProblemsExt.jl +++ /dev/null @@ -1,42 +0,0 @@ -module AbstractPPLLogDensityProblemsExt - -using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator -using ADTypes: AbstractADType -using LogDensityProblems: LogDensityProblems - -# LDP integration is restricted to vector-input evaluators; `NamedTupleEvaluator` -# does not satisfy LDP's vector-input contract. Scalar output is a runtime -# contract the user must satisfy. - -LogDensityProblems.logdensity(p::Prepared{<:AbstractADType,<:VectorEvaluator}, x) = p(x) -LogDensityProblems.logdensity(e::VectorEvaluator, x) = e(x) - -function LogDensityProblems.dimension(p::Prepared{<:AbstractADType,<:VectorEvaluator}) - return LogDensityProblems.dimension(p.evaluator) -end -LogDensityProblems.dimension(e::VectorEvaluator) = e.dim - -# A `Prepared` (i.e. `prepare(adtype, ...)`) advertises gradient capability; -# a bare evaluator (i.e. `prepare(problem, x)`, no adtype) is primal-only. -# Backends that don't implement `value_and_gradient!!` for their `Prepared` -# type will surface a `MethodError` at call time — that's the trade-off of -# this rule, in exchange for a uniform contract. -function LogDensityProblems.capabilities( - ::Type{<:Prepared{<:AbstractADType,<:VectorEvaluator}} -) - return LogDensityProblems.LogDensityOrder{1}() -end -function LogDensityProblems.capabilities(::Type{<:VectorEvaluator}) - return LogDensityProblems.LogDensityOrder{0}() -end - -function LogDensityProblems.logdensity_and_gradient( - p::Prepared{<:AbstractADType,<:VectorEvaluator}, x -) - val, grad = AbstractPPL.value_and_gradient!!(p, x) - # `value_and_gradient!!` may alias internal storage; LDP requires a stable result. - return (val, copy(grad)) -end - -end # module diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index e4d519ad..4f410f2b 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -93,7 +93,7 @@ function AbstractPPL.generate_testcases(::Val{:edge}) zeros(3), [2.0, 3.0, 4.0], (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), - Exception, + ArgumentError, ), ErrorCase( "gradient of vector-valued output, empty input", @@ -111,6 +111,22 @@ function AbstractPPL.generate_testcases(::Val{:edge}) (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), r"vector-valued", ), + ErrorCase( + "value_and_gradient!! wrong vector length", + QuadraticProblem(), + zeros(3), + [3.0, 1.0, 2.0, 99.0], + (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + DimensionMismatch, + ), + ErrorCase( + "value_and_jacobian!! wrong vector length", + VectorValuedProblem(), + zeros(3), + [2.0, 3.0, 4.0, 5.0], + (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + DimensionMismatch, + ), ) end diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index cf644ce2..bea9c16a 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -13,10 +13,26 @@ include("evaluate.jl") include("evaluators/Evaluators.jl") using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!! -# Stubs implemented by `AbstractPPLTestExt` (loaded with `Test`). Part of the -# public API for downstream AD-backend packages that want to reuse the shared -# conformance suite in their own test setups. +""" + generate_testcases(::Val{group}) + +Return a tuple of test cases for the conformance `group`. Implemented by the +`Test` extension (`AbstractPPLTestExt`). Reserved group keys (extensions must +not redefine these): `:vector` for value/gradient/jacobian round-trips on +vector-input evaluators; `:edge` for error-path cases. Downstream packages may +add their own group keys (e.g. `:my_backend_group`) by adding methods to this +function. +""" function generate_testcases end + +""" + run_testcases(::Val{group}, prepare_fn=AbstractPPL.prepare; adtype, kwargs...) + +Run the test cases produced by [`generate_testcases`](@ref) against an AD +backend, using `prepare_fn` (default `AbstractPPL.prepare`) to construct each +prepared evaluator. Implemented by the `Test` extension. See +[`generate_testcases`](@ref) for reserved group keys. +""" function run_testcases end @static if VERSION >= v"1.11.0" diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index 3e6a63e0..b33cc8c3 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -56,6 +56,11 @@ prepares gradient or jacobian machinery for vector inputs. the input shape on each call. Pass `check_dims=false` to skip the per-call check, e.g. inside an AD backend's hot path where the input shape is already guaranteed. + +The three-argument AD-aware form may invoke `problem` once during preparation +to detect output arity (scalar vs vector) and select gradient or jacobian +machinery accordingly. Avoid `prepare` calls when `problem` has side effects +that should fire only on user-driven evaluations. """ function prepare end @@ -166,13 +171,16 @@ function _reject_integer_input(x) ) end +function _check_vector_length(dim::Int, x) + length(x) == dim || throw( + DimensionMismatch("Expected a vector of length $dim, but got length $(length(x))."), + ) + return nothing +end + function (e::VectorEvaluator{true})(x::AbstractVector{T}) where {T} T <: Integer && _reject_integer_input(x) - length(x) == e.dim || throw( - DimensionMismatch( - "Expected a vector of length $(e.dim), but got length $(length(x))." - ), - ) + _check_vector_length(e.dim, x) return e.f(x) end @@ -262,8 +270,7 @@ function __init__() # `value_and_gradient!!` / `value_and_jacobian!!` are stubs until an AD # backend extension adds methods. Suppress the hint once any backend is # loaded — the standard `MethodError` candidate list is then more useful - # than a generic "load an extension" message. Also fires when reached via - # `LogDensityProblems.logdensity_and_gradient`, which delegates here. + # than a generic "load an extension" message. Base.Experimental.register_error_hint(MethodError) do io, exc, args, kwargs exc.f === value_and_gradient!! || exc.f === value_and_jacobian!! || return nothing isempty(methods(exc.f)) || return nothing diff --git a/test/ext/differentiationinterface/Project.toml b/test/ext/differentiationinterface/Project.toml index 4d8517ae..2a48e31a 100644 --- a/test/ext/differentiationinterface/Project.toml +++ b/test/ext/differentiationinterface/Project.toml @@ -3,7 +3,6 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" @@ -11,5 +10,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" ADTypes = "1" DifferentiationInterface = "0.6, 0.7" ForwardDiff = "1" -LogDensityProblems = "2" julia = "1.10" diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index f2ad9507..0637b6f6 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -3,38 +3,13 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: AbstractPPL, prepare, run_testcases +using AbstractPPL: run_testcases using ADTypes: AutoForwardDiff using DifferentiationInterface: DifferentiationInterface as DI using ForwardDiff -using LogDensityProblems: LogDensityProblems using Test @testset "AbstractPPLDifferentiationInterfaceExt" begin run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:edge); adtype=AutoForwardDiff()) - - @testset "LogDensityProblems capabilities" begin - p_scalar = prepare(AutoForwardDiff(), x -> -0.5 * sum(abs2, x), zeros(3)) - @test LogDensityProblems.capabilities(p_scalar) == - LogDensityProblems.LogDensityOrder{1}() - x = [1.0, 2.0, 3.0] - val, grad = LogDensityProblems.logdensity_and_gradient(p_scalar, x) - @test val ≈ -0.5 * sum(abs2, x) - @test grad ≈ -x - - # All DI-prepared evaluators advertise order 1; mismatched arity - # surfaces as a runtime error from `value_and_gradient!!` rather - # than a capability downgrade. - p_vec = prepare(AutoForwardDiff(), x -> [x[1] * x[2], x[2] + x[3]], zeros(3)) - @test LogDensityProblems.capabilities(p_vec) == - LogDensityProblems.LogDensityOrder{1}() - - p_empty = prepare(AutoForwardDiff(), x -> 7.5, Float64[]) - @test LogDensityProblems.capabilities(p_empty) == - LogDensityProblems.LogDensityOrder{1}() - val, grad = LogDensityProblems.logdensity_and_gradient(p_empty, Float64[]) - @test val == 7.5 - @test grad == Float64[] - end end diff --git a/test/ext/logdensityproblems/Project.toml b/test/ext/logdensityproblems/Project.toml deleted file mode 100644 index a6d089ab..00000000 --- a/test/ext/logdensityproblems/Project.toml +++ /dev/null @@ -1,11 +0,0 @@ -[deps] -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" -ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" -Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[compat] -ADTypes = "1" -LogDensityProblems = "2" -julia = "1.10" diff --git a/test/ext/logdensityproblems/main.jl b/test/ext/logdensityproblems/main.jl deleted file mode 100644 index 80d88a2a..00000000 --- a/test/ext/logdensityproblems/main.jl +++ /dev/null @@ -1,57 +0,0 @@ -using Pkg -Pkg.activate(@__DIR__) -Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) -Pkg.instantiate() - -using AbstractPPL -using AbstractPPL.Evaluators: Prepared, VectorEvaluator, NamedTupleEvaluator -using ADTypes: AbstractADType, AutoForwardDiff -using LogDensityProblems: LogDensityProblems -using Test - -# A NamedTupleEvaluator does not satisfy LDP's vector-input contract, so the -# extension does not define LDP methods for it. - -struct TestADType <: AbstractADType end - -function AbstractPPL.value_and_gradient!!( - p::Prepared{TestADType}, x::AbstractVector{<:Real} -) - return (p(x), ones(length(x))) -end - -@testset "AbstractPPLLogDensityProblemsExt" begin - @testset "VectorEvaluator" begin - ve = VectorEvaluator(sum, 3) - @test LogDensityProblems.dimension(ve) == 3 - @test LogDensityProblems.logdensity(ve, [1.0, 2.0, 3.0]) == 6.0 - # A bare evaluator (no `Prepared` wrapper) is primal-only. - @test LogDensityProblems.capabilities(ve) == LogDensityProblems.LogDensityOrder{0}() - end - - @testset "Prepared capabilities" begin - # Any `Prepared` advertises order 1 — backends that don't implement - # `value_and_gradient!!` will fail at call time, not via capabilities. - p = Prepared(AutoForwardDiff(), VectorEvaluator(sum, 3)) - @test LogDensityProblems.capabilities(p) == LogDensityProblems.LogDensityOrder{1}() - @test LogDensityProblems.capabilities(typeof(p)) == - LogDensityProblems.LogDensityOrder{1}() - - # NamedTupleEvaluator-backed Prepared has no LDP methods defined; the - # extension only integrates vector-input evaluators. - p_nt = Prepared( - AutoForwardDiff(), NamedTupleEvaluator(x -> x.a + sum(x.b), (a=0.0, b=zeros(2))) - ) - @test_throws MethodError LogDensityProblems.dimension(p_nt) - @test LogDensityProblems.capabilities(p_nt) === nothing - end - - @testset "logdensity_and_gradient" begin - f = x -> -0.5 * sum(abs2, x) - p = Prepared(TestADType(), VectorEvaluator(f, 3)) - x = [1.0, 2.0, 3.0] - val, grad = LogDensityProblems.logdensity_and_gradient(p, x) - @test val ≈ f(x) - @test grad ≈ ones(3) - end -end diff --git a/test/run_extras.jl b/test/run_extras.jl index 6b43493d..e557f363 100644 --- a/test/run_extras.jl +++ b/test/run_extras.jl @@ -2,9 +2,8 @@ # # Usage (from the repo root): # LABEL=ext/differentiationinterface julia test/run_extras.jl -# LABEL=ext/logdensityproblems julia test/run_extras.jl -const VALID_LABELS = ("ext/differentiationinterface", "ext/logdensityproblems") +const VALID_LABELS = ("ext/differentiationinterface",) label = get(ENV, "LABEL", nothing) label in VALID_LABELS || From 17ebe0e08877ce42eae2e4e6358f6c379ed2f023 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sat, 9 May 2026 00:09:05 +0100 Subject: [PATCH 14/16] Reject integer inputs in DI value_and_{gradient,jacobian}!! The bare `VectorEvaluator` rejected integer vectors with a clear `ArgumentError`, but the DI derivative entry points went straight to DI prep and surfaced an opaque `PreparationMismatchError`. Apply the same compile-time `T <: Integer` check before dispatching to DI, and add matching `:edge` testcases. Co-Authored-By: Claude Opus 4.7 (1M context) --- Project.toml | 2 +- ext/AbstractPPLDifferentiationInterfaceExt.jl | 13 ++++++------- ext/AbstractPPLTestExt.jl | 16 ++++++++++++++++ src/AbstractPPL.jl | 3 +-- src/evaluators/Evaluators.jl | 8 ++------ 5 files changed, 26 insertions(+), 16 deletions(-) diff --git a/Project.toml b/Project.toml index a07f4a20..e767cf77 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AbstractPPL" uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" -keywords = ["probablistic programming"] +keywords = ["probabilistic programming"] license = "MIT" desc = "Common interfaces for probabilistic programming" version = "0.14.3" diff --git a/ext/AbstractPPLDifferentiationInterfaceExt.jl b/ext/AbstractPPLDifferentiationInterfaceExt.jl index edcb8a42..80f95ee1 100644 --- a/ext/AbstractPPLDifferentiationInterfaceExt.jl +++ b/ext/AbstractPPLDifferentiationInterfaceExt.jl @@ -39,11 +39,8 @@ function AbstractPPL.prepare( ), ) if length(x) == 0 - # `Val(0)` marks "no DI prep, but this slot's arity is supported" — - # DI's prep paths hit errors on length-0 input (e.g. ForwardDiff - # `BoundsError`), so we bypass them. The non-`nothing` marker keeps - # the scalar-vs-vector arity check in `value_and_{gradient,jacobian}!!` - # meaningful when both prep slots would otherwise be `nothing`. + # DI prep crashes on length-0 input (e.g. ForwardDiff `BoundsError`); the + # `Val(0)` sentinel keeps the `gradient_prep === nothing` arity check meaningful. gp, jp = y isa Number ? (Val(0), nothing) : (nothing, Val(0)) return Prepared(adtype, evaluator, DICache(_call_evaluator, gp, jp, true)) end @@ -66,6 +63,7 @@ end ) where {T<:Real} p.cache.gradient_prep === nothing && throw(ArgumentError("`value_and_gradient!!` requires a scalar-valued function.")) + T <: Integer && Evaluators._reject_integer_input(x) Evaluators._check_vector_length(p.evaluator.dim, x) # Bypass DI on length-0 input — DI prep paths fail (e.g. ForwardDiff # `BoundsError`); typed `T[]` matches the caller's element type. @@ -80,10 +78,11 @@ end end @inline function AbstractPPL.value_and_jacobian!!( - p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{<:Real} -) + p::Prepared{<:AbstractADType,<:VectorEvaluator,<:DICache}, x::AbstractVector{T} +) where {T<:Real} p.cache.jacobian_prep === nothing && throw(ArgumentError("`value_and_jacobian!!` requires a vector-valued function.")) + T <: Integer && Evaluators._reject_integer_input(x) Evaluators._check_vector_length(p.evaluator.dim, x) if length(x) == 0 val = p.evaluator(x) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 4f410f2b..6f7463da 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -127,6 +127,22 @@ function AbstractPPL.generate_testcases(::Val{:edge}) (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), DimensionMismatch, ), + ErrorCase( + "value_and_gradient!! non-floating-point vector", + QuadraticProblem(), + zeros(3), + [3, 1, 2], + (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + r"floating-point", + ), + ErrorCase( + "value_and_jacobian!! non-floating-point vector", + VectorValuedProblem(), + zeros(3), + [2, 3, 4], + (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + r"floating-point", + ), ) end diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index bea9c16a..0b50ca9a 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -20,8 +20,7 @@ Return a tuple of test cases for the conformance `group`. Implemented by the `Test` extension (`AbstractPPLTestExt`). Reserved group keys (extensions must not redefine these): `:vector` for value/gradient/jacobian round-trips on vector-input evaluators; `:edge` for error-path cases. Downstream packages may -add their own group keys (e.g. `:my_backend_group`) by adding methods to this -function. +add other keys. """ function generate_testcases end diff --git a/src/evaluators/Evaluators.jl b/src/evaluators/Evaluators.jl index b33cc8c3..a88f2910 100644 --- a/src/evaluators/Evaluators.jl +++ b/src/evaluators/Evaluators.jl @@ -16,8 +16,7 @@ AD-prepared evaluator parameterised by backend type `AD`. `NamedTupleEvaluator`); forwarded on `p(x)`. - `cache` — backend-specific pre-allocated state (ForwardDiff configs, Mooncake caches, DifferentiationInterface preps, etc.). `Nothing` when the backend requires - no cached state. The two-argument constructor is for backends that allocate - fresh storage on every `value_and_gradient!!` call. + no cached state. Extension packages implement `value_and_gradient!!` (and optionally `value_and_jacobian!!`) by specialising on the `cache` type: @@ -267,10 +266,7 @@ function __init__() "\nCalling `prepare` with an AD backend requires loading the corresponding extension (e.g., `using DifferentiationInterface`).", ) end - # `value_and_gradient!!` / `value_and_jacobian!!` are stubs until an AD - # backend extension adds methods. Suppress the hint once any backend is - # loaded — the standard `MethodError` candidate list is then more useful - # than a generic "load an extension" message. + # Same fire-only-when-no-backend-loaded logic as the `prepare` hint above. Base.Experimental.register_error_hint(MethodError) do io, exc, args, kwargs exc.f === value_and_gradient!! || exc.f === value_and_jacobian!! || return nothing isempty(methods(exc.f)) || return nothing From 4fe4a9348fd4bdd7baa7ac2db35ab01acd1dc912 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Wed, 13 May 2026 13:10:06 +0530 Subject: [PATCH 15/16] Add coverage upload to ext job, bump setup-julia and codecov-action versions --- .github/workflows/CI.yml | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 1e8f957e..6a4375b1 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -33,7 +33,7 @@ jobs: steps: - uses: actions/checkout@v6 - - uses: julia-actions/setup-julia@v2 + - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.runner.version }} @@ -47,7 +47,7 @@ jobs: - uses: julia-actions/julia-processcoverage@v1 - - uses: codecov/codecov-action@v5 + - uses: codecov/codecov-action@v6 with: files: lcov.info token: ${{ secrets.CODECOV_TOKEN }} @@ -66,11 +66,17 @@ jobs: - 'min' steps: - uses: actions/checkout@v6 - - uses: julia-actions/setup-julia@v2 + - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.version }} - uses: julia-actions/cache@v3 - uses: julia-actions/julia-buildpkg@v1 - - run: julia test/run_extras.jl + - run: julia --code-coverage=user test/run_extras.jl env: LABEL: ${{ matrix.label }} + - uses: julia-actions/julia-processcoverage@v1 + - uses: codecov/codecov-action@v6 + with: + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: true From e0adffe5939660e6e1a9496d291112abb5152b21 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Wed, 13 May 2026 13:15:29 +0530 Subject: [PATCH 16/16] Add regression test for AutoReverseDiff compiled tape no-context path --- .../ext/differentiationinterface/Project.toml | 4 +++- test/ext/differentiationinterface/main.jl | 20 ++++++++++++++++--- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/test/ext/differentiationinterface/Project.toml b/test/ext/differentiationinterface/Project.toml index 2a48e31a..af28cf08 100644 --- a/test/ext/differentiationinterface/Project.toml +++ b/test/ext/differentiationinterface/Project.toml @@ -1,13 +1,15 @@ [deps] -AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] ADTypes = "1" DifferentiationInterface = "0.6, 0.7" ForwardDiff = "1" +ReverseDiff = "1" julia = "1.10" diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index 0637b6f6..4e19a8ce 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -3,13 +3,27 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: run_testcases -using ADTypes: AutoForwardDiff -using DifferentiationInterface: DifferentiationInterface as DI +using AbstractPPL: AbstractPPL, run_testcases +using ADTypes: AutoForwardDiff, AutoReverseDiff +using DifferentiationInterface using ForwardDiff +using ReverseDiff using Test @testset "AbstractPPLDifferentiationInterfaceExt" begin run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:edge); adtype=AutoForwardDiff()) + + @testset "AutoReverseDiff compiled tape (no-context path)" begin + ad = AutoReverseDiff(; compile=true) + p_scalar = AbstractPPL.prepare(ad, x -> sum(abs2, x), zeros(3)) + p_vector = AbstractPPL.prepare(ad, x -> [x[1] * x[2], x[2] + x[3]], zeros(3)) + + @test !p_scalar.cache.use_context + @test !isnothing(p_scalar.cache.gradient_prep.tape) + @test !p_vector.cache.use_context + @test !isnothing(p_vector.cache.jacobian_prep.tape) + + run_testcases(Val(:vector); adtype=ad, atol=1e-6, rtol=1e-6) + end end