From f86642edfe078663f54f18830eebb8d32aa4dd4d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 25 May 2026 13:17:52 +0000 Subject: [PATCH 01/21] Initial plan From a20e171a5e61b7388311323a8aa3837d2ae93bb3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 25 May 2026 13:28:37 +0000 Subject: [PATCH 02/21] Add ForwardDiff extension using ForwardDiff's public API Adds `AbstractPPLForwardDiffExt` that directly uses ForwardDiff's public API (gradient!, jacobian!, hessian! with DiffResults and pre-allocated configs), mirroring the Mooncake extension pattern. - ext/AbstractPPLForwardDiffExt.jl: full extension with prepare, value_and_gradient!!, value_and_jacobian!!, and value_gradient_and_hessian!! implementations - Project.toml: ForwardDiff + DiffResults as weakdeps with extension registration and compat entries - test/ext/forwarddiff/: dedicated test environment running all standard test cases plus context and empty-input tests Agent-Logs-Url: https://github.com/TuringLang/AbstractPPL.jl/sessions/2f9552bb-c72d-4891-a973-8ecc68959e06 Co-authored-by: yebai <3279477+yebai@users.noreply.github.com> --- Project.toml | 5 + ext/AbstractPPLForwardDiffExt.jl | 221 ++++++++++++++++++++++++++++++ test/ext/forwarddiff/Project.toml | 13 ++ test/ext/forwarddiff/main.jl | 76 ++++++++++ 4 files changed, 315 insertions(+) create mode 100644 ext/AbstractPPLForwardDiffExt.jl create mode 100644 test/ext/forwarddiff/Project.toml create mode 100644 test/ext/forwarddiff/main.jl diff --git a/Project.toml b/Project.toml index e0b3fa4c..f8eb8da0 100644 --- a/Project.toml +++ b/Project.toml @@ -19,14 +19,17 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [weakdeps] +DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] AbstractPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] AbstractPPLDistributionsExt = ["Distributions", "LinearAlgebra"] +AbstractPPLForwardDiffExt = ["ForwardDiff", "DiffResults"] AbstractPPLMooncakeExt = ["Mooncake"] AbstractPPLTestExt = ["Test"] @@ -36,8 +39,10 @@ AbstractMCMC = "2, 3, 4, 5" Accessors = "0.1" BangBang = "0.4" DensityInterface = "0.4" +DiffResults = "1" DifferentiationInterface = "0.6, 0.7" Distributions = "0.25" +ForwardDiff = "0.10, 1" JSON = "0.19 - 0.21, 1" LinearAlgebra = "<0.0.1, 1" MacroTools = "0.5" diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl new file mode 100644 index 00000000..4c35d37b --- /dev/null +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -0,0 +1,221 @@ +module AbstractPPLForwardDiffExt + +using AbstractPPL: AbstractPPL +using AbstractPPL.Evaluators: + Evaluators, Prepared, VectorEvaluator, _ad_output_arity +using ADTypes: AutoForwardDiff +using ForwardDiff: ForwardDiff +using DiffResults: DiffResults + +# ─── Chunk helper ───────────────────────────────────────────────────────────── +# `AutoForwardDiff{CS}` carries the chunk size as a type parameter. `nothing` +# means "let ForwardDiff pick automatically". +_fd_chunk(::AutoForwardDiff{nothing}, x) = ForwardDiff.Chunk(x) +_fd_chunk(::AutoForwardDiff{CS}, _) where {CS} = ForwardDiff.Chunk{CS}() + +# ─── Cache types ────────────────────────────────────────────────────────────── +# Each cache pre-allocates the `DiffResults` buffer and the ForwardDiff config +# so the hot path is allocation-free (modulo ForwardDiff's internals). + +struct FDGradientCache{R,C} + result::R + config::C +end + +struct FDJacobianCache{R,C} + result::R + config::C +end + +struct FDHessianCache{R,C,GR,GC} + result::R + config::C + gradient_result::GR + gradient_config::GC +end + +# ─── prepare (vector input) ────────────────────────────────────────────────── +""" + prepare(adtype::AutoForwardDiff, problem, x; check_dims=true, context::Tuple=(), order=1) + +Prepare a ForwardDiff gradient, Jacobian, or Hessian evaluator for a vector +input. `order=1` (default) picks gradient/Jacobian by output arity; +`order=2` builds Hessian machinery (`value_gradient_and_hessian!!`) and +requires a scalar-valued problem. + +`context` follows the base `prepare` contract — the prepared evaluator +computes `problem(x, context...)` with AD differentiating only `x`. +""" +function AbstractPPL.prepare( + adtype::AutoForwardDiff, + problem, + x::AbstractVector{<:Real}; + check_dims::Bool=true, + context::Tuple=(), + order::Int=1, +) + Evaluators._validate_ad_order(order) + evaluator = AbstractPPL.prepare(problem, x; check_dims, context)::VectorEvaluator + arity = _ad_output_arity(evaluator(x)) + chunk = _fd_chunk(adtype, x) + + if order == 2 + arity === :scalar || Evaluators._throw_hessian_needs_scalar() + length(x) == 0 && return Prepared( + adtype, evaluator, FDHessianCache(nothing, nothing, nothing, nothing), Val(2) + ) + # Hessian cache: DiffResults buffer for value + gradient + hessian. + hess_result = DiffResults.MutableDiffResult( + zero(eltype(x)), (similar(x), similar(x, length(x), length(x))) + ) + # ForwardDiff.HessianConfig needs the result buffer. + hess_config = ForwardDiff.HessianConfig( + _fd_target(evaluator), hess_result, x, chunk + ) + # Separate gradient cache for order=1 queries on an order=2 prep. + grad_result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) + grad_config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) + cache = FDHessianCache(hess_result, hess_config, grad_result, grad_config) + return Prepared(adtype, evaluator, cache, Val(2)) + end + + if arity === :scalar + length(x) == 0 && return Prepared( + adtype, evaluator, FDGradientCache(nothing, nothing) + ) + result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) + config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) + return Prepared(adtype, evaluator, FDGradientCache(result, config)) + else + length(x) == 0 && return Prepared( + adtype, evaluator, FDJacobianCache(nothing, nothing) + ) + y = evaluator(x) + result = DiffResults.MutableDiffResult(similar(y), (similar(y, length(y), length(x)),)) + config = ForwardDiff.JacobianConfig(_fd_target(evaluator), x, chunk) + return Prepared(adtype, evaluator, FDJacobianCache(result, config)) + end +end + +# ─── Target closure ────────────────────────────────────────────────────────── +# ForwardDiff differentiates a single-argument function; context is closed over. +@inline _fd_target(e::VectorEvaluator) = Base.Fix2(_fd_call, e) +@inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) + +# ─── value_and_gradient!! ──────────────────────────────────────────────────── + +# Empty-input shortcut. +@inline function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache{Nothing}}, + x::AbstractVector{T}, +) where {T<:Real} + Evaluators._check_ad_input(p.evaluator, x) + return (p.evaluator(x), T[]) +end + +# Scalar-gradient hot path. +@inline function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache}, + x::AbstractVector{<:Real}, +) + Evaluators._check_ad_input(p.evaluator, x) + ForwardDiff.gradient!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) +end + +# Arity mismatch: vector-valued problem queried for gradient. +@inline function AbstractPPL.value_and_gradient!!( + ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, + ::AbstractVector{<:Real}, +) + return Evaluators._throw_gradient_needs_scalar() +end + +# Order=2 prep: empty-input shortcut. +@inline function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, + x::AbstractVector{T}, +) where {T<:Real} + Evaluators._check_ad_input(p.evaluator, x) + return (p.evaluator(x), T[]) +end + +# Order=2 prep: use the dedicated gradient cache to skip Hessian work. +@inline function AbstractPPL.value_and_gradient!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, + x::AbstractVector{<:Real}, +) + Evaluators._check_ad_input(p.evaluator, x) + ForwardDiff.gradient!( + p.cache.gradient_result, _fd_target(p.evaluator), x, p.cache.gradient_config + ) + return ( + DiffResults.value(p.cache.gradient_result), + DiffResults.gradient(p.cache.gradient_result), + ) +end + +# ─── value_and_jacobian!! ──────────────────────────────────────────────────── + +# Arity mismatch: scalar-valued problem queried for jacobian. +@inline function AbstractPPL.value_and_jacobian!!( + ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDHessianCache}}, + ::AbstractVector{<:Real}, +) + return Evaluators._throw_jacobian_needs_vector() +end + +# Empty-input shortcut. +@inline function AbstractPPL.value_and_jacobian!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache{Nothing}}, + x::AbstractVector{T}, +) where {T<:Real} + Evaluators._check_ad_input(p.evaluator, x) + val = p.evaluator(x) + return (val, similar(x, length(val), 0)) +end + +# Jacobian hot path. +@inline function AbstractPPL.value_and_jacobian!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, + x::AbstractVector{<:Real}, +) + Evaluators._check_ad_input(p.evaluator, x) + ForwardDiff.jacobian!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + return (DiffResults.value(p.cache.result), DiffResults.jacobian(p.cache.result)) +end + +# ─── value_gradient_and_hessian!! ──────────────────────────────────────────── + +# Order=1 prep rejected for Hessian. +@inline function AbstractPPL.value_gradient_and_hessian!!( + ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDJacobianCache}}, + ::AbstractVector{<:Real}, +) + return Evaluators._throw_hessian_needs_order_2_prep() +end + +# Empty-input shortcut. +@inline function AbstractPPL.value_gradient_and_hessian!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, + x::AbstractVector{T}, +) where {T<:Real} + Evaluators._check_ad_input(p.evaluator, x) + return (p.evaluator(x), T[], similar(x, 0, 0)) +end + +# Hessian hot path. +@inline function AbstractPPL.value_gradient_and_hessian!!( + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, + x::AbstractVector{<:Real}, +) + Evaluators._check_ad_input(p.evaluator, x) + ForwardDiff.hessian!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + return ( + DiffResults.value(p.cache.result), + DiffResults.gradient(p.cache.result), + DiffResults.hessian(p.cache.result), + ) +end + +end # module diff --git a/test/ext/forwarddiff/Project.toml b/test/ext/forwarddiff/Project.toml new file mode 100644 index 00000000..7666f241 --- /dev/null +++ b/test/ext/forwarddiff/Project.toml @@ -0,0 +1,13 @@ +[deps] +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[compat] +ADTypes = "1" +DiffResults = "1" +ForwardDiff = "0.10, 1" +julia = "1.10" diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl new file mode 100644 index 00000000..bca8c103 --- /dev/null +++ b/test/ext/forwarddiff/main.jl @@ -0,0 +1,76 @@ +using Pkg +Pkg.activate(@__DIR__) +Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) +Pkg.instantiate() + +using AbstractPPL: + AbstractPPL, + prepare, + run_testcases, + value_and_gradient!!, + value_and_jacobian!!, + value_gradient_and_hessian!!, + order +using ADTypes: AutoForwardDiff +using ForwardDiff +using DiffResults +using Test + +@testset "AbstractPPLForwardDiffExt" begin + @testset "ForwardDiff (default chunk)" begin + run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + run_testcases(Val(:hessian); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + run_testcases(Val(:edge); adtype=AutoForwardDiff()) + end + + @testset "ForwardDiff (explicit chunk)" begin + run_testcases(Val(:vector); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6) + run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6) + end + + @testset "context-lowered gradient" begin + raw_logdensity(x::AbstractVector{<:Real}, offset) = -0.5 * (x[1] - offset)^2 + + x = [0.3] + ad = AutoForwardDiff() + + lowered = prepare(ad, raw_logdensity, x; check_dims=false, context=(0.1,)) + + # `prepared(x)` evaluates `raw_logdensity(x, context...)`. + @test lowered(x) == raw_logdensity(x, 0.1) + + # Gradient differentiates only w.r.t. `x`. + val, grad = value_and_gradient!!(lowered, x) + @test val ≈ raw_logdensity(x, 0.1) + @test grad ≈ [-(x[1] - 0.1)] atol = 1e-10 + + # Jacobian on a scalar-only lowered cache surfaces our arity-mismatch error. + @test_throws r"vector-valued" value_and_jacobian!!(lowered, x) + end + + @testset "empty input" begin + # Gradient with zero-length input. + ad = AutoForwardDiff() + f_scalar(x::AbstractVector) = sum(x; init=0.0) + prep = prepare(ad, f_scalar, Float64[]) + val, grad = value_and_gradient!!(prep, Float64[]) + @test val == 0.0 + @test grad == Float64[] + + # Jacobian with zero-length input. + f_vec(x::AbstractVector) = [sum(x; init=0.0), 1.0] + prep_j = prepare(ad, f_vec, Float64[]) + val_j, jac = value_and_jacobian!!(prep_j, Float64[]) + @test val_j == [0.0, 1.0] + @test size(jac) == (2, 0) + + # Hessian with zero-length input. + prep_h = prepare(ad, f_scalar, Float64[]; order=2) + @test order(prep_h) == 2 + val_h, grad_h, hess_h = value_gradient_and_hessian!!(prep_h, Float64[]) + @test val_h == 0.0 + @test grad_h == Float64[] + @test size(hess_h) == (0, 0) + end +end From 0d98973e421112143df4d1f23717c0e10841c48e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 14:55:22 +0100 Subject: [PATCH 03/21] Fix CI: format ext files and route DI cache structural test through AutoReverseDiff Run JuliaFormatter on ext/AbstractPPLForwardDiffExt.jl and test/ext/forwarddiff/main.jl to satisfy the Format CI job. The "DI cache encodes the call mode as a type parameter" testset asserted `DIGradientCache{0}` and `DIGradientCache{2}` cache types for `AutoForwardDiff`, but the new direct `AbstractPPLForwardDiffExt` path now takes precedence over DI when both extensions are loaded. `AutoReverseDiff()` (non-compiled) exercises the same DI constants path and keeps the structural assertion meaningful. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 26 +++++++++++++---------- test/ext/differentiationinterface/main.jl | 8 ++++--- test/ext/forwarddiff/main.jl | 8 +++++-- 3 files changed, 26 insertions(+), 16 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index 4c35d37b..f7b45337 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -1,8 +1,7 @@ module AbstractPPLForwardDiffExt using AbstractPPL: AbstractPPL -using AbstractPPL.Evaluators: - Evaluators, Prepared, VectorEvaluator, _ad_output_arity +using AbstractPPL.Evaluators: Evaluators, Prepared, VectorEvaluator, _ad_output_arity using ADTypes: AutoForwardDiff using ForwardDiff: ForwardDiff using DiffResults: DiffResults @@ -62,7 +61,10 @@ function AbstractPPL.prepare( if order == 2 arity === :scalar || Evaluators._throw_hessian_needs_scalar() length(x) == 0 && return Prepared( - adtype, evaluator, FDHessianCache(nothing, nothing, nothing, nothing), Val(2) + adtype, + evaluator, + FDHessianCache(nothing, nothing, nothing, nothing), + Val(2), ) # Hessian cache: DiffResults buffer for value + gradient + hessian. hess_result = DiffResults.MutableDiffResult( @@ -80,18 +82,18 @@ function AbstractPPL.prepare( end if arity === :scalar - length(x) == 0 && return Prepared( - adtype, evaluator, FDGradientCache(nothing, nothing) - ) + length(x) == 0 && + return Prepared(adtype, evaluator, FDGradientCache(nothing, nothing)) result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) return Prepared(adtype, evaluator, FDGradientCache(result, config)) else - length(x) == 0 && return Prepared( - adtype, evaluator, FDJacobianCache(nothing, nothing) - ) + length(x) == 0 && + return Prepared(adtype, evaluator, FDJacobianCache(nothing, nothing)) y = evaluator(x) - result = DiffResults.MutableDiffResult(similar(y), (similar(y, length(y), length(x)),)) + result = DiffResults.MutableDiffResult( + similar(y), (similar(y, length(y), length(x)),) + ) config = ForwardDiff.JacobianConfig(_fd_target(evaluator), x, chunk) return Prepared(adtype, evaluator, FDJacobianCache(result, config)) end @@ -189,7 +191,9 @@ end # Order=1 prep rejected for Hessian. @inline function AbstractPPL.value_gradient_and_hessian!!( - ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDJacobianCache}}, + ::Prepared{ + <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDJacobianCache} + }, ::AbstractVector{<:Real}, ) return Evaluators._throw_hessian_needs_order_2_prep() diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index dff6473e..53a4fba6 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -41,13 +41,15 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) # The DI cache types' `Mode` parameter is either `:closure` (compiled-tape # ReverseDiff) or the integer context length on the constants path. The # constants-path integer also documents how many `DI.Constant`s the AD - # call passes. + # call passes. `AutoReverseDiff()` (non-compiled) is used here because the + # direct `AbstractPPLForwardDiffExt` path takes precedence over DI for + # `AutoForwardDiff` when both extensions are loaded. @testset "DI cache encodes the call mode as a type parameter" begin x = [1.0, 2.0, 3.0] - prep_noctx = prepare(AutoForwardDiff(), quadratic, x) + prep_noctx = prepare(AutoReverseDiff(), quadratic, x) prep_closure = prepare(AutoReverseDiff(; compile=true), quadratic, x) affine(y, a, b) = a * sum(abs2, y) + b - prep_ctx = prepare(AutoForwardDiff(), affine, x; context=(2.0, 1.0)) + prep_ctx = prepare(AutoReverseDiff(), affine, x; context=(2.0, 1.0)) @test prep_noctx.cache isa DIExt.DIGradientCache{0} @test prep_closure.cache isa DIExt.DIGradientCache{:closure} diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index bca8c103..4a2da56b 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -25,8 +25,12 @@ using Test end @testset "ForwardDiff (explicit chunk)" begin - run_testcases(Val(:vector); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6) - run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6) + run_testcases( + Val(:vector); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6 + ) + run_testcases( + Val(:cache_reuse); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6 + ) end @testset "context-lowered gradient" begin From 3f9ad39100752f99ee2ed5efe23308c017cb5956 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:00:09 +0100 Subject: [PATCH 04/21] Apply scrutinise findings to ForwardDiff extension - Drop section banners and WHAT-comments in the extension; keep WHYs (chunk-size dispatch, fresh `Fix2` per call is Tag-type-stable, separate gradient cache on the order=2 prep). - Tighten the `prepare` docstring's second paragraph. - Remove the "empty input" testset: `run_testcases(Val(:vector))` and `run_testcases(Val(:hessian))` already cover zero-length input for every arity / order combination via `AbstractPPLTestExt`. - Remove the trailing arity-mismatch `@test_throws` from the "context-lowered gradient" testset: `run_testcases(Val(:edge))` already covers "jacobian of scalar output". - Drop now-unused imports (`value_and_jacobian!!`, `value_gradient_and_hessian!!`, `order`, `DiffResults`). Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 51 ++++++++++---------------------- test/ext/forwarddiff/main.jl | 41 ++----------------------- 2 files changed, 17 insertions(+), 75 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index f7b45337..12fdf30e 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -6,16 +6,14 @@ using ADTypes: AutoForwardDiff using ForwardDiff: ForwardDiff using DiffResults: DiffResults -# ─── Chunk helper ───────────────────────────────────────────────────────────── -# `AutoForwardDiff{CS}` carries the chunk size as a type parameter. `nothing` -# means "let ForwardDiff pick automatically". +# `AutoForwardDiff{CS}` carries the chunk size as a type parameter; `nothing` +# defers the choice to ForwardDiff. _fd_chunk(::AutoForwardDiff{nothing}, x) = ForwardDiff.Chunk(x) _fd_chunk(::AutoForwardDiff{CS}, _) where {CS} = ForwardDiff.Chunk{CS}() -# ─── Cache types ────────────────────────────────────────────────────────────── -# Each cache pre-allocates the `DiffResults` buffer and the ForwardDiff config -# so the hot path is allocation-free (modulo ForwardDiff's internals). - +# Three cache types so arity (scalar/vector) and order (1/2) are encoded in the +# type and dispatch resolves the hot path without runtime branching. The stored +# `result` aliases the arrays returned by `value_and_*!!`, per the `!!` contract. struct FDGradientCache{R,C} result::R config::C @@ -26,6 +24,8 @@ struct FDJacobianCache{R,C} config::C end +# `gradient_result` / `gradient_config` are kept alongside the Hessian cache so +# `value_and_gradient!!` on an order=2 prep skips the O(n²) Hessian work. struct FDHessianCache{R,C,GR,GC} result::R config::C @@ -33,17 +33,13 @@ struct FDHessianCache{R,C,GR,GC} gradient_config::GC end -# ─── prepare (vector input) ────────────────────────────────────────────────── """ prepare(adtype::AutoForwardDiff, problem, x; check_dims=true, context::Tuple=(), order=1) Prepare a ForwardDiff gradient, Jacobian, or Hessian evaluator for a vector -input. `order=1` (default) picks gradient/Jacobian by output arity; -`order=2` builds Hessian machinery (`value_gradient_and_hessian!!`) and -requires a scalar-valued problem. - -`context` follows the base `prepare` contract — the prepared evaluator -computes `problem(x, context...)` with AD differentiating only `x`. +input. `order=1` (default) picks gradient/Jacobian by output arity; `order=2` +builds Hessian machinery and requires a scalar-valued problem. `context` and +`check_dims` follow the base `prepare` contract. """ function AbstractPPL.prepare( adtype::AutoForwardDiff, @@ -66,15 +62,12 @@ function AbstractPPL.prepare( FDHessianCache(nothing, nothing, nothing, nothing), Val(2), ) - # Hessian cache: DiffResults buffer for value + gradient + hessian. hess_result = DiffResults.MutableDiffResult( zero(eltype(x)), (similar(x), similar(x, length(x), length(x))) ) - # ForwardDiff.HessianConfig needs the result buffer. hess_config = ForwardDiff.HessianConfig( _fd_target(evaluator), hess_result, x, chunk ) - # Separate gradient cache for order=1 queries on an order=2 prep. grad_result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) grad_config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) cache = FDHessianCache(hess_result, hess_config, grad_result, grad_config) @@ -99,14 +92,12 @@ function AbstractPPL.prepare( end end -# ─── Target closure ────────────────────────────────────────────────────────── -# ForwardDiff differentiates a single-argument function; context is closed over. +# ForwardDiff's `*Config` keys its `Tag` on the *type* of the target, so +# constructing a fresh `Fix2` per hot-path call is free — the type matches the +# one captured in the config at prep time. @inline _fd_target(e::VectorEvaluator) = Base.Fix2(_fd_call, e) @inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) -# ─── value_and_gradient!! ──────────────────────────────────────────────────── - -# Empty-input shortcut. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache{Nothing}}, x::AbstractVector{T}, @@ -115,7 +106,6 @@ end return (p.evaluator(x), T[]) end -# Scalar-gradient hot path. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache}, x::AbstractVector{<:Real}, @@ -125,7 +115,8 @@ end return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) end -# Arity mismatch: vector-valued problem queried for gradient. +# Arity-mismatch rejections live on dedicated cache types so dispatch resolves +# the failure mode at compile time. @inline function AbstractPPL.value_and_gradient!!( ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, ::AbstractVector{<:Real}, @@ -133,7 +124,6 @@ end return Evaluators._throw_gradient_needs_scalar() end -# Order=2 prep: empty-input shortcut. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, x::AbstractVector{T}, @@ -142,7 +132,6 @@ end return (p.evaluator(x), T[]) end -# Order=2 prep: use the dedicated gradient cache to skip Hessian work. @inline function AbstractPPL.value_and_gradient!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, x::AbstractVector{<:Real}, @@ -157,9 +146,6 @@ end ) end -# ─── value_and_jacobian!! ──────────────────────────────────────────────────── - -# Arity mismatch: scalar-valued problem queried for jacobian. @inline function AbstractPPL.value_and_jacobian!!( ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDHessianCache}}, ::AbstractVector{<:Real}, @@ -167,7 +153,6 @@ end return Evaluators._throw_jacobian_needs_vector() end -# Empty-input shortcut. @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache{Nothing}}, x::AbstractVector{T}, @@ -177,7 +162,6 @@ end return (val, similar(x, length(val), 0)) end -# Jacobian hot path. @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, x::AbstractVector{<:Real}, @@ -187,9 +171,6 @@ end return (DiffResults.value(p.cache.result), DiffResults.jacobian(p.cache.result)) end -# ─── value_gradient_and_hessian!! ──────────────────────────────────────────── - -# Order=1 prep rejected for Hessian. @inline function AbstractPPL.value_gradient_and_hessian!!( ::Prepared{ <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDJacobianCache} @@ -199,7 +180,6 @@ end return Evaluators._throw_hessian_needs_order_2_prep() end -# Empty-input shortcut. @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, x::AbstractVector{T}, @@ -208,7 +188,6 @@ end return (p.evaluator(x), T[], similar(x, 0, 0)) end -# Hessian hot path. @inline function AbstractPPL.value_gradient_and_hessian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, x::AbstractVector{<:Real}, diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 4a2da56b..b2baac4f 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -3,17 +3,9 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: - AbstractPPL, - prepare, - run_testcases, - value_and_gradient!!, - value_and_jacobian!!, - value_gradient_and_hessian!!, - order +using AbstractPPL: AbstractPPL, prepare, run_testcases, value_and_gradient!! using ADTypes: AutoForwardDiff using ForwardDiff -using DiffResults using Test @testset "AbstractPPLForwardDiffExt" begin @@ -33,6 +25,7 @@ using Test ) end + # `run_testcases` doesn't exercise `context=`; this fills that gap. @testset "context-lowered gradient" begin raw_logdensity(x::AbstractVector{<:Real}, offset) = -0.5 * (x[1] - offset)^2 @@ -41,40 +34,10 @@ using Test lowered = prepare(ad, raw_logdensity, x; check_dims=false, context=(0.1,)) - # `prepared(x)` evaluates `raw_logdensity(x, context...)`. @test lowered(x) == raw_logdensity(x, 0.1) - # Gradient differentiates only w.r.t. `x`. val, grad = value_and_gradient!!(lowered, x) @test val ≈ raw_logdensity(x, 0.1) @test grad ≈ [-(x[1] - 0.1)] atol = 1e-10 - - # Jacobian on a scalar-only lowered cache surfaces our arity-mismatch error. - @test_throws r"vector-valued" value_and_jacobian!!(lowered, x) - end - - @testset "empty input" begin - # Gradient with zero-length input. - ad = AutoForwardDiff() - f_scalar(x::AbstractVector) = sum(x; init=0.0) - prep = prepare(ad, f_scalar, Float64[]) - val, grad = value_and_gradient!!(prep, Float64[]) - @test val == 0.0 - @test grad == Float64[] - - # Jacobian with zero-length input. - f_vec(x::AbstractVector) = [sum(x; init=0.0), 1.0] - prep_j = prepare(ad, f_vec, Float64[]) - val_j, jac = value_and_jacobian!!(prep_j, Float64[]) - @test val_j == [0.0, 1.0] - @test size(jac) == (2, 0) - - # Hessian with zero-length input. - prep_h = prepare(ad, f_scalar, Float64[]; order=2) - @test order(prep_h) == 2 - val_h, grad_h, hess_h = value_gradient_and_hessian!!(prep_h, Float64[]) - @test val_h == 0.0 - @test grad_h == Float64[] - @test size(hess_h) == (0, 0) end end From cb6e14c3ca14714ed2d30a79b4debb104b055519 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:13:37 +0100 Subject: [PATCH 05/21] Wire ext/forwarddiff into CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `ext/forwarddiff` to `VALID_LABELS` in `test/run_extras.jl` and to the CI ext matrix so the chunk-size and context tests this branch introduces actually run on CI (they were silently skipped before — `AutoForwardDiff` was only exercised via the DI test env). Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/CI.yml | 1 + test/run_extras.jl | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 6a4189a8..a7defcfd 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -61,6 +61,7 @@ jobs: matrix: label: - ext/differentiationinterface + - ext/forwarddiff - ext/mooncake version: - '1' diff --git a/test/run_extras.jl b/test/run_extras.jl index cd2c157e..e84b5ea9 100644 --- a/test/run_extras.jl +++ b/test/run_extras.jl @@ -2,9 +2,10 @@ # # Usage (from the repo root): # LABEL=ext/differentiationinterface julia test/run_extras.jl +# LABEL=ext/forwarddiff julia test/run_extras.jl # LABEL=ext/mooncake julia test/run_extras.jl -const VALID_LABELS = ("ext/differentiationinterface", "ext/mooncake") +const VALID_LABELS = ("ext/differentiationinterface", "ext/forwarddiff", "ext/mooncake") label = get(ENV, "LABEL", nothing) label in VALID_LABELS || From dc313f7761645d3d1fc54932a52aa9265dc8d821 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:16:21 +0100 Subject: [PATCH 06/21] Unify FD gradient/Jacobian/Hessian caches into FDCache{A} MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace `FDGradientCache`, `FDJacobianCache`, and `FDHessianCache` with one parametric `FDCache{A,R,C,GR,GC}` keyed on an arity/order `Symbol` `A ∈ (:scalar, :vector, :hessian)`, mirroring the `MooncakeCache{A}` pattern. Hot paths and arity-mismatch rejections dispatch on the tag at compile time exactly as before; `result::Nothing` remains the empty-input sentinel. Verified type-stable on all four `!!` entries. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 95 ++++++++++++++++---------------- 1 file changed, 47 insertions(+), 48 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index 12fdf30e..d766fe69 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -11,26 +11,25 @@ using DiffResults: DiffResults _fd_chunk(::AutoForwardDiff{nothing}, x) = ForwardDiff.Chunk(x) _fd_chunk(::AutoForwardDiff{CS}, _) where {CS} = ForwardDiff.Chunk{CS}() -# Three cache types so arity (scalar/vector) and order (1/2) are encoded in the -# type and dispatch resolves the hot path without runtime branching. The stored -# `result` aliases the arrays returned by `value_and_*!!`, per the `!!` contract. -struct FDGradientCache{R,C} - result::R - config::C -end - -struct FDJacobianCache{R,C} - result::R - config::C -end - -# `gradient_result` / `gradient_config` are kept alongside the Hessian cache so -# `value_and_gradient!!` on an order=2 prep skips the O(n²) Hessian work. -struct FDHessianCache{R,C,GR,GC} +# `A::Symbol` ∈ `(:scalar, :vector, :hessian)` encodes both output arity +# (order=1) and order (order=2 ≡ `:hessian`), so dispatch resolves the hot path +# and the arity-mismatch failure modes at compile time without a runtime branch. +# `gradient_result` / `gradient_config` are populated only on `:hessian` caches +# so `value_and_gradient!!` on an order=2 prep skips the O(n²) Hessian work. +# `result::Nothing` is the empty-input sentinel: hot paths dispatch on +# `FDCache{A,Nothing}` to short-circuit before any ForwardDiff call (chunk +# selection `BoundsError`s on length-zero inputs). The stored `result` aliases +# the arrays returned by `value_and_*!!`, per the `!!` contract. +struct FDCache{A,R,C,GR,GC} result::R config::C gradient_result::GR gradient_config::GC + function FDCache{A}( + result::R, config::C, gradient_result::GR=nothing, gradient_config::GC=nothing + ) where {A,R,C,GR,GC} + return new{A,R,C,GR,GC}(result, config, gradient_result, gradient_config) + end end """ @@ -59,7 +58,7 @@ function AbstractPPL.prepare( length(x) == 0 && return Prepared( adtype, evaluator, - FDHessianCache(nothing, nothing, nothing, nothing), + FDCache{:hessian}(nothing, nothing, nothing, nothing), Val(2), ) hess_result = DiffResults.MutableDiffResult( @@ -70,25 +69,25 @@ function AbstractPPL.prepare( ) grad_result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) grad_config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) - cache = FDHessianCache(hess_result, hess_config, grad_result, grad_config) + cache = FDCache{:hessian}(hess_result, hess_config, grad_result, grad_config) return Prepared(adtype, evaluator, cache, Val(2)) end if arity === :scalar length(x) == 0 && - return Prepared(adtype, evaluator, FDGradientCache(nothing, nothing)) + return Prepared(adtype, evaluator, FDCache{:scalar}(nothing, nothing)) result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) - return Prepared(adtype, evaluator, FDGradientCache(result, config)) + return Prepared(adtype, evaluator, FDCache{:scalar}(result, config)) else length(x) == 0 && - return Prepared(adtype, evaluator, FDJacobianCache(nothing, nothing)) + return Prepared(adtype, evaluator, FDCache{:vector}(nothing, nothing)) y = evaluator(x) result = DiffResults.MutableDiffResult( similar(y), (similar(y, length(y), length(x)),) ) config = ForwardDiff.JacobianConfig(_fd_target(evaluator), x, chunk) - return Prepared(adtype, evaluator, FDJacobianCache(result, config)) + return Prepared(adtype, evaluator, FDCache{:vector}(result, config)) end end @@ -99,7 +98,11 @@ end @inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) @inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache{Nothing}}, + p::Prepared{ + <:AutoForwardDiff, + <:VectorEvaluator, + <:Union{FDCache{:scalar,Nothing},FDCache{:hessian,Nothing}}, + }, x::AbstractVector{T}, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) @@ -107,7 +110,7 @@ end end @inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDGradientCache}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:scalar}}, x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) @@ -115,25 +118,10 @@ end return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) end -# Arity-mismatch rejections live on dedicated cache types so dispatch resolves -# the failure mode at compile time. -@inline function AbstractPPL.value_and_gradient!!( - ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, - ::AbstractVector{<:Real}, -) - return Evaluators._throw_gradient_needs_scalar() -end - +# Order=2 prep also satisfies the order=1 gradient contract via the dedicated +# gradient cache built at prep time — skips the O(n²) Hessian work. @inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, - x::AbstractVector{T}, -) where {T<:Real} - Evaluators._check_ad_input(p.evaluator, x) - return (p.evaluator(x), T[]) -end - -@inline function AbstractPPL.value_and_gradient!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian}}, x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) @@ -146,15 +134,26 @@ end ) end +# Arity-mismatch rejections live on dedicated cache tags so dispatch resolves +# the failure mode at compile time. +@inline function AbstractPPL.value_and_gradient!!( + ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector}}, + ::AbstractVector{<:Real}, +) + return Evaluators._throw_gradient_needs_scalar() +end + @inline function AbstractPPL.value_and_jacobian!!( - ::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDHessianCache}}, + ::Prepared{ + <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDCache{:scalar},FDCache{:hessian}} + }, ::AbstractVector{<:Real}, ) return Evaluators._throw_jacobian_needs_vector() end @inline function AbstractPPL.value_and_jacobian!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache{Nothing}}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector,Nothing}}, x::AbstractVector{T}, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) @@ -163,7 +162,7 @@ end end @inline function AbstractPPL.value_and_jacobian!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDJacobianCache}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector}}, x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) @@ -173,7 +172,7 @@ end @inline function AbstractPPL.value_gradient_and_hessian!!( ::Prepared{ - <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDGradientCache,FDJacobianCache} + <:AutoForwardDiff,<:VectorEvaluator,<:Union{FDCache{:scalar},FDCache{:vector}} }, ::AbstractVector{<:Real}, ) @@ -181,7 +180,7 @@ end end @inline function AbstractPPL.value_gradient_and_hessian!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache{Nothing}}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian,Nothing}}, x::AbstractVector{T}, ) where {T<:Real} Evaluators._check_ad_input(p.evaluator, x) @@ -189,7 +188,7 @@ end end @inline function AbstractPPL.value_gradient_and_hessian!!( - p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDHessianCache}, + p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:hessian}}, x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) From ea03f724c6920522b719db9aa5dd3f9c654b81b5 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:30:37 +0100 Subject: [PATCH 07/21] Honor AutoForwardDiff tag and probe the problem once in prepare MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ext/AbstractPPLForwardDiffExt.jl`: * Thread `adtype.tag` into every `*Config` constructor via a small `_fd_tag` helper; `nothing` (the ADTypes default) reproduces ForwardDiff's per-constructor default of `Tag(target, eltype(x))`, so callers can now use `AutoForwardDiff(; tag=…)` for nested differentiation through AbstractPPL. * Hoist the arity-probe `evaluator(x)` to a single `y_probe` local and reuse it as the Jacobian-result prototype on the vector branch. The base `prepare` contract promises one prep-time call into `problem`; the vector path was invoking it twice. * Cache `target = _fd_target(evaluator)` once locally rather than reconstructing the `Fix2` per config. `test/ext/forwarddiff/main.jl`: add a regression test asserting the user-supplied tag flows into the stored config's first type parameter. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 27 ++++++++++++++++++--------- test/ext/forwarddiff/main.jl | 10 ++++++++++ 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index d766fe69..a6eee0e6 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -11,6 +11,12 @@ using DiffResults: DiffResults _fd_chunk(::AutoForwardDiff{nothing}, x) = ForwardDiff.Chunk(x) _fd_chunk(::AutoForwardDiff{CS}, _) where {CS} = ForwardDiff.Chunk{CS}() +# A user-supplied `adtype.tag` (for nested differentiation) is threaded into the +# `*Config` constructors; `nothing` (the ADTypes default) reproduces +# ForwardDiff's per-constructor default of `Tag(target, eltype(x))`. +@inline _fd_tag(adtype::AutoForwardDiff, target, x) = + adtype.tag === nothing ? ForwardDiff.Tag(target, eltype(x)) : adtype.tag + # `A::Symbol` ∈ `(:scalar, :vector, :hessian)` encodes both output arity # (order=1) and order (order=2 ≡ `:hessian`), so dispatch resolves the hot path # and the arity-mismatch failure modes at compile time without a runtime branch. @@ -50,8 +56,14 @@ function AbstractPPL.prepare( ) Evaluators._validate_ad_order(order) evaluator = AbstractPPL.prepare(problem, x; check_dims, context)::VectorEvaluator - arity = _ad_output_arity(evaluator(x)) + # Probe the output once: the value classifies arity, and the vector branch + # reuses it as the Jacobian-result prototype. The base `prepare` contract + # promises one prep-time call into `problem`. + y_probe = evaluator(x) + arity = _ad_output_arity(y_probe) chunk = _fd_chunk(adtype, x) + target = _fd_target(evaluator) + tag = _fd_tag(adtype, target, x) if order == 2 arity === :scalar || Evaluators._throw_hessian_needs_scalar() @@ -64,11 +76,9 @@ function AbstractPPL.prepare( hess_result = DiffResults.MutableDiffResult( zero(eltype(x)), (similar(x), similar(x, length(x), length(x))) ) - hess_config = ForwardDiff.HessianConfig( - _fd_target(evaluator), hess_result, x, chunk - ) + hess_config = ForwardDiff.HessianConfig(target, hess_result, x, chunk, tag) grad_result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) - grad_config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) + grad_config = ForwardDiff.GradientConfig(target, x, chunk, tag) cache = FDCache{:hessian}(hess_result, hess_config, grad_result, grad_config) return Prepared(adtype, evaluator, cache, Val(2)) end @@ -77,16 +87,15 @@ function AbstractPPL.prepare( length(x) == 0 && return Prepared(adtype, evaluator, FDCache{:scalar}(nothing, nothing)) result = DiffResults.MutableDiffResult(zero(eltype(x)), (similar(x),)) - config = ForwardDiff.GradientConfig(_fd_target(evaluator), x, chunk) + config = ForwardDiff.GradientConfig(target, x, chunk, tag) return Prepared(adtype, evaluator, FDCache{:scalar}(result, config)) else length(x) == 0 && return Prepared(adtype, evaluator, FDCache{:vector}(nothing, nothing)) - y = evaluator(x) result = DiffResults.MutableDiffResult( - similar(y), (similar(y, length(y), length(x)),) + similar(y_probe), (similar(y_probe, length(y_probe), length(x)),) ) - config = ForwardDiff.JacobianConfig(_fd_target(evaluator), x, chunk) + config = ForwardDiff.JacobianConfig(target, x, chunk, tag) return Prepared(adtype, evaluator, FDCache{:vector}(result, config)) end end diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index b2baac4f..637c1466 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -40,4 +40,14 @@ using Test @test val ≈ raw_logdensity(x, 0.1) @test grad ≈ [-(x[1] - 0.1)] atol = 1e-10 end + + # `AutoForwardDiff(; tag=...)` exists for nested differentiation. Check the + # user-supplied tag is threaded into the ForwardDiff config (the inner + # `*Config` carries the tag in its first type parameter). + @testset "custom AutoForwardDiff tag" begin + struct OuterTag end + custom = ForwardDiff.Tag{OuterTag,Float64}() + prep = prepare(AutoForwardDiff(; tag=custom), x -> sum(abs2, x), [1.0, 2.0]) + @test typeof(prep.cache.config).parameters[1] === typeof(custom) + end end From 45b1d5e1aaacf4885ec40d46c58b758f9eac8ed9 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:34:23 +0100 Subject: [PATCH 08/21] Tidy two minor redundancies in the FD extension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Drop the unused `where {T<:Real}` binding on the empty-input Jacobian method; the non-empty sibling already uses `x::AbstractVector{<:Real}` directly. * Pass two `nothing`s to `FDCache{:hessian}(nothing, nothing)` for the empty-input order=2 cache instead of four — the constructor defaults `gradient_result` and `gradient_config` to `nothing`, so the resulting type is identical and the line is consistent with the `:scalar` / `:vector` empty-input shortcuts above. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index a6eee0e6..21c35198 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -67,12 +67,8 @@ function AbstractPPL.prepare( if order == 2 arity === :scalar || Evaluators._throw_hessian_needs_scalar() - length(x) == 0 && return Prepared( - adtype, - evaluator, - FDCache{:hessian}(nothing, nothing, nothing, nothing), - Val(2), - ) + length(x) == 0 && + return Prepared(adtype, evaluator, FDCache{:hessian}(nothing, nothing), Val(2)) hess_result = DiffResults.MutableDiffResult( zero(eltype(x)), (similar(x), similar(x, length(x), length(x))) ) @@ -163,8 +159,8 @@ end @inline function AbstractPPL.value_and_jacobian!!( p::Prepared{<:AutoForwardDiff,<:VectorEvaluator,<:FDCache{:vector,Nothing}}, - x::AbstractVector{T}, -) where {T<:Real} + x::AbstractVector{<:Real}, +) Evaluators._check_ad_input(p.evaluator, x) val = p.evaluator(x) return (val, similar(x, length(val), 0)) From 29050c0305f62b3c9f08291d3bf1fa2b5f3b949f Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:41:06 +0100 Subject: [PATCH 09/21] Bump to 0.15.2 with HISTORY entry for the FD extension Co-Authored-By: Claude Opus 4.7 (1M context) --- HISTORY.md | 4 ++++ Project.toml | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index 7f5b2346..12112c61 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,7 @@ +## 0.15.2 + +Added `AbstractPPLForwardDiffExt`, a direct ForwardDiff path for `AutoForwardDiff` (gradient, Jacobian, Hessian, `context`, chunk size, custom `tag`). + ## 0.15.1 Added Hessian support to the AD interface. Pass `order=2` to `prepare(adtype, problem, x)` to build a Hessian-capable evaluator. The new `value_gradient_and_hessian!!(prepared, x)` then returns `(value, gradient, hessian)` in a single call. Both the DifferentiationInterface and Mooncake extensions implement this. diff --git a/Project.toml b/Project.toml index f8eb8da0..57c13443 100644 --- a/Project.toml +++ b/Project.toml @@ -3,7 +3,7 @@ uuid = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" keywords = ["probabilistic programming"] license = "MIT" desc = "Common interfaces for probabilistic programming" -version = "0.15.1" +version = "0.15.2" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" From e5d253bba1966730a68ebc45c24766bcfe4dd6dd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 15:51:29 +0100 Subject: [PATCH 10/21] Inline _fd_target and clarify why _fd_call must stay top-level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_fd_target(e)` was just `Base.Fix2(_fd_call, e)` — inline the five call sites and drop the helper. `_fd_call` stays as a top-level named function: ForwardDiff's `*Config` keys its `Tag` on the target type, and a closure built inside one method would have a different type from one built inside another, desyncing the per-call target from the config captured at prep time. Reworded the comment to make that constraint (not the harmless cost of per-call `Fix2`) the WHY. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index 21c35198..18f9c112 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -62,7 +62,7 @@ function AbstractPPL.prepare( y_probe = evaluator(x) arity = _ad_output_arity(y_probe) chunk = _fd_chunk(adtype, x) - target = _fd_target(evaluator) + target = Base.Fix2(_fd_call, evaluator) tag = _fd_tag(adtype, target, x) if order == 2 @@ -96,10 +96,11 @@ function AbstractPPL.prepare( end end -# ForwardDiff's `*Config` keys its `Tag` on the *type* of the target, so -# constructing a fresh `Fix2` per hot-path call is free — the type matches the -# one captured in the config at prep time. -@inline _fd_target(e::VectorEvaluator) = Base.Fix2(_fd_call, e) +# Top-level so `typeof(_fd_call)` is stable across `prepare` and the hot paths. +# ForwardDiff's `*Config` keys its `Tag` on the target type; a closure built +# inside one method would have a different type from one built inside another, +# desyncing the per-call `Base.Fix2(_fd_call, evaluator)` target from the +# config captured at prep time. @inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) @inline function AbstractPPL.value_and_gradient!!( @@ -119,7 +120,9 @@ end x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) - ForwardDiff.gradient!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + ForwardDiff.gradient!( + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + ) return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) end @@ -131,7 +134,10 @@ end ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.gradient!( - p.cache.gradient_result, _fd_target(p.evaluator), x, p.cache.gradient_config + p.cache.gradient_result, + Base.Fix2(_fd_call, p.evaluator), + x, + p.cache.gradient_config, ) return ( DiffResults.value(p.cache.gradient_result), @@ -171,7 +177,9 @@ end x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) - ForwardDiff.jacobian!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + ForwardDiff.jacobian!( + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + ) return (DiffResults.value(p.cache.result), DiffResults.jacobian(p.cache.result)) end @@ -197,7 +205,9 @@ end x::AbstractVector{<:Real}, ) Evaluators._check_ad_input(p.evaluator, x) - ForwardDiff.hessian!(p.cache.result, _fd_target(p.evaluator), x, p.cache.config) + ForwardDiff.hessian!( + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + ) return ( DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result), From 3be777f75e7e51366e6c4ee8c8754bf085f15bb8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 16:41:43 +0100 Subject: [PATCH 11/21] Share :allocations and :type_stability groups via AbstractPPLTestExt The pending changes in `test/ext/forwarddiff/main.jl` and `test/ext/mooncake/main.jl` had duplicated the same helper functions and testset bodies for "allocation-free hot paths" and "type-stable hot paths". Lift the shared logic into `AbstractPPLTestExt`: * `IdentityProblem`: allocation-free vector-output problem (avoids `VectorValuedProblem`'s result allocation masking AD-path allocations). * `_inferred_*` helpers wrap `@inferred` so it can be marked broken via `@test_broken`. * `run_testcases(Val(:allocations); ...)`: `@allocated == 0` checks on scalar gradient and vector Jacobian, with `gradient_broken` / `jacobian_broken` kwargs for backends with known regressions. * `run_testcases(Val(:type_stability); ...)`: `@inferred` checks on gradient/Jacobian/Hessian hot paths, with matching `*_broken` kwargs. Both extension test files now invoke the shared groups; Mooncake passes `jacobian_broken=true` for `:allocations` (both modes) and for `:type_stability` only on `AutoMooncakeForward` (`Tuple{Any, Union{Array{T,3}, Matrix}}` inference). Docstring on `generate_testcases` updated to list the new keys. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 83 +++++++++++++++++++++++++++++++++++- src/AbstractPPL.jl | 7 ++- test/ext/forwarddiff/main.jl | 2 + test/ext/mooncake/main.jl | 10 +++++ 4 files changed, 99 insertions(+), 3 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index c5c74727..f92d58ad 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -1,7 +1,7 @@ module AbstractPPLTestExt using AbstractPPL: AbstractPPL, generate_testcases, run_testcases -using Test: @test, @test_throws, @testset +using Test: @inferred, @test, @test_broken, @test_throws, @testset struct QuadraticProblem end (::QuadraticProblem)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) @@ -9,6 +9,11 @@ struct QuadraticProblem end struct VectorValuedProblem end (::VectorValuedProblem)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] +# Allocation-free vector-output problem for the `:allocations` group: +# `VectorValuedProblem` allocates its result vector, masking AD-path allocations. +struct IdentityProblem end +(::IdentityProblem)(x::AbstractVector{<:Real}) = x + struct ValueCase name::String f::Any @@ -330,4 +335,80 @@ function AbstractPPL.run_testcases( return nothing end +# Helpers for the `:type_stability` group: `@inferred` is a syntactic macro, so wrap +# each AD entry in a tiny named function that returns `true` on success — that +# value lets `@test` / `@test_broken` evaluate the call uniformly. +function _inferred_gradient(prepared, x) + (@inferred AbstractPPL.value_and_gradient!!(prepared, x); true) +end +function _inferred_jacobian(prepared, x) + (@inferred AbstractPPL.value_and_jacobian!!(prepared, x); true) +end +function _inferred_hessian(prepared, x) + (@inferred AbstractPPL.value_gradient_and_hessian!!(prepared, x); true) +end + +# Backends with known broken paths (e.g. Mooncake's forward-mode Jacobian) +# pass `*_broken=true` to mark the assertion as broken instead of failing. +function AbstractPPL.run_testcases( + ::Val{:allocations}, + prepare_fn=AbstractPPL.prepare; + adtype, + gradient_broken::Bool=false, + jacobian_broken::Bool=false, +) + x = [1.0, 2.0, 3.0] + @testset "scalar gradient" begin + prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) + AbstractPPL.value_and_gradient!!(prepared, x) # warm up + allocs = @allocated AbstractPPL.value_and_gradient!!(prepared, x) + gradient_broken ? (@test_broken allocs == 0) : (@test allocs == 0) + end + @testset "vector jacobian" begin + prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) + AbstractPPL.value_and_jacobian!!(prepared, x) + allocs = @allocated AbstractPPL.value_and_jacobian!!(prepared, x) + jacobian_broken ? (@test_broken allocs == 0) : (@test allocs == 0) + end + return nothing +end + +function AbstractPPL.run_testcases( + ::Val{:type_stability}, + prepare_fn=AbstractPPL.prepare; + adtype, + gradient_broken::Bool=false, + jacobian_broken::Bool=false, + hessian_broken::Bool=false, +) + x = [1.0, 2.0, 3.0] + @testset "scalar gradient" begin + prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) + if gradient_broken + (@test_broken _inferred_gradient(prepared, x)) + else + (@test _inferred_gradient(prepared, x)) + end + end + @testset "vector jacobian" begin + prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) + if jacobian_broken + (@test_broken _inferred_jacobian(prepared, x)) + else + (@test _inferred_jacobian(prepared, x)) + end + end + @testset "hessian" begin + prepared = prepare_fn( + adtype, QuadraticProblem(), zeros(3); check_dims=false, order=2 + ) + if hessian_broken + (@test_broken _inferred_hessian(prepared, x)) + else + (@test _inferred_hessian(prepared, x)) + end + end + return nothing +end + end # module diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 2a32494d..9b1f44cd 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -23,8 +23,11 @@ not redefine these): `:vector` for value/gradient/jacobian round-trips on vector-input evaluators; `:hessian` for `order=2` value/gradient/Hessian round-trips on vector-input scalar-output evaluators; `:namedtuple` for `NamedTuple`-input evaluators; `:edge` for error-path cases; `:cache_reuse` -for repeated calls against a single prepared evaluator. Downstream packages -may add other keys. +for repeated calls against a single prepared evaluator; `:allocations` and +`:type_stability` for `@allocated == 0` and `@inferred` checks on the AD hot paths +(both accept `gradient_broken`, `jacobian_broken`, and (`:type_stability` only) +`hessian_broken` kwargs for backends with known broken paths). Downstream +packages may add other keys. """ function generate_testcases end diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 637c1466..2527e399 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -14,6 +14,8 @@ using Test run_testcases(Val(:hessian); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:edge); adtype=AutoForwardDiff()) + run_testcases(Val(:allocations); adtype=AutoForwardDiff()) + run_testcases(Val(:type_stability); adtype=AutoForwardDiff()) end @testset "ForwardDiff (explicit chunk)" begin diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index 855d5f5b..f3d8f5e4 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -22,6 +22,16 @@ using Test # AutoMooncakeForward routes through the same generic Hessian path # since `Mooncake.prepare_hessian_cache` is mode-agnostic. run_testcases(Val(:hessian); adtype=adtype, atol=1e-6, rtol=1e-6) + # Mooncake's `value_and_jacobian!!` currently allocates fresh + # cotangent/Jacobian buffers each call, and the forward-mode + # Jacobian return type infers as `Tuple{Any, Union{Array{T,3}, + # Matrix}}`. Mark those known-broken; the other paths must hold. + run_testcases(Val(:allocations); adtype=adtype, jacobian_broken=true) + run_testcases( + Val(:type_stability); + adtype=adtype, + jacobian_broken=adtype isa AutoMooncakeForward, + ) end end From c246d4af10856c7557a9c0d44a1871673788ca5a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 16:45:18 +0100 Subject: [PATCH 12/21] Apply scrutinise findings to the shared test groups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Generalize the `*_broken` comment: previously cited only "Mooncake's forward-mode Jacobian", but the kwargs cover other regressions too (Mooncake's `value_and_jacobian!!` allocates on every call across both modes; only the forward-mode Jacobian *inference* is broken). * Unify the `:allocations` vs `:type_stability` branch style — both now use the same `if/else` form rather than the ternary the former was using inconsistently. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index f92d58ad..35eb0131 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -348,8 +348,9 @@ function _inferred_hessian(prepared, x) (@inferred AbstractPPL.value_gradient_and_hessian!!(prepared, x); true) end -# Backends with known broken paths (e.g. Mooncake's forward-mode Jacobian) -# pass `*_broken=true` to mark the assertion as broken instead of failing. +# Backends with known regressions (e.g. Mooncake's allocating +# `value_and_jacobian!!`, or its forward-mode Jacobian inference) pass +# `*_broken=true` to mark the assertion as broken instead of failing. function AbstractPPL.run_testcases( ::Val{:allocations}, prepare_fn=AbstractPPL.prepare; @@ -362,13 +363,21 @@ function AbstractPPL.run_testcases( prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) AbstractPPL.value_and_gradient!!(prepared, x) # warm up allocs = @allocated AbstractPPL.value_and_gradient!!(prepared, x) - gradient_broken ? (@test_broken allocs == 0) : (@test allocs == 0) + if gradient_broken + @test_broken allocs == 0 + else + @test allocs == 0 + end end @testset "vector jacobian" begin prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) AbstractPPL.value_and_jacobian!!(prepared, x) allocs = @allocated AbstractPPL.value_and_jacobian!!(prepared, x) - jacobian_broken ? (@test_broken allocs == 0) : (@test allocs == 0) + if jacobian_broken + @test_broken allocs == 0 + else + @test allocs == 0 + end end return nothing end @@ -385,17 +394,17 @@ function AbstractPPL.run_testcases( @testset "scalar gradient" begin prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) if gradient_broken - (@test_broken _inferred_gradient(prepared, x)) + @test_broken _inferred_gradient(prepared, x) else - (@test _inferred_gradient(prepared, x)) + @test _inferred_gradient(prepared, x) end end @testset "vector jacobian" begin prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) if jacobian_broken - (@test_broken _inferred_jacobian(prepared, x)) + @test_broken _inferred_jacobian(prepared, x) else - (@test _inferred_jacobian(prepared, x)) + @test _inferred_jacobian(prepared, x) end end @testset "hessian" begin @@ -403,9 +412,9 @@ function AbstractPPL.run_testcases( adtype, QuadraticProblem(), zeros(3); check_dims=false, order=2 ) if hessian_broken - (@test_broken _inferred_hessian(prepared, x)) + @test_broken _inferred_hessian(prepared, x) else - (@test _inferred_hessian(prepared, x)) + @test _inferred_hessian(prepared, x) end end return nothing From 97788ad71dd5b60d78975fc626a7c76ee37ab680 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 16:53:20 +0100 Subject: [PATCH 13/21] Fix CI: format the ext, gate :allocations on Julia 1.10, add :context group * Format: match the CI JuliaFormatter v1.0.62 baseline (the local one I was using is on 2.x and disagrees on `return`-keyword placement). * `:allocations` group: Julia 1.10 heap-allocates `Fix2`/closure captures that 1.11+ elides. Mark `gradient_broken=VERSION < v"1.11"` (and `jacobian_broken=VERSION < v"1.11"` on FD) so min CI doesn't flag the older runtime as a regression. * New `:context` group: lifts the inline "context-lowered gradient" testset from the FD test into `AbstractPPLTestExt`. Verifies `prepare(adtype, f, x; context=(c,))` lowers the context out of the gradient. FD calls it in place of the inline testset; Mooncake adds it alongside its richer Mooncake-specific context testset (forward parity, vector arity rejection, empty input with context). Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 24 +++++++++++++++++++++--- src/AbstractPPL.jl | 5 +++-- test/ext/forwarddiff/main.jl | 26 +++++++++++--------------- test/ext/mooncake/main.jl | 10 +++++++++- 4 files changed, 44 insertions(+), 21 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 35eb0131..54345000 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -339,13 +339,13 @@ end # each AD entry in a tiny named function that returns `true` on success — that # value lets `@test` / `@test_broken` evaluate the call uniformly. function _inferred_gradient(prepared, x) - (@inferred AbstractPPL.value_and_gradient!!(prepared, x); true) + return (@inferred AbstractPPL.value_and_gradient!!(prepared, x); true) end function _inferred_jacobian(prepared, x) - (@inferred AbstractPPL.value_and_jacobian!!(prepared, x); true) + return (@inferred AbstractPPL.value_and_jacobian!!(prepared, x); true) end function _inferred_hessian(prepared, x) - (@inferred AbstractPPL.value_gradient_and_hessian!!(prepared, x); true) + return (@inferred AbstractPPL.value_gradient_and_hessian!!(prepared, x); true) end # Backends with known regressions (e.g. Mooncake's allocating @@ -420,4 +420,22 @@ function AbstractPPL.run_testcases( return nothing end +function AbstractPPL.run_testcases( + ::Val{:context}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 +) + # `prepare(adtype, f, x; context=(c,))` builds an evaluator that computes + # `f(x, context...)` with AD differentiating only `x`. + f(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 + x = [0.3] + c = 0.1 + @testset "scalar gradient with context" begin + prepared = prepare_fn(adtype, f, x; check_dims=false, context=(c,)) + @test prepared(x) ≈ f(x, c) atol = atol rtol = rtol + val, grad = AbstractPPL.value_and_gradient!!(prepared, x) + @test val ≈ f(x, c) atol = atol rtol = rtol + @test grad ≈ [-(x[1] - c)] atol = 1e-10 rtol = rtol + end + return nothing +end + end # module diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 9b1f44cd..3578c394 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -26,8 +26,9 @@ round-trips on vector-input scalar-output evaluators; `:namedtuple` for for repeated calls against a single prepared evaluator; `:allocations` and `:type_stability` for `@allocated == 0` and `@inferred` checks on the AD hot paths (both accept `gradient_broken`, `jacobian_broken`, and (`:type_stability` only) -`hessian_broken` kwargs for backends with known broken paths). Downstream -packages may add other keys. +`hessian_broken` kwargs for backends with known broken paths); `:context` for +the `prepare(adtype, f, x; context=(c,))` lowering on a scalar gradient. +Downstream packages may add other keys. """ function generate_testcases end diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 2527e399..815b4062 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -14,7 +14,15 @@ using Test run_testcases(Val(:hessian); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) run_testcases(Val(:edge); adtype=AutoForwardDiff()) - run_testcases(Val(:allocations); adtype=AutoForwardDiff()) + # Julia 1.10 heap-allocates some `Fix2`/closure captures that 1.11+ + # elides. Mark `:allocations` broken on min to flag the regression + # detection without failing the suite on the older runtime. + run_testcases( + Val(:allocations); + adtype=AutoForwardDiff(), + gradient_broken=VERSION < v"1.11", + jacobian_broken=VERSION < v"1.11", + ) run_testcases(Val(:type_stability); adtype=AutoForwardDiff()) end @@ -27,20 +35,8 @@ using Test ) end - # `run_testcases` doesn't exercise `context=`; this fills that gap. - @testset "context-lowered gradient" begin - raw_logdensity(x::AbstractVector{<:Real}, offset) = -0.5 * (x[1] - offset)^2 - - x = [0.3] - ad = AutoForwardDiff() - - lowered = prepare(ad, raw_logdensity, x; check_dims=false, context=(0.1,)) - - @test lowered(x) == raw_logdensity(x, 0.1) - - val, grad = value_and_gradient!!(lowered, x) - @test val ≈ raw_logdensity(x, 0.1) - @test grad ≈ [-(x[1] - 0.1)] atol = 1e-10 + @testset "AutoForwardDiff context" begin + run_testcases(Val(:context); adtype=AutoForwardDiff(), atol=1e-10, rtol=1e-10) end # `AutoForwardDiff(; tag=...)` exists for nested differentiation. Check the diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index f3d8f5e4..3691fe97 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -26,12 +26,20 @@ using Test # cotangent/Jacobian buffers each call, and the forward-mode # Jacobian return type infers as `Tuple{Any, Union{Array{T,3}, # Matrix}}`. Mark those known-broken; the other paths must hold. - run_testcases(Val(:allocations); adtype=adtype, jacobian_broken=true) + # Julia 1.10 also heap-allocates `Fix2`/closure captures the AD + # path uses, so scalar gradient is marked broken on min. + run_testcases( + Val(:allocations); + adtype=adtype, + gradient_broken=VERSION < v"1.11", + jacobian_broken=true, + ) run_testcases( Val(:type_stability); adtype=adtype, jacobian_broken=adtype isa AutoMooncakeForward, ) + run_testcases(Val(:context); adtype=adtype, atol=1e-6, rtol=1e-6) end end From 0bfe1e89dd69fc9135e0495ff0c23acf895ed0ca Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 17:14:13 +0100 Subject: [PATCH 14/21] Skip ForwardDiff.checktag so custom Tag sentinels work in hot paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The custom-tag path was only structurally tested (the tag flowed into the config type parameter), not exercised through an actual AD call. DynamicPPL's downstream tests caught the gap: `AutoForwardDiff(; tag=Tag{DynamicPPLTag,Float64}())` carries a sentinel tag whose first type parameter is *not* `typeof(target)`, so ForwardDiff's default `checktag` errors when the hot path calls `ForwardDiff.gradient!`. Pass `Val(false)` to skip `checktag` in all four hot-path calls (this is what DifferentiationInterface does). The tag is purely a label for the outer Dual scope; the config built at prep time already encodes the right tag, so the check is redundant and harmful in the custom-tag case. Strengthen the regression test to actually run `value_and_gradient!!` on a prep built with a custom sentinel tag and assert the gradient matches the analytic value — would have caught the original bug. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLForwardDiffExt.jl | 13 ++++++++++--- test/ext/forwarddiff/main.jl | 14 ++++++++++---- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/ext/AbstractPPLForwardDiffExt.jl b/ext/AbstractPPLForwardDiffExt.jl index 18f9c112..da8f7c5c 100644 --- a/ext/AbstractPPLForwardDiffExt.jl +++ b/ext/AbstractPPLForwardDiffExt.jl @@ -103,6 +103,12 @@ end # config captured at prep time. @inline _fd_call(x, e::VectorEvaluator) = e.f(x, e.context...) +# `Val(false)` on every hot-path call below skips `ForwardDiff.checktag`. A +# user-supplied `adtype.tag` (e.g. DynamicPPL's `DynamicPPLTag` sentinel for +# nested AD) has a tag-type parameter that does not equal `typeof(target)`, so +# the default check would error. The tag's role is only to label the outer +# Dual scope; the config we built at prep time already encodes the right tag. + @inline function AbstractPPL.value_and_gradient!!( p::Prepared{ <:AutoForwardDiff, @@ -121,7 +127,7 @@ end ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.gradient!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) ) return (DiffResults.value(p.cache.result), DiffResults.gradient(p.cache.result)) end @@ -138,6 +144,7 @@ end Base.Fix2(_fd_call, p.evaluator), x, p.cache.gradient_config, + Val(false), ) return ( DiffResults.value(p.cache.gradient_result), @@ -178,7 +185,7 @@ end ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.jacobian!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) ) return (DiffResults.value(p.cache.result), DiffResults.jacobian(p.cache.result)) end @@ -206,7 +213,7 @@ end ) Evaluators._check_ad_input(p.evaluator, x) ForwardDiff.hessian!( - p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config + p.cache.result, Base.Fix2(_fd_call, p.evaluator), x, p.cache.config, Val(false) ) return ( DiffResults.value(p.cache.result), diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 815b4062..86071539 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -39,13 +39,19 @@ using Test run_testcases(Val(:context); adtype=AutoForwardDiff(), atol=1e-10, rtol=1e-10) end - # `AutoForwardDiff(; tag=...)` exists for nested differentiation. Check the - # user-supplied tag is threaded into the ForwardDiff config (the inner - # `*Config` carries the tag in its first type parameter). + # `AutoForwardDiff(; tag=...)` exists for nested differentiation. The tag's + # type parameter is a sentinel chosen by the caller (e.g. DynamicPPL's + # `DynamicPPLTag`); it intentionally does not equal `typeof(target)`, so + # the hot path must skip `ForwardDiff.checktag` to avoid a false error. @testset "custom AutoForwardDiff tag" begin struct OuterTag end custom = ForwardDiff.Tag{OuterTag,Float64}() - prep = prepare(AutoForwardDiff(; tag=custom), x -> sum(abs2, x), [1.0, 2.0]) + x = [1.0, 2.0] + prep = prepare(AutoForwardDiff(; tag=custom), x -> sum(abs2, x), x) @test typeof(prep.cache.config).parameters[1] === typeof(custom) + # The actual AD call must succeed despite the sentinel tag. + val, grad = value_and_gradient!!(prep, x) + @test val ≈ 5.0 + @test grad ≈ [2.0, 4.0] end end From 268e5454bb86af3300db8e67084f14efe305c51d Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 17:17:57 +0100 Subject: [PATCH 15/21] Honor caller atol in the :context group's gradient assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gradient comparison was hardcoding `atol = 1e-10` while the value comparison above it (and every other group in this file) used `atol = atol`. The hardcoded value silently overrode the caller's kwarg — Mooncake calls with `atol = 1e-6` were getting the tighter 1e-10. Use `atol = atol` to match the surrounding pattern. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 54345000..a7ff36a5 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -433,7 +433,7 @@ function AbstractPPL.run_testcases( @test prepared(x) ≈ f(x, c) atol = atol rtol = rtol val, grad = AbstractPPL.value_and_gradient!!(prepared, x) @test val ≈ f(x, c) atol = atol rtol = rtol - @test grad ≈ [-(x[1] - c)] atol = 1e-10 rtol = rtol + @test grad ≈ [-(x[1] - c)] atol = atol rtol = rtol end return nothing end From c35ac4a1a9f87c375c328114d1d1506550648cb1 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 17:48:26 +0100 Subject: [PATCH 16/21] Unify the conformance harness into TestCase + run_testcase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collapse the four case structs (`ValueCase`, `HessianCase`, `ErrorCase`, `CacheReuseCase`) and seven `run_testcases(Val(:group); ...)` methods into a single tagged `TestCase` and a single `run_testcase(case; ...)` that dispatches on `case.tag` via `Val`. Each backend's test script is now a single uniform loop: ```julia for case in generate_testcases() run_testcase(case; adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6, allocations=:test, type_stability=:test) end ``` Tags are `:vector`, `:hessian`, `:context`, `:edge`, `:cache_reuse`, `:namedtuple`. NamedTuple-input cases live in `generate_namedtuple_testcases()` so backends that don't support that input shape don't need to filter at the call site. `allocations` / `type_stability` accept `:skip` / `:test` / `:broken` (`:broken` wraps as `@test_broken` for known regressions). Per-case `allocations_safe::Bool` defaults to `true`; cases with allocating primals (`VectorValuedProblem` result vector, empty-input shortcuts, hessian scratch, cache-reuse loops) opt out so the runner skips the alloc check regardless of caller intent. Case types and stubs (`TestCase`, `generate_testcases`, `generate_namedtuple_testcases`, `run_testcase`) live in `src/AbstractPPL.jl`; the generators and dispatched runner live in `ext/AbstractPPLTestExt.jl`. The old `Val{group}` API and the standalone `IdentityProblem` fixture are gone. Backend-specific broken predicates (`_mooncake_alloc`, `_mooncake_inferred`) sit next to the loop they drive — they encode Mooncake's known issues (allocating Jacobian, forward-mode Jacobian and context inference) without touching the shared harness. Local: FD 111/111, Mooncake 149 pass + 3 broken, DI 115/115. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 653 +++++++++++----------- src/AbstractPPL.jl | 116 +++- test/ext/differentiationinterface/main.jl | 24 +- test/ext/forwarddiff/main.jl | 49 +- test/ext/mooncake/main.jl | 77 ++- 5 files changed, 506 insertions(+), 413 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index a7ff36a5..ea48787e 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -1,6 +1,7 @@ module AbstractPPLTestExt -using AbstractPPL: AbstractPPL, generate_testcases, run_testcases +using AbstractPPL: + AbstractPPL, generate_testcases, generate_namedtuple_testcases, run_testcase, TestCase using Test: @inferred, @test, @test_broken, @test_throws, @testset struct QuadraticProblem end @@ -9,433 +10,429 @@ struct QuadraticProblem end struct VectorValuedProblem end (::VectorValuedProblem)(x::AbstractVector{<:Real}) = [x[1] * x[2], x[2] + x[3]] -# Allocation-free vector-output problem for the `:allocations` group: -# `VectorValuedProblem` allocates its result vector, masking AD-path allocations. -struct IdentityProblem end -(::IdentityProblem)(x::AbstractVector{<:Real}) = x +_context_problem(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 -struct ValueCase - name::String - f::Any - x_proto::Any - x::Any - value::Any - gradient::Any - jacobian::Any -end - -struct HessianCase - name::String - f::Any - x_proto::Any - x::Any - value::Any - gradient::Any - hessian::Any -end - -struct ErrorCase - name::String - f::Any - x_proto::Any - x::Any - op::Any - exception::Any -end - -function AbstractPPL.generate_testcases(::Val{:vector}) +function AbstractPPL.generate_testcases() return ( - ValueCase( + TestCase( "quadratic (scalar output)", + :vector, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0], - 14.0, - [6.0, 2.0, 4.0], - nothing, + zeros(3); + x=[3.0, 1.0, 2.0], + value=14.0, + gradient=[6.0, 2.0, 4.0], ), - ValueCase( + TestCase( "vector-valued (vector output)", + :vector, 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], + zeros(3); + x=[2.0, 3.0, 4.0], + value=[6.0, 7.0], + jacobian=[3.0 2.0 0.0; 0.0 1.0 1.0], + allocations_safe=false, # primal allocates its result vector ), - ValueCase( + TestCase( "empty input, scalar output", + :vector, x -> 7.5, - Float64[], - Float64[], - 7.5, - Float64[], - nothing, + Float64[]; + x=Float64[], + value=7.5, + gradient=Float64[], + allocations_safe=false, # empty-input shortcut returns fresh `T[]` ), - ValueCase( + TestCase( "empty input, vector output", + :vector, x -> [2.0, 3.0], - Float64[], - Float64[], - [2.0, 3.0], - nothing, - zeros(2, 0), + Float64[]; + x=Float64[], + value=[2.0, 3.0], + jacobian=zeros(2, 0), + allocations_safe=false, # empty-input shortcut allocates empty matrix ), - ) -end - -function AbstractPPL.generate_testcases(::Val{:hessian}) - return ( - HessianCase( - "quadratic (scalar output)", + TestCase( + "scalar gradient with context", + :context, + _context_problem, + [0.3]; + x=[0.3], + value=_context_problem([0.3], 0.1), + gradient=[-(0.3 - 0.1)], + context=(0.1,), + ), + TestCase( + "quadratic (hessian)", + :hessian, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0], - 14.0, - [6.0, 2.0, 4.0], - [2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 2.0], + zeros(3); + x=[3.0, 1.0, 2.0], + value=14.0, + gradient=[6.0, 2.0, 4.0], + hessian=[2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 2.0], + allocations_safe=false, # ForwardDiff/Mooncake hessian path allocates scratch ), - HessianCase( - "empty input, scalar output", + TestCase( + "empty input, hessian", + :hessian, x -> 7.5, - Float64[], - Float64[], - 7.5, - Float64[], - zeros(0, 0), + Float64[]; + x=Float64[], + value=7.5, + gradient=Float64[], + hessian=zeros(0, 0), + allocations_safe=false, ), - ) -end - -function AbstractPPL.generate_testcases(::Val{:hessian_edge}) - return ( - # `value_gradient_and_hessian!!` rejects order=1 preps regardless of - # the underlying problem arity — both paths share the same dispatch - # so one case suffices. - ErrorCase( + # value_gradient_and_hessian!! rejects order=1 preps regardless of arity; + # both paths share the dispatch so one case suffices. + TestCase( "value_gradient_and_hessian!! on order=1 prep", + :edge, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0], - (prepared, x) -> AbstractPPL.value_gradient_and_hessian!!(prepared, x), - r"order=2", + zeros(3); + x=[3.0, 1.0, 2.0], + op=(prepared, x) -> AbstractPPL.value_gradient_and_hessian!!(prepared, x), + exception=r"order=2", ), - ) -end - -function AbstractPPL.generate_testcases(::Val{:edge}) - return ( - ErrorCase( + TestCase( "wrong vector length", + :edge, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0, 99.0], - (prepared, x) -> prepared(x), - DimensionMismatch, + zeros(3); + x=[3.0, 1.0, 2.0, 99.0], + op=(prepared, x) -> prepared(x), + exception=DimensionMismatch, ), - ErrorCase( + TestCase( "non-floating-point vector", + :edge, QuadraticProblem(), - zeros(3), - [3, 1, 2], - (prepared, x) -> prepared(x), - r"floating-point", + zeros(3); + x=[3, 1, 2], + op=(prepared, x) -> prepared(x), + exception=r"floating-point", ), - ErrorCase( + TestCase( "gradient of vector-valued output", + :edge, VectorValuedProblem(), - zeros(3), - [2.0, 3.0, 4.0], - (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), - r"scalar-valued", + zeros(3); + x=[2.0, 3.0, 4.0], + op=(prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + exception=r"scalar-valued", ), - ErrorCase( + TestCase( "jacobian of scalar output", + :edge, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0], - (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), - r"vector-valued", + zeros(3); + x=[3.0, 1.0, 2.0], + op=(prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + exception=r"vector-valued", ), - ErrorCase( + TestCase( "gradient of vector-valued output, empty input", + :edge, x -> [2.0, 3.0], - Float64[], - Float64[], - (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), - r"scalar-valued", + Float64[]; + x=Float64[], + op=(prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + exception=r"scalar-valued", ), - ErrorCase( + TestCase( "jacobian of scalar output, empty input", + :edge, x -> 7.5, - Float64[], - Float64[], - (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), - r"vector-valued", + Float64[]; + x=Float64[], + op=(prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + exception=r"vector-valued", ), - ErrorCase( + TestCase( "value_and_gradient!! wrong vector length", + :edge, QuadraticProblem(), - zeros(3), - [3.0, 1.0, 2.0, 99.0], - (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), - DimensionMismatch, + zeros(3); + x=[3.0, 1.0, 2.0, 99.0], + op=(prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + exception=DimensionMismatch, ), - ErrorCase( + TestCase( "value_and_jacobian!! wrong vector length", + :edge, VectorValuedProblem(), - zeros(3), - [2.0, 3.0, 4.0, 5.0], - (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), - DimensionMismatch, + zeros(3); + x=[2.0, 3.0, 4.0, 5.0], + op=(prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + exception=DimensionMismatch, ), - ErrorCase( + TestCase( "value_and_gradient!! non-floating-point vector", + :edge, QuadraticProblem(), - zeros(3), - [3, 1, 2], - (prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), - r"floating-point", + zeros(3); + x=[3, 1, 2], + op=(prepared, x) -> AbstractPPL.value_and_gradient!!(prepared, x), + exception=r"floating-point", ), - ErrorCase( + TestCase( "value_and_jacobian!! non-floating-point vector", + :edge, VectorValuedProblem(), - zeros(3), - [2, 3, 4], - (prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), - r"floating-point", + zeros(3); + x=[2, 3, 4], + op=(prepared, x) -> AbstractPPL.value_and_jacobian!!(prepared, x), + exception=r"floating-point", + ), + TestCase( + "scalar output, cache reuse", + :cache_reuse, + QuadraticProblem(), + zeros(3); + inputs=[ + (x=[1.0, 2.0, 3.0], value=14.0, gradient=[2.0, 4.0, 6.0]), + (x=[4.0, 5.0, 6.0], value=77.0, gradient=[8.0, 10.0, 12.0]), + (x=[0.5, -1.0, 2.0], value=5.25, gradient=[1.0, -2.0, 4.0]), + ], + allocations_safe=false, + ), + TestCase( + "vector output, cache reuse", + :cache_reuse, + VectorValuedProblem(), + zeros(3); + inputs=[ + (x=[2.0, 3.0, 4.0], value=[6.0, 7.0], jacobian=[3.0 2.0 0.0; 0.0 1.0 1.0]), + (x=[5.0, 1.0, 7.0], value=[5.0, 8.0], jacobian=[1.0 5.0 0.0; 0.0 1.0 1.0]), + (x=[0.0, 4.0, -2.0], value=[0.0, 2.0], jacobian=[4.0 0.0 0.0; 0.0 1.0 1.0]), + ], + allocations_safe=false, ), ) end -function AbstractPPL.generate_testcases(::Val{:namedtuple}) +function AbstractPPL.generate_namedtuple_testcases() return ( - ValueCase( + TestCase( "scalar output over (x::Real, y::Vector)", + :namedtuple, 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, + (x=0.0, y=zeros(2)); + x=(x=3.0, y=[1.0, 2.0]), + value=14.0, + gradient=(x=6.0, y=[2.0, 4.0]), ), ) 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 AbstractPPL.order(prepared) == 1 - @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 +# ----- helpers ----- + +# NamedTuple gradients compare per-key (some backends return Mooncake-tagged +# tangents that aren't directly `≈`-comparable as a whole). +function _compare_derivative(actual::NamedTuple, expected::NamedTuple; atol, rtol) + for k in keys(expected) + @test getproperty(actual, k) ≈ getproperty(expected, k) atol = atol rtol = rtol end - return nothing +end +function _compare_derivative(actual, expected; atol, rtol) + @test actual ≈ expected atol = atol rtol = rtol end -function AbstractPPL.run_testcases( - ::Val{:hessian}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 -) - for case in generate_testcases(Val(:hessian)) - @testset "$(case.name)" begin - prepared = prepare_fn(adtype, case.f, case.x_proto; order=2) - @test AbstractPPL.order(prepared) == 2 - @test prepared(case.x) ≈ case.value atol = atol rtol = rtol - val, grad, hess = AbstractPPL.value_gradient_and_hessian!!(prepared, case.x) - @test val ≈ case.value atol = atol rtol = rtol - @test grad ≈ case.gradient atol = atol rtol = rtol - @test hess ≈ case.hessian atol = atol rtol = rtol - # Order=2 prep also satisfies the order=1 gradient contract. - val1, grad1 = AbstractPPL.value_and_gradient!!(prepared, case.x) - @test val1 ≈ case.value atol = atol rtol = rtol - @test grad1 ≈ case.gradient atol = atol rtol = rtol - end - end - for case in generate_testcases(Val(:hessian_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 +function _record_alloc!(state::Symbol, allocs::Integer) + state === :test && @test allocs == 0 + state === :broken && @test_broken allocs == 0 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 +# `@inferred` is syntactic and throws on failure; wrap so we can pin `op`'s +# type via an F-parameter and convert the throw into a Bool. +function _is_inferred(op::F, args...) where {F} + try + @inferred op(args...) + return true + catch + return false end +end + +function _record_inferred!(state::Symbol, inferred::Bool) + state === :test && @test inferred + state === :broken && @test_broken inferred return nothing end -function AbstractPPL.run_testcases( - ::Val{:namedtuple}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 -) - for case in generate_testcases(Val(:namedtuple)) - @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 - for k in keys(case.gradient) - @test getproperty(grad, k) ≈ getproperty(case.gradient, k) atol = atol rtol = - rtol - end - end - end +# ----- runner ----- + +function AbstractPPL.run_testcase(case::TestCase; kwargs...) + @testset "$(case.name)" begin + _run(Val(case.tag), case; kwargs...) end return nothing end -# Drive `value_and_{gradient,jacobian}!!` twice with different inputs against -# the same `prepared` evaluator to exercise cache reuse — catches backends -# whose cache state is corrupted by a prior call. -function AbstractPPL.run_testcases( - ::Val{:cache_reuse}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 +function _run( + ::Val{:vector}, + case; + adtype, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + check_dims::Bool=true, + type_stability::Symbol=:skip, + allocations::Symbol=:skip, ) - @testset "scalar output, repeated calls" begin - prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3)) - for (x, value, gradient) in ( - ([1.0, 2.0, 3.0], 14.0, [2.0, 4.0, 6.0]), - ([4.0, 5.0, 6.0], 77.0, [8.0, 10.0, 12.0]), - ([0.5, -1.0, 2.0], 5.25, [1.0, -2.0, 4.0]), + prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims) + @test AbstractPPL.order(prepared) == 1 + @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 + _compare_derivative(grad, case.gradient; atol, rtol) + _maybe_check_alloc!( + case, allocations, AbstractPPL.value_and_gradient!!, prepared, case.x + ) + _maybe_check_inferred!( + type_stability, AbstractPPL.value_and_gradient!!, prepared, case.x ) - val, grad = AbstractPPL.value_and_gradient!!(prepared, x) - @test val ≈ value atol = atol rtol = rtol - @test grad ≈ gradient atol = atol rtol = rtol - end end - @testset "vector output, repeated calls" begin - prepared = prepare_fn(adtype, VectorValuedProblem(), zeros(3)) - for (x, value, jacobian) in ( - ([2.0, 3.0, 4.0], [6.0, 7.0], [3.0 2.0 0.0; 0.0 1.0 1.0]), - ([5.0, 1.0, 7.0], [5.0, 8.0], [1.0 5.0 0.0; 0.0 1.0 1.0]), - ([0.0, 4.0, -2.0], [0.0, 2.0], [4.0 0.0 0.0; 0.0 1.0 1.0]), + + 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 + _maybe_check_alloc!( + case, allocations, AbstractPPL.value_and_jacobian!!, prepared, case.x + ) + _maybe_check_inferred!( + type_stability, AbstractPPL.value_and_jacobian!!, prepared, case.x ) - val, jac = AbstractPPL.value_and_jacobian!!(prepared, x) - @test val ≈ value atol = atol rtol = rtol - @test jac ≈ jacobian atol = atol rtol = rtol - end end return nothing end -# Helpers for the `:type_stability` group: `@inferred` is a syntactic macro, so wrap -# each AD entry in a tiny named function that returns `true` on success — that -# value lets `@test` / `@test_broken` evaluate the call uniformly. -function _inferred_gradient(prepared, x) - return (@inferred AbstractPPL.value_and_gradient!!(prepared, x); true) -end -function _inferred_jacobian(prepared, x) - return (@inferred AbstractPPL.value_and_jacobian!!(prepared, x); true) -end -function _inferred_hessian(prepared, x) - return (@inferred AbstractPPL.value_gradient_and_hessian!!(prepared, x); true) +function _run( + ::Val{:context}, + case; + adtype, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + check_dims::Bool=true, + type_stability::Symbol=:skip, + allocations::Symbol=:skip, +) + prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims, context=case.context) + @test AbstractPPL.order(prepared) == 1 + @test prepared(case.x) ≈ case.value atol = atol rtol = rtol + 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 + _maybe_check_alloc!( + case, allocations, AbstractPPL.value_and_gradient!!, prepared, case.x + ) + _maybe_check_inferred!( + type_stability, AbstractPPL.value_and_gradient!!, prepared, case.x + ) + return nothing end -# Backends with known regressions (e.g. Mooncake's allocating -# `value_and_jacobian!!`, or its forward-mode Jacobian inference) pass -# `*_broken=true` to mark the assertion as broken instead of failing. -function AbstractPPL.run_testcases( - ::Val{:allocations}, - prepare_fn=AbstractPPL.prepare; +function _run( + ::Val{:hessian}, + case; adtype, - gradient_broken::Bool=false, - jacobian_broken::Bool=false, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + check_dims::Bool=true, + type_stability::Symbol=:skip, + allocations::Symbol=:skip, ) - x = [1.0, 2.0, 3.0] - @testset "scalar gradient" begin - prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) - AbstractPPL.value_and_gradient!!(prepared, x) # warm up - allocs = @allocated AbstractPPL.value_and_gradient!!(prepared, x) - if gradient_broken - @test_broken allocs == 0 - else - @test allocs == 0 - end - end - @testset "vector jacobian" begin - prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) - AbstractPPL.value_and_jacobian!!(prepared, x) - allocs = @allocated AbstractPPL.value_and_jacobian!!(prepared, x) - if jacobian_broken - @test_broken allocs == 0 - else - @test allocs == 0 - end - end + prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims, order=2) + @test AbstractPPL.order(prepared) == 2 + @test prepared(case.x) ≈ case.value atol = atol rtol = rtol + + val, grad, hess = AbstractPPL.value_gradient_and_hessian!!(prepared, case.x) + @test val ≈ case.value atol = atol rtol = rtol + @test grad ≈ case.gradient atol = atol rtol = rtol + @test hess ≈ case.hessian atol = atol rtol = rtol + + # Order=2 prep also satisfies the order=1 gradient contract. + val1, grad1 = AbstractPPL.value_and_gradient!!(prepared, case.x) + @test val1 ≈ case.value atol = atol rtol = rtol + @test grad1 ≈ case.gradient atol = atol rtol = rtol + + _maybe_check_alloc!( + case, allocations, AbstractPPL.value_gradient_and_hessian!!, prepared, case.x + ) + _maybe_check_inferred!( + type_stability, AbstractPPL.value_gradient_and_hessian!!, prepared, case.x + ) + return nothing +end + +function _run(::Val{:edge}, case; adtype, prepare_fn=AbstractPPL.prepare, kwargs...) + prepared = prepare_fn(adtype, case.f, case.x_proto) + @test_throws case.exception case.op(prepared, case.x) return nothing end -function AbstractPPL.run_testcases( - ::Val{:type_stability}, - prepare_fn=AbstractPPL.prepare; +function _run( + ::Val{:cache_reuse}, + case; adtype, - gradient_broken::Bool=false, - jacobian_broken::Bool=false, - hessian_broken::Bool=false, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + kwargs..., ) - x = [1.0, 2.0, 3.0] - @testset "scalar gradient" begin - prepared = prepare_fn(adtype, QuadraticProblem(), zeros(3); check_dims=false) - if gradient_broken - @test_broken _inferred_gradient(prepared, x) - else - @test _inferred_gradient(prepared, x) - end - end - @testset "vector jacobian" begin - prepared = prepare_fn(adtype, IdentityProblem(), zeros(3); check_dims=false) - if jacobian_broken - @test_broken _inferred_jacobian(prepared, x) - else - @test _inferred_jacobian(prepared, x) - end - end - @testset "hessian" begin - prepared = prepare_fn( - adtype, QuadraticProblem(), zeros(3); check_dims=false, order=2 - ) - if hessian_broken - @test_broken _inferred_hessian(prepared, x) + prepared = prepare_fn(adtype, case.f, case.x_proto) + for input in case.inputs + if haskey(input, :gradient) + val, grad = AbstractPPL.value_and_gradient!!(prepared, input.x) + @test val ≈ input.value atol = atol rtol = rtol + @test grad ≈ input.gradient atol = atol rtol = rtol else - @test _inferred_hessian(prepared, x) + val, jac = AbstractPPL.value_and_jacobian!!(prepared, input.x) + @test val ≈ input.value atol = atol rtol = rtol + @test jac ≈ input.jacobian atol = atol rtol = rtol end end return nothing end -function AbstractPPL.run_testcases( - ::Val{:context}, prepare_fn=AbstractPPL.prepare; adtype, atol=0, rtol=1e-10 +function _run( + ::Val{:namedtuple}, + case; + adtype, + prepare_fn=AbstractPPL.prepare, + atol=0, + rtol=1e-10, + kwargs..., ) - # `prepare(adtype, f, x; context=(c,))` builds an evaluator that computes - # `f(x, context...)` with AD differentiating only `x`. - f(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 - x = [0.3] - c = 0.1 - @testset "scalar gradient with context" begin - prepared = prepare_fn(adtype, f, x; check_dims=false, context=(c,)) - @test prepared(x) ≈ f(x, c) atol = atol rtol = rtol - val, grad = AbstractPPL.value_and_gradient!!(prepared, x) - @test val ≈ f(x, c) atol = atol rtol = rtol - @test grad ≈ [-(x[1] - c)] atol = atol rtol = rtol - end + prepared = prepare_fn(adtype, case.f, case.x_proto) + @test prepared(case.x) ≈ case.value atol = atol rtol = rtol + val, grad = AbstractPPL.value_and_gradient!!(prepared, case.x) + @test val ≈ case.value atol = atol rtol = rtol + _compare_derivative(grad, case.gradient; atol, rtol) return nothing end +_resolve_alloc_state(case::TestCase, state::Symbol) = case.allocations_safe ? state : :skip + +function _maybe_check_alloc!(case::TestCase, state::Symbol, op::F, prepared, x) where {F} + effective = _resolve_alloc_state(case, state) + effective === :skip && return nothing + op(prepared, x) # warm up + allocs = @allocated op(prepared, x) + return _record_alloc!(effective, allocs) +end + +function _maybe_check_inferred!(state::Symbol, op::F, prepared, x) where {F} + state === :skip && return nothing + return _record_inferred!(state, _is_inferred(op, prepared, x)) +end + end # module diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 3578c394..4aad5891 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -15,37 +15,111 @@ using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!!, value_gradient_and_hessian!!, order """ - 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; `:hessian` for `order=2` value/gradient/Hessian -round-trips on vector-input scalar-output evaluators; `:namedtuple` for -`NamedTuple`-input evaluators; `:edge` for error-path cases; `:cache_reuse` -for repeated calls against a single prepared evaluator; `:allocations` and -`:type_stability` for `@allocated == 0` and `@inferred` checks on the AD hot paths -(both accept `gradient_broken`, `jacobian_broken`, and (`:type_stability` only) -`hessian_broken` kwargs for backends with known broken paths); `:context` for -the `prepare(adtype, f, x; context=(c,))` lowering on a scalar gradient. -Downstream packages may add other keys. + TestCase(name, tag, f, x_proto; x, value, gradient, jacobian, hessian, + context=(), op, exception, inputs) + +Single tagged case for AD conformance testing. The `tag::Symbol` selects how +the case is run; the kwargs populate only the fields the tag uses. + +Reserved tags (recognised by [`run_testcase`](@ref)): + + - `:vector` — vector input, scalar output (`gradient`) or vector output + (`jacobian`). + - `:hessian` — order=2 round-trip on scalar output. + - `:context` — scalar-output gradient with a non-empty `context::Tuple` + passed to `prepare`. + - `:edge` — error case; `op(prepared, x)` must throw `exception`. + - `:cache_reuse` — multiple inputs against a single prepared evaluator + (`inputs::Vector{<:NamedTuple}`, with `(x=, value=, + gradient=)` or `(x=, value=, jacobian=)` per row). + - `:namedtuple` — NamedTuple input and gradient; Mooncake-only. +""" +struct TestCase + name::String + tag::Symbol + f::Any + x_proto::Any + x::Any + value::Any + gradient::Any + jacobian::Any + hessian::Any + context::Tuple + op::Any + exception::Any + inputs::Any + # Cases with an allocating primal (vector-output result vectors, the + # empty-input shortcut's `T[]`) or shapes the original `:allocations` group + # never covered (hessian, cache-reuse, edge) set this to `false` — the + # runner then skips the `allocations=` check regardless of caller intent. + allocations_safe::Bool +end +function TestCase( + name, + tag::Symbol, + f, + x_proto; + x=nothing, + value=nothing, + gradient=nothing, + jacobian=nothing, + hessian=nothing, + context::Tuple=(), + op=nothing, + exception=nothing, + inputs=nothing, + allocations_safe::Bool=true, +) + return TestCase( + name, + tag, + f, + x_proto, + x, + value, + gradient, + jacobian, + hessian, + context, + op, + exception, + inputs, + allocations_safe, + ) +end + +""" + generate_testcases() + +Return a tuple of conformance [`TestCase`](@ref)s for vector-input AD +backends. Iterate and pass each to [`run_testcase`](@ref). """ function generate_testcases end """ - run_testcases(::Val{group}, prepare_fn=AbstractPPL.prepare; adtype, kwargs...) + generate_namedtuple_testcases() + +Like [`generate_testcases`](@ref) but for evaluators with `NamedTuple` input. +""" +function generate_namedtuple_testcases end + +""" + run_testcase(case; adtype, prepare_fn=AbstractPPL.prepare, atol=0, rtol=1e-10, + check_dims=true, type_stability=:skip, allocations=:skip) -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. +Run a single [`TestCase`](@ref) against an AD backend. `type_stability` and +`allocations` accept `:skip` / `:test` / `:broken` — `:test` asserts the +invariant, `:broken` marks it `@test_broken` (use for backends with known +regressions). Implemented by the `Test` extension. """ -function run_testcases end +function run_testcase end @static if VERSION >= v"1.11.0" eval( Meta.parse( - "public prepare, value_and_gradient!!, value_and_jacobian!!, value_gradient_and_hessian!!, order, generate_testcases, run_testcases", + "public prepare, value_and_gradient!!, value_and_jacobian!!, " * + "value_gradient_and_hessian!!, order, " * + "generate_testcases, generate_namedtuple_testcases, run_testcase, TestCase", ), ) end diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index 53a4fba6..e9a5cf05 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -6,7 +6,8 @@ Pkg.instantiate() using AbstractPPL: AbstractPPL, prepare, - run_testcases, + generate_testcases, + run_testcase, value_and_gradient!!, value_gradient_and_hessian!!, order @@ -22,20 +23,21 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) @testset "AbstractPPLDifferentiationInterfaceExt" begin @testset "ForwardDiff" begin - run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:hessian); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:edge); adtype=AutoForwardDiff()) + for case in generate_testcases() + run_testcase(case; adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) + end end - # Compiled-tape ReverseDiff goes through the `_di_call_shape(::AutoReverseDiff{true}, …)` - # specialisation that closes the evaluator into a `Base.Fix2` target — the - # `:cache_reuse` group exercises that path across multiple inputs. + # Compiled-tape ReverseDiff closes the evaluator into a `Base.Fix2` target + # via `_di_call_shape(::AutoReverseDiff{true}, …)`; the `:cache_reuse` + # cases exercise that path across multiple inputs. Skip `:hessian` + # (compiled tape doesn't support `prepare_hessian`). @testset "ReverseDiff (compiled tape)" begin adtype = AutoReverseDiff(; compile=true) - run_testcases(Val(:vector); adtype=adtype, atol=1e-6, rtol=1e-6) - run_testcases(Val(:cache_reuse); adtype=adtype, atol=1e-6, rtol=1e-6) - run_testcases(Val(:edge); adtype=adtype) + for case in generate_testcases() + case.tag === :hessian && continue + run_testcase(case; adtype, atol=1e-6, rtol=1e-6) + end end # The DI cache types' `Mode` parameter is either `:closure` (compiled-tape diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 86071539..435fa176 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -3,40 +3,38 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: AbstractPPL, prepare, run_testcases, value_and_gradient!! +using AbstractPPL: + AbstractPPL, prepare, generate_testcases, run_testcase, value_and_gradient!! using ADTypes: AutoForwardDiff using ForwardDiff using Test @testset "AbstractPPLForwardDiffExt" begin + # Julia 1.10 heap-allocates closure captures the 1.11+ runtime elides; mark + # allocations broken on min so the regression check stays honest on latest. + alloc_state = VERSION < v"1.11" ? :broken : :test + @testset "ForwardDiff (default chunk)" begin - run_testcases(Val(:vector); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:hessian); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:cache_reuse); adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) - run_testcases(Val(:edge); adtype=AutoForwardDiff()) - # Julia 1.10 heap-allocates some `Fix2`/closure captures that 1.11+ - # elides. Mark `:allocations` broken on min to flag the regression - # detection without failing the suite on the older runtime. - run_testcases( - Val(:allocations); - adtype=AutoForwardDiff(), - gradient_broken=VERSION < v"1.11", - jacobian_broken=VERSION < v"1.11", - ) - run_testcases(Val(:type_stability); adtype=AutoForwardDiff()) + for case in generate_testcases() + run_testcase( + case; + adtype=AutoForwardDiff(), + atol=1e-6, + rtol=1e-6, + allocations=alloc_state, + type_stability=:test, + ) + end end + # `chunksize=2` needs x with at least two elements; skip the `:context` + # case (x of length 1) and `:edge` cases (chunk doesn't apply). @testset "ForwardDiff (explicit chunk)" begin - run_testcases( - Val(:vector); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6 - ) - run_testcases( - Val(:cache_reuse); adtype=AutoForwardDiff(; chunksize=2), atol=1e-6, rtol=1e-6 - ) - end - - @testset "AutoForwardDiff context" begin - run_testcases(Val(:context); adtype=AutoForwardDiff(), atol=1e-10, rtol=1e-10) + ad = AutoForwardDiff(; chunksize=2) + for case in generate_testcases() + case.tag ∈ (:vector, :cache_reuse, :hessian) || continue + run_testcase(case; adtype=ad, atol=1e-6, rtol=1e-6) + end end # `AutoForwardDiff(; tag=...)` exists for nested differentiation. The tag's @@ -49,7 +47,6 @@ using Test x = [1.0, 2.0] prep = prepare(AutoForwardDiff(; tag=custom), x -> sum(abs2, x), x) @test typeof(prep.cache.config).parameters[1] === typeof(custom) - # The actual AD call must succeed despite the sentinel tag. val, grad = value_and_gradient!!(prep, x) @test val ≈ 5.0 @test grad ≈ [2.0, 4.0] diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index 3691fe97..eda4e254 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -3,43 +3,66 @@ Pkg.activate(@__DIR__) Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() -using AbstractPPL: AbstractPPL, prepare, run_testcases, value_and_gradient!! +using AbstractPPL: + AbstractPPL, + prepare, + generate_testcases, + generate_namedtuple_testcases, + run_testcase, + value_and_gradient!! using ADTypes: AutoMooncake, AutoMooncakeForward using Mooncake using Test +# Known-broken paths in Mooncake: +# * `value_and_jacobian!!` allocates fresh cotangent/Jacobian buffers on +# every call (both modes); forward-mode Jacobian return type infers as +# `Tuple{Any, Union{Array{T,3}, Matrix}}`. +# * `value_and_gradient!!` on a context-lowered prep splats `args_to_zero` +# per call (reverse mode allocates; forward mode also fails inference). +# Julia 1.10 also heap-allocates `Fix2`/closure captures that 1.11+ elides. +function _mooncake_alloc(case, adtype) + if case.tag === :vector && case.jacobian !== nothing + return :broken + elseif case.tag === :context && adtype isa AutoMooncakeForward + return :broken + elseif VERSION < v"1.11" + return :broken + else + return :test + end +end +# The forward-mode Jacobian inference issue only affects non-empty input; +# the empty-input shortcut bypasses Mooncake and is inferable on either mode. +function _mooncake_inferred(case, adtype) + is_jac_inf_broken = + case.tag === :vector && + case.jacobian !== nothing && + length(case.x) > 0 && + adtype isa AutoMooncakeForward + is_ctx_inf_broken = case.tag === :context && adtype isa AutoMooncakeForward + return (is_jac_inf_broken || is_ctx_inf_broken) ? :broken : :test +end + @testset "AbstractPPLMooncakeExt" begin for (label, adtype) in ( ("Mooncake (reverse)", AutoMooncake()), ("Mooncake (forward)", AutoMooncakeForward()), ) @testset "$label" begin - run_testcases(Val(:vector); adtype=adtype, atol=1e-6, rtol=1e-6) - run_testcases(Val(:namedtuple); adtype=adtype, atol=1e-6, rtol=1e-6) - run_testcases(Val(:cache_reuse); adtype=adtype, atol=1e-6, rtol=1e-6) - run_testcases(Val(:edge); adtype=adtype) - # Hessian (`order=2`) is reverse-mode only on the AutoMooncake side; - # AutoMooncakeForward routes through the same generic Hessian path - # since `Mooncake.prepare_hessian_cache` is mode-agnostic. - run_testcases(Val(:hessian); adtype=adtype, atol=1e-6, rtol=1e-6) - # Mooncake's `value_and_jacobian!!` currently allocates fresh - # cotangent/Jacobian buffers each call, and the forward-mode - # Jacobian return type infers as `Tuple{Any, Union{Array{T,3}, - # Matrix}}`. Mark those known-broken; the other paths must hold. - # Julia 1.10 also heap-allocates `Fix2`/closure captures the AD - # path uses, so scalar gradient is marked broken on min. - run_testcases( - Val(:allocations); - adtype=adtype, - gradient_broken=VERSION < v"1.11", - jacobian_broken=true, - ) - run_testcases( - Val(:type_stability); - adtype=adtype, - jacobian_broken=adtype isa AutoMooncakeForward, - ) - run_testcases(Val(:context); adtype=adtype, atol=1e-6, rtol=1e-6) + for case in generate_testcases() + run_testcase( + case; + adtype, + atol=1e-6, + rtol=1e-6, + allocations=_mooncake_alloc(case, adtype), + type_stability=_mooncake_inferred(case, adtype), + ) + end + for case in generate_namedtuple_testcases() + run_testcase(case; adtype, atol=1e-6, rtol=1e-6) + end end end From 4246f6a12594693d91bc7ee0863930326549c5bd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 17:54:14 +0100 Subject: [PATCH 17/21] Move TestCase to the Test extension; one Val-dispatched generator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `src/AbstractPPL.jl` keeps only the two function stubs `generate_testcases` and `run_testcase`. The `TestCase` struct (and its keyword-arg constructor) moves into `ext/AbstractPPLTestExt.jl` — test scripts only access `case.tag` via field access, so the type itself doesn't need to live in main. Collapse the two separate generators into a single Val-dispatched function: generate_testcases(Val(:vector)) — all vector-input cases generate_testcases(Val(:namedtuple)) — NamedTuple-input cases Backends iterate `generate_testcases(Val(:vector))` (and Mooncake also `Val(:namedtuple)`). Local: FD 111/111, Mooncake 149 pass + 3 broken, DI 115/115 — no behavioural change. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 80 ++++++++++++++++++- src/AbstractPPL.jl | 93 ++--------------------- test/ext/differentiationinterface/main.jl | 4 +- test/ext/forwarddiff/main.jl | 4 +- test/ext/mooncake/main.jl | 11 +-- 5 files changed, 90 insertions(+), 102 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index ea48787e..430fb7fc 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -1,9 +1,81 @@ module AbstractPPLTestExt -using AbstractPPL: - AbstractPPL, generate_testcases, generate_namedtuple_testcases, run_testcase, TestCase +using AbstractPPL: AbstractPPL, generate_testcases, run_testcase using Test: @inferred, @test, @test_broken, @test_throws, @testset +""" + TestCase(name, tag, f, x_proto; x, value, gradient, jacobian, hessian, + context=(), op, exception, inputs, allocations_safe=true) + +Single tagged case for AD conformance testing. The `tag::Symbol` selects how +the case is run; the kwargs populate only the fields the tag uses. + +Reserved tags (recognised by [`run_testcase`](@ref)): + + - `:vector` — vector input, scalar output (`gradient`) or vector output + (`jacobian`). + - `:hessian` — order=2 round-trip on scalar output. + - `:context` — scalar-output gradient with a non-empty `context::Tuple` + passed to `prepare`. + - `:edge` — error case; `op(prepared, x)` must throw `exception`. + - `:cache_reuse` — multiple inputs against a single prepared evaluator + (`inputs::Vector{<:NamedTuple}`, with `(x=, value=, + gradient=)` or `(x=, value=, jacobian=)` per row). + - `:namedtuple` — NamedTuple input and gradient; Mooncake-only. + +`allocations_safe=false` opts the case out of the alloc check +(cases with an allocating primal or empty-input shortcuts that allocate). +""" +struct TestCase + name::String + tag::Symbol + f::Any + x_proto::Any + x::Any + value::Any + gradient::Any + jacobian::Any + hessian::Any + context::Tuple + op::Any + exception::Any + inputs::Any + allocations_safe::Bool +end +function TestCase( + name, + tag::Symbol, + f, + x_proto; + x=nothing, + value=nothing, + gradient=nothing, + jacobian=nothing, + hessian=nothing, + context::Tuple=(), + op=nothing, + exception=nothing, + inputs=nothing, + allocations_safe::Bool=true, +) + return TestCase( + name, + tag, + f, + x_proto, + x, + value, + gradient, + jacobian, + hessian, + context, + op, + exception, + inputs, + allocations_safe, + ) +end + struct QuadraticProblem end (::QuadraticProblem)(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) @@ -12,7 +84,7 @@ struct VectorValuedProblem end _context_problem(y::AbstractVector{<:Real}, offset) = -0.5 * (y[1] - offset)^2 -function AbstractPPL.generate_testcases() +function AbstractPPL.generate_testcases(::Val{:vector}) return ( TestCase( "quadratic (scalar output)", @@ -213,7 +285,7 @@ function AbstractPPL.generate_testcases() ) end -function AbstractPPL.generate_namedtuple_testcases() +function AbstractPPL.generate_testcases(::Val{:namedtuple}) return ( TestCase( "scalar output over (x::Real, y::Vector)", diff --git a/src/AbstractPPL.jl b/src/AbstractPPL.jl index 4aad5891..78f0748f 100644 --- a/src/AbstractPPL.jl +++ b/src/AbstractPPL.jl @@ -15,99 +15,20 @@ using .Evaluators: prepare, value_and_gradient!!, value_and_jacobian!!, value_gradient_and_hessian!!, order """ - TestCase(name, tag, f, x_proto; x, value, gradient, jacobian, hessian, - context=(), op, exception, inputs) - -Single tagged case for AD conformance testing. The `tag::Symbol` selects how -the case is run; the kwargs populate only the fields the tag uses. - -Reserved tags (recognised by [`run_testcase`](@ref)): - - - `:vector` — vector input, scalar output (`gradient`) or vector output - (`jacobian`). - - `:hessian` — order=2 round-trip on scalar output. - - `:context` — scalar-output gradient with a non-empty `context::Tuple` - passed to `prepare`. - - `:edge` — error case; `op(prepared, x)` must throw `exception`. - - `:cache_reuse` — multiple inputs against a single prepared evaluator - (`inputs::Vector{<:NamedTuple}`, with `(x=, value=, - gradient=)` or `(x=, value=, jacobian=)` per row). - - `:namedtuple` — NamedTuple input and gradient; Mooncake-only. -""" -struct TestCase - name::String - tag::Symbol - f::Any - x_proto::Any - x::Any - value::Any - gradient::Any - jacobian::Any - hessian::Any - context::Tuple - op::Any - exception::Any - inputs::Any - # Cases with an allocating primal (vector-output result vectors, the - # empty-input shortcut's `T[]`) or shapes the original `:allocations` group - # never covered (hessian, cache-reuse, edge) set this to `false` — the - # runner then skips the `allocations=` check regardless of caller intent. - allocations_safe::Bool -end -function TestCase( - name, - tag::Symbol, - f, - x_proto; - x=nothing, - value=nothing, - gradient=nothing, - jacobian=nothing, - hessian=nothing, - context::Tuple=(), - op=nothing, - exception=nothing, - inputs=nothing, - allocations_safe::Bool=true, -) - return TestCase( - name, - tag, - f, - x_proto, - x, - value, - gradient, - jacobian, - hessian, - context, - op, - exception, - inputs, - allocations_safe, - ) -end + generate_testcases(::Val{group}) -""" - generate_testcases() - -Return a tuple of conformance [`TestCase`](@ref)s for vector-input AD -backends. Iterate and pass each to [`run_testcase`](@ref). +Return a tuple of AD conformance test cases for the input-shape `group`. +Reserved groups: `:vector` (vector input) and `:namedtuple` (NamedTuple +input; Mooncake-only). Iterate and pass each to [`run_testcase`](@ref). +Implemented by the `Test` extension (`AbstractPPLTestExt`). """ function generate_testcases end -""" - generate_namedtuple_testcases() - -Like [`generate_testcases`](@ref) but for evaluators with `NamedTuple` input. -""" -function generate_namedtuple_testcases end - """ run_testcase(case; adtype, prepare_fn=AbstractPPL.prepare, atol=0, rtol=1e-10, check_dims=true, type_stability=:skip, allocations=:skip) -Run a single [`TestCase`](@ref) against an AD backend. `type_stability` and +Run a single conformance case against an AD backend. `type_stability` and `allocations` accept `:skip` / `:test` / `:broken` — `:test` asserts the invariant, `:broken` marks it `@test_broken` (use for backends with known regressions). Implemented by the `Test` extension. @@ -119,7 +40,7 @@ function run_testcase end Meta.parse( "public prepare, value_and_gradient!!, value_and_jacobian!!, " * "value_gradient_and_hessian!!, order, " * - "generate_testcases, generate_namedtuple_testcases, run_testcase, TestCase", + "generate_testcases, run_testcase", ), ) end diff --git a/test/ext/differentiationinterface/main.jl b/test/ext/differentiationinterface/main.jl index e9a5cf05..e709bbd9 100644 --- a/test/ext/differentiationinterface/main.jl +++ b/test/ext/differentiationinterface/main.jl @@ -23,7 +23,7 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) @testset "AbstractPPLDifferentiationInterfaceExt" begin @testset "ForwardDiff" begin - for case in generate_testcases() + for case in generate_testcases(Val(:vector)) run_testcase(case; adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6) end end @@ -34,7 +34,7 @@ quadratic(x::AbstractVector{<:Real}) = sum(xi -> xi^2, x) # (compiled tape doesn't support `prepare_hessian`). @testset "ReverseDiff (compiled tape)" begin adtype = AutoReverseDiff(; compile=true) - for case in generate_testcases() + for case in generate_testcases(Val(:vector)) case.tag === :hessian && continue run_testcase(case; adtype, atol=1e-6, rtol=1e-6) end diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 435fa176..1ff6a03a 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -15,7 +15,7 @@ using Test alloc_state = VERSION < v"1.11" ? :broken : :test @testset "ForwardDiff (default chunk)" begin - for case in generate_testcases() + for case in generate_testcases(Val(:vector)) run_testcase( case; adtype=AutoForwardDiff(), @@ -31,7 +31,7 @@ using Test # case (x of length 1) and `:edge` cases (chunk doesn't apply). @testset "ForwardDiff (explicit chunk)" begin ad = AutoForwardDiff(; chunksize=2) - for case in generate_testcases() + for case in generate_testcases(Val(:vector)) case.tag ∈ (:vector, :cache_reuse, :hessian) || continue run_testcase(case; adtype=ad, atol=1e-6, rtol=1e-6) end diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index eda4e254..a8bfef35 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -4,12 +4,7 @@ Pkg.develop(; path=joinpath(@__DIR__, "..", "..", "..")) Pkg.instantiate() using AbstractPPL: - AbstractPPL, - prepare, - generate_testcases, - generate_namedtuple_testcases, - run_testcase, - value_and_gradient!! + AbstractPPL, prepare, generate_testcases, run_testcase, value_and_gradient!! using ADTypes: AutoMooncake, AutoMooncakeForward using Mooncake using Test @@ -50,7 +45,7 @@ end ("Mooncake (forward)", AutoMooncakeForward()), ) @testset "$label" begin - for case in generate_testcases() + for case in generate_testcases(Val(:vector)) run_testcase( case; adtype, @@ -60,7 +55,7 @@ end type_stability=_mooncake_inferred(case, adtype), ) end - for case in generate_namedtuple_testcases() + for case in generate_testcases(Val(:namedtuple)) run_testcase(case; adtype, atol=1e-6, rtol=1e-6) end end From 650d443d13fbb1e2c7b620534538f424310418b0 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 17:57:41 +0100 Subject: [PATCH 18/21] Merge :vector and :context runners The two `_run(Val{...})` methods differed only by passing `context=case.context` to `prepare` (no-op for `:vector` cases since `case.context` defaults to `()`). Collapse into one method with `Union{Val{:vector},Val{:context}}` dispatch. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 33 +++++---------------------------- 1 file changed, 5 insertions(+), 28 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 430fb7fc..7cb147c1 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -344,8 +344,11 @@ function AbstractPPL.run_testcase(case::TestCase; kwargs...) return nothing end +# `:vector` and `:context` share a runner — `case.context` defaults to `()` so +# threading it through `prepare` is a no-op on `:vector` cases that don't set +# it. function _run( - ::Val{:vector}, + ::Union{Val{:vector},Val{:context}}, case; adtype, prepare_fn=AbstractPPL.prepare, @@ -355,7 +358,7 @@ function _run( type_stability::Symbol=:skip, allocations::Symbol=:skip, ) - prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims) + prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims, context=case.context) @test AbstractPPL.order(prepared) == 1 @test prepared(case.x) ≈ case.value atol = atol rtol = rtol @@ -385,32 +388,6 @@ function _run( return nothing end -function _run( - ::Val{:context}, - case; - adtype, - prepare_fn=AbstractPPL.prepare, - atol=0, - rtol=1e-10, - check_dims::Bool=true, - type_stability::Symbol=:skip, - allocations::Symbol=:skip, -) - prepared = prepare_fn(adtype, case.f, case.x_proto; check_dims, context=case.context) - @test AbstractPPL.order(prepared) == 1 - @test prepared(case.x) ≈ case.value atol = atol rtol = rtol - 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 - _maybe_check_alloc!( - case, allocations, AbstractPPL.value_and_gradient!!, prepared, case.x - ) - _maybe_check_inferred!( - type_stability, AbstractPPL.value_and_gradient!!, prepared, case.x - ) - return nothing -end - function _run( ::Val{:hessian}, case; From 80d8ade49078e8724dd13227727f3c390c0f2e4a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 18:01:12 +0100 Subject: [PATCH 19/21] Document the remaining allocations_safe=false reasons The empty-input hessian case and both cache-reuse cases set `allocations_safe=false` without an inline reason, while the other four instances do. Add brief explanations to match. Co-Authored-By: Claude Opus 4.7 (1M context) --- ext/AbstractPPLTestExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/AbstractPPLTestExt.jl b/ext/AbstractPPLTestExt.jl index 7cb147c1..05f9cdd1 100644 --- a/ext/AbstractPPLTestExt.jl +++ b/ext/AbstractPPLTestExt.jl @@ -155,7 +155,7 @@ function AbstractPPL.generate_testcases(::Val{:vector}) value=7.5, gradient=Float64[], hessian=zeros(0, 0), - allocations_safe=false, + allocations_safe=false, # empty-input hessian shortcut allocates ), # value_gradient_and_hessian!! rejects order=1 preps regardless of arity; # both paths share the dispatch so one case suffices. @@ -268,7 +268,7 @@ function AbstractPPL.generate_testcases(::Val{:vector}) (x=[4.0, 5.0, 6.0], value=77.0, gradient=[8.0, 10.0, 12.0]), (x=[0.5, -1.0, 2.0], value=5.25, gradient=[1.0, -2.0, 4.0]), ], - allocations_safe=false, + allocations_safe=false, # cache-reuse loops aren't single-call alloc tests ), TestCase( "vector output, cache reuse", @@ -280,7 +280,7 @@ function AbstractPPL.generate_testcases(::Val{:vector}) (x=[5.0, 1.0, 7.0], value=[5.0, 8.0], jacobian=[1.0 5.0 0.0; 0.0 1.0 1.0]), (x=[0.0, 4.0, -2.0], value=[0.0, 2.0], jacobian=[4.0 0.0 0.0; 0.0 1.0 1.0]), ], - allocations_safe=false, + allocations_safe=false, # cache-reuse loops aren't single-call alloc tests ), ) end From 1ea0a78d79123c65d613ebee2b97f4d158790c48 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 18:13:52 +0100 Subject: [PATCH 20/21] Drop Julia 1.10 broken-marker on scalar-gradient allocations CI reported `Unexpected Pass` on Julia 1.10 for `quadratic (scalar output)` (FD + Mooncake) and `scalar gradient with context` (FD): the recent FD-ext tweaks (skip-checktag, hoisted target/tag locals) made these paths alloc-free on 1.10 too. Drop the `VERSION < v"1.11"` gating; the per-case `allocations_safe=false` still filters the genuinely-allocating paths (vector jacobian, empty-input shortcuts, hessian, cache-reuse loops). Co-Authored-By: Claude Opus 4.7 (1M context) --- test/ext/forwarddiff/main.jl | 6 +----- test/ext/mooncake/main.jl | 7 ++----- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/test/ext/forwarddiff/main.jl b/test/ext/forwarddiff/main.jl index 1ff6a03a..5173339b 100644 --- a/test/ext/forwarddiff/main.jl +++ b/test/ext/forwarddiff/main.jl @@ -10,10 +10,6 @@ using ForwardDiff using Test @testset "AbstractPPLForwardDiffExt" begin - # Julia 1.10 heap-allocates closure captures the 1.11+ runtime elides; mark - # allocations broken on min so the regression check stays honest on latest. - alloc_state = VERSION < v"1.11" ? :broken : :test - @testset "ForwardDiff (default chunk)" begin for case in generate_testcases(Val(:vector)) run_testcase( @@ -21,7 +17,7 @@ using Test adtype=AutoForwardDiff(), atol=1e-6, rtol=1e-6, - allocations=alloc_state, + allocations=:test, type_stability=:test, ) end diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index a8bfef35..a1e7173e 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -13,16 +13,13 @@ using Test # * `value_and_jacobian!!` allocates fresh cotangent/Jacobian buffers on # every call (both modes); forward-mode Jacobian return type infers as # `Tuple{Any, Union{Array{T,3}, Matrix}}`. -# * `value_and_gradient!!` on a context-lowered prep splats `args_to_zero` -# per call (reverse mode allocates; forward mode also fails inference). -# Julia 1.10 also heap-allocates `Fix2`/closure captures that 1.11+ elides. +# * `value_and_gradient!!` on a forward-mode context-lowered prep splats +# `args_to_zero` per call and allocates; forward mode also fails inference. function _mooncake_alloc(case, adtype) if case.tag === :vector && case.jacobian !== nothing return :broken elseif case.tag === :context && adtype isa AutoMooncakeForward return :broken - elseif VERSION < v"1.11" - return :broken else return :test end From d84606a7eb91a2755374554f20a0d76d7d86e683 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Mon, 25 May 2026 18:24:45 +0100 Subject: [PATCH 21/21] =?UTF-8?q?Skip=20Mooncake=20:alloc=20checks=20on=20?= =?UTF-8?q?Julia=201.10=20=E2=80=94=20they're=20resolver-flaky?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI reported back-to-back inconsistent results on Julia 1.10 for the same code: one run had Mooncake's scalar-gradient `@allocated` come out 0 (an Unexpected Pass when marked `:broken`), the next run had it at 256 (a Test Failed when marked `:test`). The dependency resolver picks slightly different Mooncake versions between runs, and Mooncake 0.5.x's allocation behaviour on 1.10 isn't stable across them. Set `_mooncake_alloc` to return `:skip` on Julia 1.10 instead of either `:test` or `:broken` — that way the check doesn't fire on min, regardless of which Mooncake version the resolver picked. Latest-Julia coverage is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) --- test/ext/mooncake/main.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/ext/mooncake/main.jl b/test/ext/mooncake/main.jl index a1e7173e..50c2f2d3 100644 --- a/test/ext/mooncake/main.jl +++ b/test/ext/mooncake/main.jl @@ -20,6 +20,10 @@ function _mooncake_alloc(case, adtype) return :broken elseif case.tag === :context && adtype isa AutoMooncakeForward return :broken + elseif VERSION < v"1.11" + # Mooncake's value_and_gradient!! allocations are flaky on Julia 1.10 + # (resolver-dependent: some Mooncake versions alloc, others don't). + return :skip else return :test end