From 795cb5968ebf53c70f554bd2f9eb419bb57bcbe3 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 19 May 2026 21:37:38 +0100 Subject: [PATCH 1/2] Use AbstractPPL AD interface in test infra; drop DI from package deps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `src/vector/test_utils.jl` and `test/test_resources.jl` route all AD calls through `AbstractPPL.prepare` + `value_and_jacobian!!` / `value_and_gradient!!` in place of `DifferentiationInterface.jacobian` / `DI.gradient`. The package no longer depends on DifferentiationInterface directly; it remains a transitive test dep so AbstractPPL can dispatch to the DI extension for the `AutoForwardDiff` reference and for any other adtype that lacks a direct AbstractPPL extension. Other changes that follow from the routing switch: - `Project.toml` and the per-integration-suite `Project.toml`s bump `AbstractPPL = "0.15"` and add `ADTypes = "1"`; DI moves to a test-only dep. - `src/vector/test_utils.jl`: `test_optics` and `test_logjac` skip the AD-based check for discrete distributions — `AbstractPPL.prepare` rejects the integer vectors that `to_vec(d)` returns, and those checks are tautological for identity link transforms anyway. - `test/integration_tests/mooncake/main.jl`: PlanarLayer-inverse cases under `AutoMooncakeForward` hit Mooncake's batched NDual code path, which has no rule for `find_alpha`, so mark them `broken=true`. - `test/integration_tests/enzyme/main.jl`: the triple-nested product cases throw an Enzyme runtime `TypeError` that `@test_broken` doesn't reliably catch; switch to `skip=true` so non-AD checks still run. - `docs/src/defining_examples.md`: the StereographicProj example now uses `ForwardDiff.jacobian` directly instead of `DI.jacobian`. Co-Authored-By: Claude Opus 4.7 (1M context) --- CLAUDE.md | 2 +- Project.toml | 6 +- docs/Project.toml | 1 - docs/src/defining_examples.md | 8 +- src/vector/test_utils.jl | 155 ++++++++++++++---- test/Project.toml | 6 +- test/integration_tests/enzyme/Project.toml | 4 +- test/integration_tests/enzyme/main.jl | 26 +-- test/integration_tests/mooncake/Project.toml | 4 +- test/integration_tests/mooncake/main.jl | 12 +- .../reversediff/Project.toml | 4 +- test/integration_tests/reversediff/main.jl | 6 +- test/runtests.jl | 10 +- test/test_resources.jl | 42 +++-- 14 files changed, 204 insertions(+), 82 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index c8ba52766..43a5a59b9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -15,7 +15,7 @@ The code in `src/` outside of `src/vector/` is largely legacy. It still works an ## Testing -Run tests with `julia --project -e 'using Pkg; Pkg.test()'`. +Run tests with `julia --project -e 'using Pkg; Pkg.test()'`. The Enzyme integration suite lives in `test/integration_tests/enzyme/` and runs separately (see `.github/workflows/Enzyme.yml`). Always refresh each environment (`Pkg.update()` / `up`) before tests or doc builds — a stale manifest can cause subtle resolution and loading issues. `Bijectors.VectorBijectors.test_all(d)` is the single source of truth for whether a VectorBijectors implementation is correct. If it passes, the implementation is almost certainly fine. It checks roundtrips, type stability, vector lengths, optics, allocations, log-Jacobian correctness (against ForwardDiff), and AD compatibility across backends. diff --git a/Project.toml b/Project.toml index ea949b413..5edeacd37 100644 --- a/Project.toml +++ b/Project.toml @@ -4,10 +4,10 @@ version = "0.15.24" [deps] AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" @@ -40,12 +40,12 @@ BijectorsReverseDiffChainRulesExt = ["ChainRules", "ReverseDiff"] BijectorsReverseDiffExt = "ReverseDiff" [compat] -AbstractPPL = "0.14" +AbstractPPL = "0.15" +ADTypes = "1" ArgCheck = "1, 2" ChainRules = "1" ChainRulesCore = "0.10.11, 1" ChangesOfVariables = "0.1" -DifferentiationInterface = "0.7.14" Distributions = "0.25.33" DocStringExtensions = "0.9" EnzymeCore = "0.8.15" diff --git a/docs/Project.toml b/docs/Project.toml index 9698442af..6e0d207a8 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,7 +1,6 @@ [deps] Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" diff --git a/docs/src/defining_examples.md b/docs/src/defining_examples.md index 95305ae69..00252f6e4 100644 --- a/docs/src/defining_examples.md +++ b/docs/src/defining_examples.md @@ -199,6 +199,8 @@ If we try to calculate a Jacobian for `StereographicProj()`, we will just get a So, we need to take an extra step to map from the independent coordinates of $x$ (i.e., $x_1$ and $x_2$) to the full 3D coordinates, and _then_ to the plane: ```@example stereographic +using ForwardDiff, LinearAlgebra + sgn = 1 function full_transform(x12) @@ -207,12 +209,8 @@ function full_transform(x12) return StereographicProj()(x123) end -import DifferentiationInterface as DI -using FiniteDifferences, LinearAlgebra x = [0.3, 0.4, sgn * sqrt(1 - 0.3^2 - 0.4^2)] - -adtype = DI.AutoFiniteDifferences(; fdm=central_fdm(5, 1)) -jac = DI.jacobian(full_transform, adtype, x[1:2]) +jac = ForwardDiff.jacobian(full_transform, x[1:2]) logjac = logabsdet(jac)[1] ``` diff --git a/src/vector/test_utils.jl b/src/vector/test_utils.jl index fe73053f9..60544d157 100644 --- a/src/vector/test_utils.jl +++ b/src/vector/test_utils.jl @@ -1,14 +1,14 @@ using Test using LinearAlgebra: logabsdet, Cholesky, UpperTriangular, LowerTriangular -import DifferentiationInterface as DI - -# Would like to use FiniteDifferences, but very easy to run into issues with -# https://juliadiff.org/FiniteDifferences.jl/latest/#Dealing-with-Singularities -const ref_adtype = DI.AutoForwardDiff() +using ADTypes +using AbstractPPL: AbstractPPL # ForwardDiff is the baked-in reference; per-backend integration suites under -# `test/integration_tests/` pass their own `adtypes` to `test_all`. -const default_adtypes = [ref_adtype] +# `test/integration_tests/` pass their own `adtypes` to `test_all`. Routing the reference +# through `AbstractPPL.prepare(ref_adtype, ...)` keeps the call path uniform with the +# backend AD calls. +const ref_adtype = AutoForwardDiff() +const default_adtypes = ADTypes.AbstractADType[] _get_value_support(::D.Distribution{<:Any,VS}) where {VS<:D.ValueSupport} = VS @@ -228,7 +228,7 @@ function test_all( d::D.Distribution; expected_zero_allocs=(), adtypes=default_adtypes, - broken_adtypes=DI.AbstractADType[], + broken_adtypes=ADTypes.AbstractADType[], ad_atol=1e-10, ad_rtol=sqrt(eps()), roundtrip_atol=1e-10, @@ -409,6 +409,10 @@ function test_optics(d::D.Distribution) end end + # Discrete: link transform is identity, so the AD Jacobian sparsity test below is + # trivial — skip it. + d isa D.Distribution{<:Any,D.Discrete} && return nothing + @testset "linked_optic_vec: $(_name(d))" begin # This is a lot harder to test. What we need to prove is that # x = rand(d) @@ -422,7 +426,10 @@ function test_optics(d::D.Distribution) x = rand(d) xvec = to_vec(d)(x) yvec = to_linked_vec(d)(x) - J = DI.jacobian(to_linked_vec(d) ∘ from_vec(d), ref_adtype, xvec) + f = to_linked_vec(d) ∘ from_vec(d) + J = last( + AbstractPPL.value_and_jacobian!!(AbstractPPL.prepare(ref_adtype, f, xvec), xvec) + ) o = optic_vec(d) lo = linked_optic_vec(d) for i in 1:length(yvec) @@ -537,6 +544,10 @@ function test_logjac(d::D.Distribution, atol, rtol) end end + # Discrete: link transform is identity, so the linked logjac is trivially zero on + # both sides of the comparison — skip the AD-based check. + d isa D.Distribution{<:Any,D.Discrete} && return nothing + # Link logjacs will not be zero, so we need to check against a chosen backend. Because # Jacobians need to map from vector to vector, here we test the transformation of the # vectorised form to the linked vectorised form via the original sample. @@ -566,7 +577,15 @@ function test_logjac(d::D.Distribution, atol, rtol) # to make sure that the Jacobian is square. ad_xvec = to_vec_for_logjac_test(d)(x) ad_ffwd = to_linked_vec(d) ∘ from_vec_for_logjac_test(d) - ad_logjac = first(logabsdet(DI.jacobian(ad_ffwd, ref_adtype, ad_xvec))) + ad_logjac = first( + logabsdet( + last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, ad_ffwd, ad_xvec), ad_xvec + ), + ), + ), + ) @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol end @@ -579,7 +598,15 @@ function test_logjac(d::D.Distribution, atol, rtol) # For the AD calculation we need to use to/from_vec_for_logjac_test instead, # to make sure that the Jacobian is square. ad_frvs = to_vec_for_logjac_test(d) ∘ from_linked_vec(d) - ad_logjac = first(logabsdet(DI.jacobian(ad_frvs, ref_adtype, yvec))) + ad_logjac = first( + logabsdet( + last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, ad_frvs, yvec), yvec + ), + ), + ), + ) @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol end end @@ -590,7 +617,7 @@ end Test that various AD backends can differentiate the conversions to and from vector and linked vector forms for the given distribution `d`. """ -function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol) +function test_ad(d::D.Distribution, adtypes::Vector{<:ADTypes.AbstractADType}, atol, rtol) # Discrete-distribution link transforms are identities; nothing AD-meaningful to check. d isa D.Distribution{<:Any,D.Discrete} && return nothing @@ -598,18 +625,34 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, x = _rand_safe_ad(d) xvec = to_vec(d)(x) ffwd = to_linked_vec(d) ∘ from_vec(d) - ref_jac = DI.jacobian(ffwd, ref_adtype, xvec) + ref_jac = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, ffwd, xvec), xvec + ), + ) ladj(xvec) = last(with_logabsdet_jacobian(ffwd, xvec)) - ref_grad_ladj = DI.gradient(ladj, ref_adtype, xvec) + ref_grad_ladj = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(ref_adtype, ladj, xvec), xvec + ), + ) for adtype in adtypes @testset let x = x, adtype = adtype, d = d - ad_jac = DI.jacobian(ffwd, adtype, xvec) + ad_jac = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(adtype, ffwd, xvec), xvec + ), + ) @test ref_jac ≈ ad_jac atol = atol rtol = rtol end @testset let x = x, adtype = adtype, d = d - ad_grad_ladj = DI.gradient(ladj, adtype, xvec) + ad_grad_ladj = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(adtype, ladj, xvec), xvec + ), + ) @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol end end @@ -619,19 +662,35 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, x = _rand_safe_ad(d) yvec = to_linked_vec(d)(x) frvs = to_vec(d) ∘ from_linked_vec(d) - ref_jac = DI.jacobian(frvs, ref_adtype, yvec) + ref_jac = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, frvs, yvec), yvec + ), + ) ladj(yvec) = last(with_logabsdet_jacobian(frvs, yvec)) - ref_grad_ladj = DI.gradient(ladj, ref_adtype, yvec) + ref_grad_ladj = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(ref_adtype, ladj, yvec), yvec + ), + ) for adtype in adtypes @testset let x = x, adtype = adtype, d = d - ad_jac = DI.jacobian(frvs, adtype, yvec) + ad_jac = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(adtype, frvs, yvec), yvec + ), + ) @test ref_jac ≈ ad_jac atol = atol rtol = rtol end @testset let x = x, adtype = adtype, d = d - ad_grad_ladj = DI.gradient(ladj, adtype, yvec) + ad_grad_ladj = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(adtype, ladj, yvec), yvec + ), + ) @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol end end @@ -641,7 +700,7 @@ end # Broken-mode counterpart to `test_ad`: collapses the four per-adtype AD comparisons into # one `@test_broken` per adtype so a partially-broken backend stays cleanly broken. function _test_ad_broken( - d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol + d::D.Distribution, adtypes::Vector{<:ADTypes.AbstractADType}, atol, rtol ) d isa D.Distribution{<:Any,D.Discrete} && return nothing @testset "AD (broken): $(_name(d))" begin @@ -652,27 +711,65 @@ function _test_ad_broken( frvs = to_vec(d) ∘ from_linked_vec(d) ladj_fwd(xv) = last(with_logabsdet_jacobian(ffwd, xv)) ladj_rev(yv) = last(with_logabsdet_jacobian(frvs, yv)) - ref_jac_fwd = DI.jacobian(ffwd, ref_adtype, xvec) - ref_grad_ladj_fwd = DI.gradient(ladj_fwd, ref_adtype, xvec) - ref_jac_rev = DI.jacobian(frvs, ref_adtype, yvec) - ref_grad_ladj_rev = DI.gradient(ladj_rev, ref_adtype, yvec) + ref_jac_fwd = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, ffwd, xvec), xvec + ), + ) + ref_grad_ladj_fwd = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(ref_adtype, ladj_fwd, xvec), xvec + ), + ) + ref_jac_rev = last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(ref_adtype, frvs, yvec), yvec + ), + ) + ref_grad_ladj_rev = last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(ref_adtype, ladj_rev, yvec), yvec + ), + ) for adtype in adtypes @testset let x = x, adtype = adtype, d = d @test_broken ( isapprox( - DI.jacobian(ffwd, adtype, xvec), ref_jac_fwd; atol=atol, rtol=rtol + last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(adtype, ffwd, xvec), xvec + ), + ), + ref_jac_fwd; + atol=atol, + rtol=rtol, ) && isapprox( - DI.gradient(ladj_fwd, adtype, xvec), + last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(adtype, ladj_fwd, xvec), xvec + ), + ), ref_grad_ladj_fwd; atol=atol, rtol=rtol, ) && isapprox( - DI.jacobian(frvs, adtype, yvec), ref_jac_rev; atol=atol, rtol=rtol + last( + AbstractPPL.value_and_jacobian!!( + AbstractPPL.prepare(adtype, frvs, yvec), yvec + ), + ), + ref_jac_rev; + atol=atol, + rtol=rtol, ) && isapprox( - DI.gradient(ladj_rev, adtype, yvec), + last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(adtype, ladj_rev, yvec), yvec + ), + ), ref_grad_ladj_rev; atol=atol, rtol=rtol, diff --git a/test/Project.toml b/test/Project.toml index 2a5e13a6d..3facb3994 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" @@ -27,15 +28,16 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" [compat] +ADTypes = "1" AbstractMCMC = "5" -AbstractPPL = "0.14" +AbstractPPL = "0.15" AdvancedHMC = "0.6, 0.7, 0.8" Bijectors = "0.15" ChainRules = "1" ChainRulesTestUtils = "0.7, 1" ChangesOfVariables = "0.1" Combinatorics = "1.0.2" -DifferentiationInterface = "0.7.7" +DifferentiationInterface = "0.6, 0.7" Distributions = "0.25" FillArrays = "1" FiniteDifferences = "0.11, 0.12" diff --git a/test/integration_tests/enzyme/Project.toml b/test/integration_tests/enzyme/Project.toml index 36e6c82dd..5220e13c9 100644 --- a/test/integration_tests/enzyme/Project.toml +++ b/test/integration_tests/enzyme/Project.toml @@ -1,12 +1,12 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" -FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -17,4 +17,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Bijectors = {path = "../../.."} [compat] +AbstractPPL = "0.15" +DifferentiationInterface = "0.6, 0.7" Enzyme = "0.13.111" diff --git a/test/integration_tests/enzyme/main.jl b/test/integration_tests/enzyme/main.jl index bdd7495d3..c4d13fcc1 100644 --- a/test/integration_tests/enzyme/main.jl +++ b/test/integration_tests/enzyme/main.jl @@ -1,11 +1,11 @@ using ADTypes +using AbstractPPL: AbstractPPL using Bijectors -using DifferentiationInterface +using DifferentiationInterface: DifferentiationInterface using Distributions using Enzyme: Enzyme, set_runtime_activity, Forward, Reverse, Const using EnzymeTestUtils: test_forward, test_reverse using FillArrays: Fill -using FiniteDifferences using ForwardDiff: ForwardDiff using LinearAlgebra using PDMats @@ -26,22 +26,24 @@ const ENZYME_REVERSE = AutoEnzyme(; ) const adtypes = [ENZYME_FORWARD, ENZYME_REVERSE] -# Enzyme cannot differentiate through triple-nested tuple-of-products (e.g. -# `product_distribution(p1t, p1t, p1t)`); identify them structurally. +# Triple-nested tuple-of-products (e.g. `product_distribution(p1t, p1t, p1t)`) throw an +# Enzyme runtime activity TypeError that `@test_broken` doesn't reliably catch — skip the +# AD section for these cases (non-AD tests still run). function _enzyme_failing_product(d) d isa Distributions.ProductDistribution || return false d.dists isa Tuple || return false return first(d.dists) isa Union{Distributions.Product,Distributions.ProductDistribution} end -# Return the subset of `adtypes` that is known-broken for this case. Only Reverse mode -# hits https://github.com/EnzymeAD/Enzyme.jl/issues/2987 on `:reshaped_beta_special` on -# Julia 1.10 — Forward mode passes. Triple-nested products in `:type_unstable_products` -# defeat activity inference for both modes. +function skip_enzyme_ad(c::VectorTestCase) + return c.tag === :type_unstable_products && _enzyme_failing_product(c.dist) +end + +# Reverse mode hits https://github.com/EnzymeAD/Enzyme.jl/issues/2987 on +# `:reshaped_beta_special` on Julia 1.10 — Forward mode passes. function vector_broken_adtypes(c::VectorTestCase) c.tag === :reshaped_beta_special && VERSION < v"1.11-" && return [ENZYME_REVERSE] - c.tag === :type_unstable_products && _enzyme_failing_product(c.dist) && return adtypes - return DI.AbstractADType[] + return ADTypes.AbstractADType[] end # This entire test suite is broken on 1.11. @@ -109,6 +111,8 @@ end @testset "Enzyme vector test_all" begin for c in generate_vector_testcases() - run_vector_case(c, adtypes; broken_adtypes=vector_broken_adtypes(c)) + run_vector_case( + c, adtypes; broken_adtypes=vector_broken_adtypes(c), skip=skip_enzyme_ad(c) + ) end end diff --git a/test/integration_tests/mooncake/Project.toml b/test/integration_tests/mooncake/Project.toml index a914acb95..c036282e3 100644 --- a/test/integration_tests/mooncake/Project.toml +++ b/test/integration_tests/mooncake/Project.toml @@ -1,10 +1,10 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" -FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" @@ -17,4 +17,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Bijectors = {path = "../../.."} [compat] +AbstractPPL = "0.15" +DifferentiationInterface = "0.6, 0.7" Mooncake = "0.4, 0.5" diff --git a/test/integration_tests/mooncake/main.jl b/test/integration_tests/mooncake/main.jl index 77b416bc2..7ee68e203 100644 --- a/test/integration_tests/mooncake/main.jl +++ b/test/integration_tests/mooncake/main.jl @@ -1,9 +1,9 @@ using ADTypes +using AbstractPPL: AbstractPPL using Bijectors -using DifferentiationInterface +using DifferentiationInterface: DifferentiationInterface using Distributions using FillArrays: Fill -using FiniteDifferences using ForwardDiff: ForwardDiff using LinearAlgebra using Mooncake: Mooncake @@ -42,7 +42,13 @@ end @testset "Mooncake bijector AD" begin for c in generate_ad_testcases(), adtype in adtypes - run_ad_case(c, adtype) + # AbstractPPL's Mooncake extension routes `AutoMooncakeForward` through + # `Mooncake.prepare_derivative_cache` (batched NDual mode). `BijectorsMooncakeExt` + # only has `find_alpha` rules for `Mooncake.Dual`, so the PlanarLayer inverse path + # (which calls `find_alpha`) trips on `ceil(::Int, ::NDual)` inside Roots.ITP. + is_broken = + adtype isa AutoMooncakeForward && startswith(c.name, "PlanarLayer inverse") + run_ad_case(c, adtype; broken=is_broken) end end diff --git a/test/integration_tests/reversediff/Project.toml b/test/integration_tests/reversediff/Project.toml index 2178ac33d..c8cea4354 100644 --- a/test/integration_tests/reversediff/Project.toml +++ b/test/integration_tests/reversediff/Project.toml @@ -1,11 +1,11 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" -FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -17,4 +17,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Bijectors = {path = "../../.."} [compat] +AbstractPPL = "0.15" +DifferentiationInterface = "0.6, 0.7" ReverseDiff = "1.4.2" diff --git a/test/integration_tests/reversediff/main.jl b/test/integration_tests/reversediff/main.jl index f92298948..42970d7dd 100644 --- a/test/integration_tests/reversediff/main.jl +++ b/test/integration_tests/reversediff/main.jl @@ -1,10 +1,10 @@ using ADTypes +using AbstractPPL: AbstractPPL using Bijectors using ChainRules: ChainRules -using DifferentiationInterface +using DifferentiationInterface: DifferentiationInterface using Distributions using FillArrays: Fill -using FiniteDifferences using ForwardDiff: ForwardDiff using LinearAlgebra using PDMats @@ -21,7 +21,7 @@ const adtypes = [AutoReverseDiff(), AutoReverseDiff(; compile=true)] const _BROKEN_VECTOR_TAGS = (:lkj_matrix_dists, :order_joint) function vector_broken_adtypes(c::VectorTestCase) - return c.tag in _BROKEN_VECTOR_TAGS ? adtypes : DI.AbstractADType[] + return c.tag in _BROKEN_VECTOR_TAGS ? adtypes : ADTypes.AbstractADType[] end @testset "ReverseDiff bijector AD" begin diff --git a/test/runtests.jl b/test/runtests.jl index 522069ea4..6182aa0ef 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,7 +2,9 @@ using Bijectors using ChainRulesTestUtils using Combinatorics -using DifferentiationInterface +using ADTypes +using AbstractPPL: AbstractPPL +using DifferentiationInterface: DifferentiationInterface using FiniteDifferences using ForwardDiff using Functors @@ -56,12 +58,6 @@ include("bijectors/utils.jl") include("bijectors/simplex.jl") include("bijectors/equality.jl") include("bijectors/scale.jl") - - @testset "ForwardDiff bijector AD" begin - for c in generate_ad_testcases() - run_ad_case(c, AutoForwardDiff()) - end - end end if GROUP == "All" || GROUP == "Vector" || GROUP == "VectorProduct" diff --git a/test/test_resources.jl b/test/test_resources.jl index 5409364d5..ed71100bd 100644 --- a/test/test_resources.jl +++ b/test/test_resources.jl @@ -9,26 +9,25 @@ # Each integration suite iterates these once with a local `is_broken(c)` predicate so the # main file becomes one adtype list + one filter + two short loops. # -# `ADTestCase` cases run via `run_ad_case` (gradient comparison against a FiniteDifferences +# `ADTestCase` cases run via `run_ad_case` (gradient comparison against a ForwardDiff # reference). `VectorTestCase` cases run via `run_vector_case`, which delegates to # `VectorBijectors.test_all`; pass an explicit `adtypes` to exercise non-ForwardDiff -# backends (`VectorBijectors.test_all` defaults to `[AutoForwardDiff()]`). +# backends. `VectorBijectors.test_all` defaults to an empty adtype list — the integration +# suites pass their own. # -# Everything `VectorTestCase`-related lives in `test/vector_bijectors.jl`, which this file -# includes at the bottom — so a single `include("test_resources.jl")` is all integration -# tests need. +# `VectorTestCase` generators live in `test/vector/*.jl`, which this file includes at the +# bottom — so a single `include("test_resources.jl")` is all integration tests need. using Bijectors using Bijectors: ordered import Bijectors.VectorBijectors using Bijectors.VectorBijectors: from_linked_vec, from_vec, linked_vec_length, to_linked_vec, to_vec -using DifferentiationInterface -import DifferentiationInterface as DI +using ADTypes +using AbstractPPL: AbstractPPL using Distributions const _D = Distributions using FillArrays: Fill -using FiniteDifferences: central_fdm using LinearAlgebra using PDMats using StableRNGs: StableRNG @@ -79,12 +78,16 @@ generate_ad_testcases() = reduce(vcat, generate_testcases(Val(t)) for t in _AD_T # ===== ADTestCase runner ===== -const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) +const REF_BACKEND = AutoForwardDiff() function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) @info "testing AD for function $f with $backend" - ref_gradient = gradient(f, REF_BACKEND, x) - ad_gradient = gradient(f, backend, x) + ref_gradient = last( + AbstractPPL.value_and_gradient!!(AbstractPPL.prepare(REF_BACKEND, f, x), x) + ) + ad_gradient = last( + AbstractPPL.value_and_gradient!!(AbstractPPL.prepare(backend, f, x), x) + ) @test isapprox(ad_gradient, ref_gradient; rtol=rtol, atol=atol) end @@ -96,8 +99,16 @@ function run_ad_case(c::ADTestCase, adtype; broken::Bool=false, rtol=1e-6, atol= # the case flips to "unexpectedly passing" and the maintainer gets a nudge. @info "testing (broken) AD for function $(c.func) with $adtype" @test_broken isapprox( - gradient(c.func, adtype, c.arg), - gradient(c.func, REF_BACKEND, c.arg); + last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(adtype, c.func, c.arg), c.arg + ), + ), + last( + AbstractPPL.value_and_gradient!!( + AbstractPPL.prepare(REF_BACKEND, c.func, c.arg), c.arg + ), + ); rtol=rtol, atol=atol, ) @@ -320,7 +331,10 @@ generate_vector_testcases() = reduce(vcat, generate_testcases(Val(t)) for t in _ # ===== VectorTestCase runner ===== function run_vector_case( - c::VectorTestCase, adtypes=nothing; broken_adtypes=DI.AbstractADType[], skip::Bool=false + c::VectorTestCase, + adtypes=nothing; + broken_adtypes=ADTypes.AbstractADType[], + skip::Bool=false, ) if adtypes === nothing VectorBijectors.test_all( From 2c49e3f4cf4d70fbb6194fb8fddc056425e73e1a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 21 May 2026 10:47:23 +0100 Subject: [PATCH 2/2] Fix vector.md docs: default_adtypes is empty Co-Authored-By: Claude Opus 4.7 (1M context) --- docs/src/vector.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/vector.md b/docs/src/vector.md index 975f74cce..31b33993a 100644 --- a/docs/src/vector.md +++ b/docs/src/vector.md @@ -117,7 +117,7 @@ For more information about generally testing bijectors (and in particular how to One of the most tricky parts of testing Bijectors is ensuring that the transforms are compatible with automatic differentiation. This is important for DynamicPPL: we need to be able to compute the gradient of the log-density with respect to (possibly transformed) parameters, which may include the log-abs-det-Jacobian of the transformation. -`test_all` accepts an `adtypes` keyword argument and defaults to `[AutoForwardDiff()]`, so the main test suite exercises ForwardDiff as the reference backend. +`test_all` accepts an `adtypes` keyword argument that defaults to an empty list, so the main test suite does not exercise any AD backend directly — ForwardDiff is still used internally as the reference for Jacobian correctness via `ref_adtype`. ReverseDiff, Mooncake, and Enzyme are covered by standalone integration suites under `test/integration_tests//`, each of which passes its own `adtypes` list through to `test_all`. It is acceptable to skip tests for a particular backend if there are genuine upstream bugs, especially with ReverseDiff, which is not actively maintained. Where possible it is best to ensure that all backends are supported, and to use `@test_broken` to mark any known issues with specific backends.