diff --git a/CLAUDE.md b/CLAUDE.md index c8ba5276..43a5a59b 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 bba7f9ae..6a7efb06 100644 --- a/Project.toml +++ b/Project.toml @@ -4,9 +4,9 @@ version = "0.15.24" [deps] AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -42,12 +42,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 9698442a..6e0d207a 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 95305ae6..00252f6e 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/docs/src/vector.md b/docs/src/vector.md index 975f74cc..31b33993 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. diff --git a/src/vector/test_utils.jl b/src/vector/test_utils.jl index fe73053f..60544d15 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 2a5e13a6..3facb399 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 36e6c82d..5220e13c 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 38254885..c4d13fcc 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,9 +26,9 @@ const ENZYME_REVERSE = AutoEnzyme(; ) const adtypes = [ENZYME_FORWARD, ENZYME_REVERSE] -# Triple-nested tuple-of-products (e.g. `product_distribution(p1t, p1t, p1t)`) trip the -# Enzyme runtime activity inference. Skip the AD section for these cases (non-AD tests -# still run). +# 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 @@ -43,7 +43,7 @@ end # `: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] - return DI.AbstractADType[] + return ADTypes.AbstractADType[] end # This entire test suite is broken on 1.11. diff --git a/test/integration_tests/mooncake/Project.toml b/test/integration_tests/mooncake/Project.toml index a914acb9..c036282e 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 02586355..88df74ee 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 @@ -59,7 +59,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 2178ac33..c8cea435 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 f9229894..42970d7d 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 522069ea..6182aa0e 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 5409364d..ed71100b 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(