From 117b8e85c106ad62c01b5a50a28f12a0ac2c4110 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 13 Aug 2025 17:29:09 +0100 Subject: [PATCH 01/28] Remove test_ad calls from interface tests --- test/ad/corr.jl | 50 ++++++++++++++++++++++-------------------- test/ad/utils.jl | 3 --- test/bijectors/corr.jl | 4 ---- 3 files changed, 26 insertions(+), 31 deletions(-) diff --git a/test/ad/corr.jl b/test/ad/corr.jl index c2c99033..a1d01b51 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -1,35 +1,37 @@ @testset "AD for VecCorrBijector" begin - d = 4 - dist = LKJ(d, 2.0) - b = bijector(dist) - binv = inverse(b) + @testset "d = $d" for d in (1, 2, 4) + dist = LKJ(d, 2.0) + b = bijector(dist) + binv = inverse(b) - x = rand(dist) - y = b(x) + x = rand(dist) + y = b(x) - test_ad(y) do x - sum(transform(b, binv(x))) - end - - test_ad(y) do y - sum(transform(binv, y)) + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), y) + # inverse only + test_ad(y -> sum(transform(binv, y)), y) end end @testset "AD for VecCholeskyBijector" begin - d = 4 - dist = LKJCholesky(d, 2.0) - b = bijector(dist) - binv = inverse(b) + @testset "d = $d, uplo = $uplo" for d in (1, 2, 4), uplo in ('U', 'L') + dist = LKJCholesky(d, 2.0, uplo) + b = bijector(dist) + binv = inverse(b) - x = rand(dist) - y = b(x) - - test_ad(y) do y - sum(transform(b, binv(y))) - end + x = rand(dist) + y = b(x) - test_ad(y) do y - sum(Bijectors.cholesky_upper(transform(binv, y))) + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), y) + # inverse only + test_ad(y -> sum(transform(binv, y)), y) + # additionally check that cholesky_{upper,lower} is differentiable + if uplo == 'U' + test_ad(y -> sum(Bijectors.cholesky_upper(transform(b, y))), y) + else + test_ad(y -> sum(Bijectors.cholesky_lower(transform(b, y))), y) + end end end diff --git a/test/ad/utils.jl b/test/ad/utils.jl index 2abf7622..2f9679e8 100644 --- a/test/ad/utils.jl +++ b/test/ad/utils.jl @@ -1,6 +1,3 @@ -# Figure out which AD backend to test -const AD = get(ENV, "AD", "All") - function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6) for b in broken if !( diff --git a/test/bijectors/corr.jl b/test/bijectors/corr.jl index 8a423bc3..6e153fab 100644 --- a/test/bijectors/corr.jl +++ b/test/bijectors/corr.jl @@ -32,8 +32,6 @@ using Bijectors: VecCorrBijector, VecCholeskyBijector, CorrBijector test_bijector(b, x; test_not_identity=d != 1, changes_of_variables_test=false) test_bijector(bvec, x; test_not_identity=d != 1, changes_of_variables_test=false) - test_ad(x -> sum(bvec(bvecinv(x))), yvec) - # Check that output sizes are computed correctly. tdist = transformed(dist) @test length(tdist) == length(yvec) @@ -64,8 +62,6 @@ end @test xinv.U ≈ cholesky(xinv_lkj).U - test_ad(x -> sum(b(binv(x))), y) - # test_bijector is commented out for now, # as isapprox is not defined for ::Cholesky types (the domain of LKJCholesky) # test_bijector(b, x; test_not_identity=d != 1, changes_of_variables_test=false) From d0e6a6fe1b484bfb6b1940bbf988cf29e9a31b9a Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 13 Aug 2025 17:30:43 +0100 Subject: [PATCH 02/28] remove dead code --- test/bijectors/utils.jl | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/test/bijectors/utils.jl b/test/bijectors/utils.jl index 8c31ccec..dbec0c9c 100644 --- a/test/bijectors/utils.jl +++ b/test/bijectors/utils.jl @@ -108,19 +108,3 @@ function test_functor(x, xs) @test x == re(_xs) @test _xs == xs end - -function test_bijector_parameter_gradient(b::Bijectors.Transform, x, y=b(x)) - args, re = Functors.functor(b) - recon(k, param) = re(merge(args, NamedTuple{(k,)}((param,)))) - - # Compute the gradient wrt. one argument at the time. - for (k, v) in pairs(args) - test_ad(p -> sum(transform(recon(k, p), x)), v) - test_ad(p -> logabsdetjac(recon(k, p), x), v) - - if Bijectors.isinvertible(b) - test_ad(p -> sum(transform(inv(recon(k, p)), y)), v) - test_ad(p -> logabsdetjac(inv(recon(k, p)), y), v) - end - end -end From fb6d71c40bbaebe006eacfb79e21cbf55bee477b Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 13 Aug 2025 17:34:07 +0100 Subject: [PATCH 03/28] don't remove that --- test/ad/utils.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/ad/utils.jl b/test/ad/utils.jl index 2f9679e8..2abf7622 100644 --- a/test/ad/utils.jl +++ b/test/ad/utils.jl @@ -1,3 +1,6 @@ +# Figure out which AD backend to test +const AD = get(ENV, "AD", "All") + function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6) for b in broken if !( From b020f9e3cbe8a46c0b66385abb5881b278378817 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 13 Aug 2025 17:34:43 +0100 Subject: [PATCH 04/28] Remove macOS from CI --- .github/workflows/AD.yml | 1 - .github/workflows/Interface.yml | 2 -- 2 files changed, 3 deletions(-) diff --git a/.github/workflows/AD.yml b/.github/workflows/AD.yml index c5c434c0..d56aa9bc 100644 --- a/.github/workflows/AD.yml +++ b/.github/workflows/AD.yml @@ -23,7 +23,6 @@ jobs: - '1' os: - ubuntu-latest - - macOS-latest AD: - Enzyme - ForwardDiff diff --git a/.github/workflows/Interface.yml b/.github/workflows/Interface.yml index f8f325c5..7ed12425 100644 --- a/.github/workflows/Interface.yml +++ b/.github/workflows/Interface.yml @@ -20,11 +20,9 @@ jobs: matrix: version: - 'min' - - 'lts' - '1' os: - ubuntu-latest - - macOS-latest steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 From cd4e62fec9eeb944aef9de14a9d9a14f44e42ec6 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Wed, 13 Aug 2025 17:43:43 +0100 Subject: [PATCH 05/28] fix AD --- test/ad/corr.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/ad/corr.jl b/test/ad/corr.jl index a1d01b51..42fefdfd 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -25,13 +25,13 @@ end # roundtrip test_ad(y -> sum(transform(b, binv(y))), y) - # inverse only - test_ad(y -> sum(transform(binv, y)), y) - # additionally check that cholesky_{upper,lower} is differentiable + # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, + # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't + # give a scalar) if uplo == 'U' - test_ad(y -> sum(Bijectors.cholesky_upper(transform(b, y))), y) + test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), y) else - test_ad(y -> sum(Bijectors.cholesky_lower(transform(b, y))), y) + test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), y) end end end From 78f3ef96addf5fb712e030ec225488f8c26cf319 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 18:45:50 +0100 Subject: [PATCH 06/28] Refactor AD tests --- .github/workflows/AD.yml | 45 --------- .github/workflows/{Interface.yml => CI.yml} | 9 +- test/ad/corr.jl | 14 +-- test/ad/flows.jl | 16 ++- test/ad/pd.jl | 24 ++--- test/ad/stacked.jl | 20 +--- test/ad/utils.jl | 103 ++------------------ test/runtests.jl | 15 ++- 8 files changed, 54 insertions(+), 192 deletions(-) delete mode 100644 .github/workflows/AD.yml rename .github/workflows/{Interface.yml => CI.yml} (83%) diff --git a/.github/workflows/AD.yml b/.github/workflows/AD.yml deleted file mode 100644 index d56aa9bc..00000000 --- a/.github/workflows/AD.yml +++ /dev/null @@ -1,45 +0,0 @@ -name: AD tests - -on: - push: - branches: - - main - pull_request: - -concurrency: - # Skip intermediate builds: always. - # Cancel intermediate builds: only if it is a pull request build. - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} - -jobs: - test: - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - version: - - 'min' - - '1' - os: - - ubuntu-latest - AD: - - Enzyme - - ForwardDiff - - Mooncake - - Tracker - - ReverseDiff - steps: - - uses: actions/checkout@v4 - - - uses: julia-actions/setup-julia@v2 - with: - version: ${{ matrix.version }} - - - uses: julia-actions/julia-buildpkg@v1 - - - uses: julia-actions/julia-runtest@v1 - continue-on-error: ${{ matrix.AD == 'Enzyme' && matrix.version == '1' }} - env: - GROUP: AD - AD: ${{ matrix.AD }} diff --git a/.github/workflows/Interface.yml b/.github/workflows/CI.yml similarity index 83% rename from .github/workflows/Interface.yml rename to .github/workflows/CI.yml index 7ed12425..b690de15 100644 --- a/.github/workflows/Interface.yml +++ b/.github/workflows/CI.yml @@ -1,4 +1,4 @@ -name: Interface tests +name: CI on: push: @@ -23,12 +23,17 @@ jobs: - '1' os: - ubuntu-latest + group: + - 'Interface' + - 'AD' + steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} + - uses: julia-actions/cache@v2 - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 env: - GROUP: Interface + GROUP: ${{ matrix.group }} diff --git a/test/ad/corr.jl b/test/ad/corr.jl index 42fefdfd..49e35c0c 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -1,4 +1,4 @@ -@testset "AD for VecCorrBijector" begin +@testset "VecCorrBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES @testset "d = $d" for d in (1, 2, 4) dist = LKJ(d, 2.0) b = bijector(dist) @@ -8,13 +8,13 @@ y = b(x) # roundtrip - test_ad(y -> sum(transform(b, binv(y))), y) + test_ad(y -> sum(transform(b, binv(y))), adtype, y) # inverse only - test_ad(y -> sum(transform(binv, y)), y) + test_ad(y -> sum(transform(binv, y)), adtype, y) end end -@testset "AD for VecCholeskyBijector" begin +@testset "VecCholeskyBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES @testset "d = $d, uplo = $uplo" for d in (1, 2, 4), uplo in ('U', 'L') dist = LKJCholesky(d, 2.0, uplo) b = bijector(dist) @@ -24,14 +24,14 @@ end y = b(x) # roundtrip - test_ad(y -> sum(transform(b, binv(y))), y) + test_ad(y -> sum(transform(b, binv(y))), adtype, y) # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't # give a scalar) if uplo == 'U' - test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), y) + test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), adtype, y) else - test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), y) + test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), adtype, y) end end end diff --git a/test/ad/flows.jl b/test/ad/flows.jl index bfcbaacc..34ad90aa 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -1,29 +1,35 @@ -@testset "PlanarLayer" begin +@testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES # logpdf of a flow with a planar layer and two-dimensional inputs - test_ad(randn(7)) do θ + function f(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) flow = transformed(MvNormal(zeros(2), I), layer) x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - test_ad(randn(11)) do θ + test_ad(f, adtype, randn(7)) + + function g(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) flow = transformed(MvNormal(zeros(2), I), layer) x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end + test_ad(g, adtype, randn(11)) # logpdf of a flow with the inverse of a planar layer and two-dimensional inputs - test_ad(randn(7)) do θ + function finv(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) flow = transformed(MvNormal(zeros(2), I), inverse(layer)) x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - test_ad(randn(11)) do θ + test_ad(finv, adtype, randn(7)) + + function ginv(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) flow = transformed(MvNormal(zeros(2), I), inverse(layer)) x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end + test_ad(ginv, adtype, randn(11)) end diff --git a/test/ad/pd.jl b/test/ad/pd.jl index 61f5cbd1..daa244c3 100644 --- a/test/ad/pd.jl +++ b/test/ad/pd.jl @@ -1,6 +1,6 @@ _topd(x) = x * x' + I -@testset "AD for PDVecBijector" begin +@testset "PDVecBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES d = 4 b = Bijectors.PDVecBijector() binv = inverse(b) @@ -9,21 +9,11 @@ _topd(x) = x * x' + I x = _topd(z) y = b(x) - test_ad(vec(z)) do x - sum(transform(b, _topd(reshape(x, d, d)))) - end + test_ad(x -> sum(transform(b, _topd(reshape(x, d, d)))), adtype, vec(z)) + test_ad(y -> sum(transform(binv, y)), adtype, y) - test_ad(y) do y - sum(transform(binv, y)) - end - - if AD == "ReverseDiff" # `AD` is defined in `test/ad/utils.jl`. - test_ad(y) do y - sum(Bijectors.cholesky_lower(transform(binv, y))) - end - - test_ad(y) do y - sum(Bijectors.cholesky_upper(transform(binv, y))) - end - end + # if occursin("ReverseDiff", backend_name) + test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), adtype, y) + test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), adtype, y) + # end end diff --git a/test/ad/stacked.jl b/test/ad/stacked.jl index 89f5f4be..fe86fcda 100644 --- a/test/ad/stacked.jl +++ b/test/ad/stacked.jl @@ -1,4 +1,4 @@ -@testset "AD for StackedBijector" begin +@testset "StackedBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES dist1 = Dirichlet(4, 1.0) b1 = bijector(dist1) @@ -17,22 +17,12 @@ y = vcat(y1, [y2]) x = binv(y) - test_ad(y) do x - sum(transform(b, binv(x))) - end - - test_ad(y) do y - sum(transform(binv, y)) - end + test_ad(y -> sum(transform(b, binv(y))), adtype, y) + test_ad(y -> sum(transform(binv, y)), adtype, y) bvec = Stacked([b1, b2], [1:4, 5:5]) bvec_inv = inverse(bvec) - test_ad(y) do x - sum(transform(bvec, binv(x))) - end - - test_ad(y) do y - sum(transform(bvec_inv, y)) - end + test_ad(y -> sum(transform(bvec, binv(y))), adtype, y) + test_ad(y -> sum(transform(bvec_inv, y)), adtype, y) end diff --git a/test/ad/utils.jl b/test/ad/utils.jl index 2abf7622..9ea29272 100644 --- a/test/ad/utils.jl +++ b/test/ad/utils.jl @@ -1,100 +1,9 @@ -# Figure out which AD backend to test -const AD = get(ENV, "AD", "All") +using DifferentiationInterface -function test_ad(f, x, broken=(); rtol=1e-6, atol=1e-6) - for b in broken - if !( - b in ( - :ForwardDiff, - :Mooncake, - :ReverseDiff, - :Enzyme, - :EnzymeForward, - :EnzymeReverse, - # The `Crash` ones indicate that the error will cause a Julia crash, and - # thus we can't even run `@test_broken on it. - :EnzymeForwardCrash, - :EnzymeReverseCrash, - ) - ) - error("Unknown broken AD backend: $b") - end - end +const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) - finitediff = FiniteDifferences.grad(central_fdm(5, 1), f, x)[1] - - if AD == "All" || AD == "ForwardDiff" - if :ForwardDiff in broken - @test_broken ForwardDiff.gradient(f, x) ≈ finitediff rtol = rtol atol = atol - else - @test ForwardDiff.gradient(f, x) ≈ finitediff rtol = rtol atol = atol - end - end - - if AD == "All" || AD == "ReverseDiff" - if :ReverseDiff in broken - @test_broken ReverseDiff.gradient(f, x) ≈ finitediff rtol = rtol atol = atol - else - @test ReverseDiff.gradient(f, x) ≈ finitediff rtol = rtol atol = atol - end - end - - if AD == "All" || AD == "Enzyme" - forward_broken = :EnzymeForward in broken || :Enzyme in broken - reverse_broken = :EnzymeReverse in broken || :Enzyme in broken - if !(:EnzymeForwardCrash in broken) - if forward_broken - @test_broken( - Enzyme.gradient(Forward, Enzyme.Const(f), x)[1] ≈ finitediff, - rtol = rtol, - atol = atol - ) - else - @test( - Enzyme.gradient(Forward, Enzyme.Const(f), x)[1] ≈ finitediff, - rtol = rtol, - atol = atol - ) - end - end - - if !(:EnzymeReverseCrash in broken) - if reverse_broken - @test_broken( - Enzyme.gradient(set_runtime_activity(Reverse), Enzyme.Const(f), x)[1] ≈ - finitediff, - rtol = rtol, - atol = atol - ) - else - @test( - Enzyme.gradient(set_runtime_activity(Reverse), Enzyme.Const(f), x)[1] ≈ - finitediff, - rtol = rtol, - atol = atol - ) - end - end - end - - if AD == "All" || AD == "Mooncake" - rule = Mooncake.build_rrule(f, x) - if :Mooncake in broken - @test_broken isapprox( - Mooncake.value_and_gradient!!(rule, f, x)[2][2], - finitediff; - rtol=rtol, - atol=atol, - ) - else - @test isapprox( - Mooncake.value_and_gradient!!(rule, f, x)[2][2], - finitediff; - rtol=rtol, - atol=atol, - ) - end - end - - return nothing +function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) + ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) + gradient = DifferentiationInterface.gradient(f, backend, x) + @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end diff --git a/test/runtests.jl b/test/runtests.jl index ca7131d0..f0412a61 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -33,6 +33,7 @@ using InverseFunctions: InverseFunctions using LazyArrays: LazyArrays const GROUP = get(ENV, "GROUP", "All") +const IS_PRERELEASE = !isempty(VERSION.prerelease) # Always include this since it can be useful for other tests. include("ad/utils.jl") @@ -52,7 +53,6 @@ if GROUP == "All" || GROUP == "Interface" include("bijectors/reshape.jl") include("bijectors/corr.jl") include("bijectors/product_bijector.jl") - include("distributionsad.jl") @testset "doctests" begin @@ -68,13 +68,20 @@ if GROUP == "All" || GROUP == "Interface" end end +TEST_ADTYPES = [ + ("ForwardDiff", AutoForwardDiff()), + ("ReverseDiff", AutoReverseDiff(; compile=false)), + ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), +] +if !IS_PRERELEASE + push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) +end + if GROUP == "All" || GROUP == "AD" include("ad/chainrules.jl") - if get(ENV, "AD", "All") in ("All", "Enzyme") - include("ad/enzyme.jl") - end include("ad/flows.jl") include("ad/pd.jl") include("ad/corr.jl") include("ad/stacked.jl") + include("ad/enzyme.jl") end From 982d19b976ecd6f36b32bd7c6e76860688b914dc Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 18:56:28 +0100 Subject: [PATCH 07/28] add doctests, fix other things --- .github/workflows/DocTests.yml | 41 +++++++++++++++++++++++++++++++ docs/Project.toml | 4 ++-- test/Project.toml | 1 + test/ad/utils.jl | 9 ------- test/runtests.jl | 44 +++++++++++++++++++--------------- 5 files changed, 69 insertions(+), 30 deletions(-) create mode 100644 .github/workflows/DocTests.yml delete mode 100644 test/ad/utils.jl diff --git a/.github/workflows/DocTests.yml b/.github/workflows/DocTests.yml new file mode 100644 index 00000000..73f337e1 --- /dev/null +++ b/.github/workflows/DocTests.yml @@ -0,0 +1,41 @@ +# We want to only run doctests on a single version of Julia, because +# things like error messages / output can change between versions and +# is fragile to test against. +name: Doctests + +on: + push: + branches: + - main + pull_request: + merge_group: + types: [checks_requested] + +# needed to allow julia-actions/cache to delete old caches that it has created +permissions: + actions: write + contents: read + +# Cancel existing tests on the same PR if a new commit is added to a pull request +concurrency: + group: ${{ github.workflow }}-${{ github.ref || github.run_id }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: julia-actions/setup-julia@v2 + with: + version: '1' + + - uses: julia-actions/cache@v2 + + - uses: julia-actions/julia-buildpkg@v1 + + - uses: julia-actions/julia-runtest@v1 + env: + GROUP: Doctests diff --git a/docs/Project.toml b/docs/Project.toml index e64b9b85..b235005d 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -4,6 +4,6 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" [compat] -Documenter = "0.27" +Documenter = "1" Functors = "0.3" -StableRNGs = "1" \ No newline at end of file +StableRNGs = "1" diff --git a/test/Project.toml b/test/Project.toml index 447c36dd..979b70a8 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" diff --git a/test/ad/utils.jl b/test/ad/utils.jl deleted file mode 100644 index 9ea29272..00000000 --- a/test/ad/utils.jl +++ /dev/null @@ -1,9 +0,0 @@ -using DifferentiationInterface - -const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) - -function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) - ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) - gradient = DifferentiationInterface.gradient(f, backend, x) - @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) -end diff --git a/test/runtests.jl b/test/runtests.jl index f0412a61..7632bd67 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2,6 +2,7 @@ using Bijectors using ChainRulesTestUtils using Combinatorics +using DifferentiationInterface using DistributionsAD using Documenter: Documenter using Enzyme @@ -34,9 +35,16 @@ using LazyArrays: LazyArrays const GROUP = get(ENV, "GROUP", "All") const IS_PRERELEASE = !isempty(VERSION.prerelease) +TEST_ADTYPES = [ + ("ForwardDiff", AutoForwardDiff()), + ("ReverseDiff", AutoReverseDiff(; compile=false)), + ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), +] +if !IS_PRERELEASE + push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) +end # Always include this since it can be useful for other tests. -include("ad/utils.jl") include("bijectors/utils.jl") if GROUP == "All" || GROUP == "Interface" @@ -54,7 +62,22 @@ if GROUP == "All" || GROUP == "Interface" include("bijectors/corr.jl") include("bijectors/product_bijector.jl") include("distributionsad.jl") - +end +if GROUP == "All" || GROUP == "AD" + const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) + function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) + ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) + gradient = DifferentiationInterface.gradient(f, backend, x) + @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) + end + include("ad/chainrules.jl") + include("ad/flows.jl") + include("ad/pd.jl") + include("ad/corr.jl") + include("ad/stacked.jl") + include("ad/enzyme.jl") +end +if GROUP == "All" || GROUP == "Doctests" @testset "doctests" begin Documenter.DocMeta.setdocmeta!( Bijectors, :DocTestSetup, :(using Bijectors); recursive=true @@ -68,20 +91,3 @@ if GROUP == "All" || GROUP == "Interface" end end -TEST_ADTYPES = [ - ("ForwardDiff", AutoForwardDiff()), - ("ReverseDiff", AutoReverseDiff(; compile=false)), - ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), -] -if !IS_PRERELEASE - push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) -end - -if GROUP == "All" || GROUP == "AD" - include("ad/chainrules.jl") - include("ad/flows.jl") - include("ad/pd.jl") - include("ad/corr.jl") - include("ad/stacked.jl") - include("ad/enzyme.jl") -end From 1b9964aeb3fe82c4c1e39c0e832e730d2dcd60cc Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 19:03:24 +0100 Subject: [PATCH 08/28] fix --- docs/make.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/make.jl b/docs/make.jl index 76e93b39..c59fed7f 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -14,6 +14,8 @@ makedocs(; "Distributions.jl integration" => "distributions.md", "Examples" => "examples.md", ], - strict=false, checkdocs=:exports, + doctest=false, + # TODO: Remove this + warnonly=true, ) From 4f1004c7f92f1d938b3e993c5d12d904ed5471a8 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 19:03:39 +0100 Subject: [PATCH 09/28] format --- test/runtests.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 7632bd67..172dab77 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -90,4 +90,3 @@ if GROUP == "All" || GROUP == "Doctests" Documenter.doctest(Bijectors; manual=false, doctestfilters=doctestfilters) end end - From 481c5ca5906cf2ad17006140d53583c06585772b Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 23:10:33 +0100 Subject: [PATCH 10/28] Disable failing Enzyme rule tests --- test/ad/enzyme.jl | 51 ------------------------------ test/ad/enzyme_rules.jl | 70 +++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 62 ++++++++++++++++++------------------ 3 files changed, 101 insertions(+), 82 deletions(-) delete mode 100644 test/ad/enzyme.jl create mode 100644 test/ad/enzyme_rules.jl diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl deleted file mode 100644 index 78fcab05..00000000 --- a/test/ad/enzyme.jl +++ /dev/null @@ -1,51 +0,0 @@ -@testset "Enzyme: Bijectors.find_alpha" begin - x = randn() - y = expm1(randn()) - z = randn() - - @testset "forward" begin - # No batches - @testset for RT in (Const, Duplicated, DuplicatedNoNeed), - Tx in (Const, Duplicated), - Ty in (Const, Duplicated), - Tz in (Const, Duplicated) - - # Rule not picked up by Enzyme on Julia 1.11?! - # Ref https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 - if VERSION >= v"1.11" && Tx <: Const && Ty <: Const && Tz <: Const - continue - end - - test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) - end - - # Batches - @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed), - Tx in (Const, BatchDuplicated), - Ty in (Const, BatchDuplicated), - Tz in (Const, BatchDuplicated) - - # Rule not picked up by Enzyme on Julia 1.11?! - # Ref https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 - if VERSION >= v"1.11" && Tx <: Const && Ty <: Const && Tz <: Const - continue - end - - test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) - end - end - @testset "reverse" begin - # No batches - @testset for RT in (Const, Active), - Tx in (Const, Active), - Ty in (Const, Active), - Tz in (Const, Active) - - test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) - end - - # TODO: Test batch mode - # This is a bit problematic since Enzyme does not support all combinations of activities currently - # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728 - end -end diff --git a/test/ad/enzyme_rules.jl b/test/ad/enzyme_rules.jl new file mode 100644 index 00000000..ff56c641 --- /dev/null +++ b/test/ad/enzyme_rules.jl @@ -0,0 +1,70 @@ +module BijectorsEnzymeRulesTests + +using Bijectors +using Enzyme +using EnzymeTestUtils: test_forward, test_reverse +using Test + +@testset "Enzyme: Bijectors.find_alpha" begin + x = randn() + y = expm1(randn()) + z = randn() + + @testset "forward" begin + # No batches + @testset for RT in (Const, Duplicated, DuplicatedNoNeed), + Tx in (Const, Duplicated), + Ty in (Const, Duplicated), + Tz in (Const, Duplicated) + + if VERSION >= v"1.11" && (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) + # https://github.com/EnzymeAD/Enzyme.jl/issues/2121 + # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 + # + # Ideally we'd use `@test_throws`. However, that doesn't work + # because `test_forward` itself calls `@test`, and the error is + # captured by that `@test`, not our `@test_throws`. + # Consequently `@test_throws` doesn't actually see any error. + # Weird Julia behaviour. + continue + else + test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + end + + # Batches + @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed), + Tx in (Const, BatchDuplicated), + Ty in (Const, BatchDuplicated), + Tz in (Const, BatchDuplicated) + + if VERSION >= v"1.11" && (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) + # See above + continue + else + test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + end + end + @testset "reverse" begin + # No batches + @testset for RT in (Const, Active), + Tx in (Const, Active), + Ty in (Const, Active), + Tz in (Const, Active) + + if VERSION >= v"1.11" + # See above + continue + else + test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + end + + # TODO: Test batch mode + # This is a bit problematic since Enzyme does not support all combinations of activities currently + # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728 + end +end + +end diff --git a/test/runtests.jl b/test/runtests.jl index 172dab77..1029d2ea 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -48,20 +48,20 @@ end include("bijectors/utils.jl") if GROUP == "All" || GROUP == "Interface" - include("interface.jl") - include("transform.jl") - include("norm_flows.jl") - include("bijectors/permute.jl") - include("bijectors/rational_quadratic_spline.jl") - include("bijectors/named_bijector.jl") - include("bijectors/leaky_relu.jl") - include("bijectors/coupling.jl") - include("bijectors/ordered.jl") - include("bijectors/pd.jl") - include("bijectors/reshape.jl") - include("bijectors/corr.jl") - include("bijectors/product_bijector.jl") - include("distributionsad.jl") + # include("interface.jl") + # include("transform.jl") + # include("norm_flows.jl") + # include("bijectors/permute.jl") + # include("bijectors/rational_quadratic_spline.jl") + # include("bijectors/named_bijector.jl") + # include("bijectors/leaky_relu.jl") + # include("bijectors/coupling.jl") + # include("bijectors/ordered.jl") + # include("bijectors/pd.jl") + # include("bijectors/reshape.jl") + # include("bijectors/corr.jl") + # include("bijectors/product_bijector.jl") + # include("distributionsad.jl") end if GROUP == "All" || GROUP == "AD" const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) @@ -70,23 +70,23 @@ if GROUP == "All" || GROUP == "AD" gradient = DifferentiationInterface.gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end - include("ad/chainrules.jl") - include("ad/flows.jl") - include("ad/pd.jl") - include("ad/corr.jl") - include("ad/stacked.jl") - include("ad/enzyme.jl") + # include("ad/chainrules.jl") + # include("ad/flows.jl") + # include("ad/pd.jl") + # include("ad/corr.jl") + # include("ad/stacked.jl") + include("ad/enzyme_rules.jl") end if GROUP == "All" || GROUP == "Doctests" - @testset "doctests" begin - Documenter.DocMeta.setdocmeta!( - Bijectors, :DocTestSetup, :(using Bijectors); recursive=true - ) - doctestfilters = [ - # Ignore the source of a warning in the doctest output, since this is dependent - # on host. This is a line that starts with "└ @ " and ends with the line number. - r"└ @ .+:[0-9]+", - ] - Documenter.doctest(Bijectors; manual=false, doctestfilters=doctestfilters) - end + # @testset "doctests" begin + # Documenter.DocMeta.setdocmeta!( + # Bijectors, :DocTestSetup, :(using Bijectors); recursive=true + # ) + # doctestfilters = [ + # # Ignore the source of a warning in the doctest output, since this is dependent + # # on host. This is a line that starts with "└ @ " and ends with the line number. + # r"└ @ .+:[0-9]+", + # ] + # Documenter.doctest(Bijectors; manual=false, doctestfilters=doctestfilters) + # end end From eb4bdfb77ba24825c9e9c3e1b3d0e68561caeaaa Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 23:28:25 +0100 Subject: [PATCH 11/28] Re-enable Enzyme tests --- test/ad/{enzyme_rules.jl => enzyme.jl} | 0 test/runtests.jl | 69 ++++++++++++++------------ 2 files changed, 36 insertions(+), 33 deletions(-) rename test/ad/{enzyme_rules.jl => enzyme.jl} (100%) diff --git a/test/ad/enzyme_rules.jl b/test/ad/enzyme.jl similarity index 100% rename from test/ad/enzyme_rules.jl rename to test/ad/enzyme.jl diff --git a/test/runtests.jl b/test/runtests.jl index 1029d2ea..edc9d5a7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,8 +5,7 @@ using Combinatorics using DifferentiationInterface using DistributionsAD using Documenter: Documenter -using Enzyme -using EnzymeTestUtils +using Enzyme: set_runtime_activity, Forward, Reverse using FiniteDifferences using ForwardDiff using Functors @@ -42,27 +41,30 @@ TEST_ADTYPES = [ ] if !IS_PRERELEASE push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) + push!(TEST_ADTYPES, ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward)))) + push!(TEST_ADTYPES, ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse)))) end # Always include this since it can be useful for other tests. include("bijectors/utils.jl") if GROUP == "All" || GROUP == "Interface" - # include("interface.jl") - # include("transform.jl") - # include("norm_flows.jl") - # include("bijectors/permute.jl") - # include("bijectors/rational_quadratic_spline.jl") - # include("bijectors/named_bijector.jl") - # include("bijectors/leaky_relu.jl") - # include("bijectors/coupling.jl") - # include("bijectors/ordered.jl") - # include("bijectors/pd.jl") - # include("bijectors/reshape.jl") - # include("bijectors/corr.jl") - # include("bijectors/product_bijector.jl") - # include("distributionsad.jl") + include("interface.jl") + include("transform.jl") + include("norm_flows.jl") + include("bijectors/permute.jl") + include("bijectors/rational_quadratic_spline.jl") + include("bijectors/named_bijector.jl") + include("bijectors/leaky_relu.jl") + include("bijectors/coupling.jl") + include("bijectors/ordered.jl") + include("bijectors/pd.jl") + include("bijectors/reshape.jl") + include("bijectors/corr.jl") + include("bijectors/product_bijector.jl") + include("distributionsad.jl") end + if GROUP == "All" || GROUP == "AD" const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) @@ -70,23 +72,24 @@ if GROUP == "All" || GROUP == "AD" gradient = DifferentiationInterface.gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end - # include("ad/chainrules.jl") - # include("ad/flows.jl") - # include("ad/pd.jl") - # include("ad/corr.jl") - # include("ad/stacked.jl") - include("ad/enzyme_rules.jl") + include("ad/chainrules.jl") + include("ad/flows.jl") + include("ad/pd.jl") + include("ad/corr.jl") + include("ad/stacked.jl") + include("ad/enzyme.jl") end + if GROUP == "All" || GROUP == "Doctests" - # @testset "doctests" begin - # Documenter.DocMeta.setdocmeta!( - # Bijectors, :DocTestSetup, :(using Bijectors); recursive=true - # ) - # doctestfilters = [ - # # Ignore the source of a warning in the doctest output, since this is dependent - # # on host. This is a line that starts with "└ @ " and ends with the line number. - # r"└ @ .+:[0-9]+", - # ] - # Documenter.doctest(Bijectors; manual=false, doctestfilters=doctestfilters) - # end + @testset "doctests" begin + Documenter.DocMeta.setdocmeta!( + Bijectors, :DocTestSetup, :(using Bijectors); recursive=true + ) + doctestfilters = [ + # Ignore the source of a warning in the doctest output, since this is dependent + # on host. This is a line that starts with "└ @ " and ends with the line number. + r"└ @ .+:[0-9]+", + ] + Documenter.doctest(Bijectors; manual=false, doctestfilters=doctestfilters) + end end From 1400d75a4eb0fca508424d3f54a80997ff480120 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 23:28:49 +0100 Subject: [PATCH 12/28] format --- test/ad/enzyme.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl index ff56c641..c6a87863 100644 --- a/test/ad/enzyme.jl +++ b/test/ad/enzyme.jl @@ -17,7 +17,8 @@ using Test Ty in (Const, Duplicated), Tz in (Const, Duplicated) - if VERSION >= v"1.11" && (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) + if VERSION >= v"1.11" && + (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) # https://github.com/EnzymeAD/Enzyme.jl/issues/2121 # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 # @@ -38,7 +39,8 @@ using Test Ty in (Const, BatchDuplicated), Tz in (Const, BatchDuplicated) - if VERSION >= v"1.11" && (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) + if VERSION >= v"1.11" && + (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) # See above continue else From c993dfb23c508e7583d4c51591a3493734478bff Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 14 Aug 2025 23:46:55 +0100 Subject: [PATCH 13/28] separate mooncake rule test into its own thing --- test/ad/chainrules.jl | 47 ++++++++----------------------------------- test/ad/mooncake.jl | 47 +++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 8 +++++--- 3 files changed, 60 insertions(+), 42 deletions(-) create mode 100644 test/ad/mooncake.jl diff --git a/test/ad/chainrules.jl b/test/ad/chainrules.jl index 5c938502..d798d639 100644 --- a/test/ad/chainrules.jl +++ b/test/ad/chainrules.jl @@ -1,6 +1,10 @@ -using Random: Xoshiro +module BijectorsChainRulesTests + +using Bijectors using LinearAlgebra -using ChainRulesTestUtils: ChainRulesCore +using ChainRulesTestUtils: ChainRulesCore, test_frule, test_rrule, ⊢, ChainRulesTestUtils +using Random: Xoshiro +using Test # HACK: This is a workaround to test `Bijectors._inv_link_chol_lkj` which produces an # upper-triangular `Matrix`, leading to `test_rrule` comaring the _full_ `Matrix`, @@ -29,43 +33,6 @@ end test_frule(Bijectors.find_alpha, x, y, z) test_rrule(Bijectors.find_alpha, x, y, z) - if @isdefined Mooncake - rng = Xoshiro(123456) - # TODO: Enable Mooncake.ForwardMode as well. - @testset "$mode" for mode in (Mooncake.ReverseMode,) - Mooncake.TestUtils.test_rule( - rng, - Bijectors.find_alpha, - x, - y, - z; - is_primitive=true, - perf_flag=:none, - mode=mode, - ) - Mooncake.TestUtils.test_rule( - rng, - Bijectors.find_alpha, - x, - y, - 3; - is_primitive=true, - perf_flag=:none, - mode=mode, - ) - Mooncake.TestUtils.test_rule( - rng, - Bijectors.find_alpha, - x, - y, - UInt32(3); - is_primitive=true, - perf_flag=:none, - mode=mode, - ) - end - end - test_rrule( Bijectors.combine, Bijectors.PartitionMask(3, [1], [2]) ⊢ ChainRulesTestUtils.NoTangent(), @@ -183,3 +150,5 @@ end end end end + +end # module BijectorsChainRulesTests diff --git a/test/ad/mooncake.jl b/test/ad/mooncake.jl new file mode 100644 index 00000000..2431e520 --- /dev/null +++ b/test/ad/mooncake.jl @@ -0,0 +1,47 @@ +module BijectorsMooncakeTests + +using Bijectors: Bijectors +using Mooncake: Mooncake +using Random: Xoshiro +using Test + +x = randn() +y = expm1(randn()) +z = randn() +rng = Xoshiro(123456) + +# TODO: Enable Mooncake.ForwardMode as well. +@testset "Mooncake $mode: find_alpha" for mode in (Mooncake.ReverseMode,) + Mooncake.TestUtils.test_rule( + rng, + Bijectors.find_alpha, + x, + y, + z; + is_primitive=true, + perf_flag=:none, + mode=mode, + ) + Mooncake.TestUtils.test_rule( + rng, + Bijectors.find_alpha, + x, + y, + 3; + is_primitive=true, + perf_flag=:none, + mode=mode, + ) + Mooncake.TestUtils.test_rule( + rng, + Bijectors.find_alpha, + x, + y, + UInt32(3); + is_primitive=true, + perf_flag=:none, + mode=mode, + ) +end + +end # module diff --git a/test/runtests.jl b/test/runtests.jl index edc9d5a7..aadb4f26 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -40,9 +40,10 @@ TEST_ADTYPES = [ ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), ] if !IS_PRERELEASE - push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) - push!(TEST_ADTYPES, ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward)))) - push!(TEST_ADTYPES, ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse)))) + println("ok") + # push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) + # push!(TEST_ADTYPES, ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward)))) + # push!(TEST_ADTYPES, ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse)))) end # Always include this since it can be useful for other tests. @@ -78,6 +79,7 @@ if GROUP == "All" || GROUP == "AD" include("ad/corr.jl") include("ad/stacked.jl") include("ad/enzyme.jl") + include("ad/mooncake.jl") end if GROUP == "All" || GROUP == "Doctests" From ef21f50cbc6bba86d93106ab7ac4bab08bf32975 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 15 Aug 2025 01:07:40 +0100 Subject: [PATCH 14/28] Skip broken Enzyme tests --- test/ad/chainrules.jl | 1 + test/ad/corr.jl | 37 +++++++++++++++++++++++++------------ test/ad/flows.jl | 26 ++++++++++++++++++++++---- test/ad/pd.jl | 3 --- test/runtests.jl | 25 +++++++++++++++++-------- 5 files changed, 65 insertions(+), 27 deletions(-) diff --git a/test/ad/chainrules.jl b/test/ad/chainrules.jl index d798d639..ffc453db 100644 --- a/test/ad/chainrules.jl +++ b/test/ad/chainrules.jl @@ -3,6 +3,7 @@ module BijectorsChainRulesTests using Bijectors using LinearAlgebra using ChainRulesTestUtils: ChainRulesCore, test_frule, test_rrule, ⊢, ChainRulesTestUtils +using FiniteDifferences: FiniteDifferences using Random: Xoshiro using Test diff --git a/test/ad/corr.jl b/test/ad/corr.jl index 49e35c0c..f498ebc1 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -1,3 +1,5 @@ +using Enzyme: ForwardMode + @testset "VecCorrBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES @testset "d = $d" for d in (1, 2, 4) dist = LKJ(d, 2.0) @@ -7,10 +9,17 @@ x = rand(dist) y = b(x) - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse only - test_ad(y -> sum(transform(binv, y)), adtype, y) + if adtype isa AutoEnzyme{<:ForwardMode} && d == 1 + # For d == 1, y has length 0, and DI doesn't handle this well + # https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/802 + @test_throws DivideError test_ad(y -> sum(transform(b, binv(y))), adtype, y) + @test_throws DivideError test_ad(y -> sum(transform(binv, y)), adtype, y) + else + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), adtype, y) + # inverse only + test_ad(y -> sum(transform(binv, y)), adtype, y) + end end end @@ -22,16 +31,20 @@ end x = rand(dist) y = b(x) + cholesky_to_triangular = uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, - # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't - # give a scalar) - if uplo == 'U' - test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), adtype, y) + if adtype isa AutoEnzyme{<:ForwardMode} && d == 1 + # For d == 1, y has length 0, and DI doesn't handle this well + # https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/802 + @test_throws DivideError test_ad(y -> sum(transform(b, binv(y))), adtype, y) + @test_throws DivideError test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) else - test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), adtype, y) + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), adtype, y) + # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, + # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't + # give a scalar) + test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) end end end diff --git a/test/ad/flows.jl b/test/ad/flows.jl index 34ad90aa..f984e72b 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -1,3 +1,5 @@ +using Enzyme: ForwardMode + @testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES # logpdf of a flow with a planar layer and two-dimensional inputs function f(θ) @@ -6,7 +8,11 @@ x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - test_ad(f, adtype, randn(7)) + if adtype isa AutoEnzyme{<:ForwardMode} + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) + else + test_ad(f, adtype, randn(7)) + end function g(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) @@ -14,7 +20,11 @@ x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end - test_ad(g, adtype, randn(11)) + if adtype isa AutoEnzyme{<:ForwardMode} + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) + else + test_ad(g, adtype, randn(11)) + end # logpdf of a flow with the inverse of a planar layer and two-dimensional inputs function finv(θ) @@ -23,7 +33,11 @@ x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - test_ad(finv, adtype, randn(7)) + if adtype isa AutoEnzyme{<:ForwardMode} + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) + else + test_ad(f, adtype, randn(7)) + end function ginv(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) @@ -31,5 +45,9 @@ x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end - test_ad(ginv, adtype, randn(11)) + if adtype isa AutoEnzyme{<:ForwardMode} + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) + else + test_ad(g, adtype, randn(11)) + end end diff --git a/test/ad/pd.jl b/test/ad/pd.jl index daa244c3..acfe42bd 100644 --- a/test/ad/pd.jl +++ b/test/ad/pd.jl @@ -11,9 +11,6 @@ _topd(x) = x * x' + I test_ad(x -> sum(transform(b, _topd(reshape(x, d, d)))), adtype, vec(z)) test_ad(y -> sum(transform(binv, y)), adtype, y) - - # if occursin("ReverseDiff", backend_name) test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), adtype, y) test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), adtype, y) - # end end diff --git a/test/runtests.jl b/test/runtests.jl index aadb4f26..697a8e22 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,7 +5,7 @@ using Combinatorics using DifferentiationInterface using DistributionsAD using Documenter: Documenter -using Enzyme: set_runtime_activity, Forward, Reverse +using Enzyme: Enzyme, set_runtime_activity, Forward, Reverse, Const using FiniteDifferences using ForwardDiff using Functors @@ -40,10 +40,13 @@ TEST_ADTYPES = [ ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), ] if !IS_PRERELEASE - println("ok") - # push!(TEST_ADTYPES, ("Mooncake", AutoMooncake())) - # push!(TEST_ADTYPES, ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward)))) - # push!(TEST_ADTYPES, ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse)))) + # Mooncake and Enzyme tend to be unstable on prerelease, so only + # run these on stable Julia releases + TEST_ADTYPES = [TEST_ADTYPES..., + ("Mooncake", AutoMooncake()), + ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const)), + ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const)), + ] end # Always include this since it can be useful for other tests. @@ -67,19 +70,25 @@ if GROUP == "All" || GROUP == "Interface" end if GROUP == "All" || GROUP == "AD" + # These tests specifically check the implementation of AD backend rules. + include("ad/chainrules.jl") + if !IS_PRERELEASE + include("ad/enzyme.jl") + include("ad/mooncake.jl") + end + + # These tests check that AD can differentiate through Bijectors + # functionality without explicit rules. const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) gradient = DifferentiationInterface.gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) end - include("ad/chainrules.jl") include("ad/flows.jl") include("ad/pd.jl") include("ad/corr.jl") include("ad/stacked.jl") - include("ad/enzyme.jl") - include("ad/mooncake.jl") end if GROUP == "All" || GROUP == "Doctests" From e81001e0817d21451712037c499ae9bced30d94d Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 15 Aug 2025 01:09:06 +0100 Subject: [PATCH 15/28] add macOS and windows test --- .github/workflows/CI.yml | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index b690de15..06eb07ee 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -14,15 +14,19 @@ concurrency: jobs: test: - runs-on: ${{ matrix.os }} + runs-on: ${{ matrix.runner.os }} strategy: fail-fast: false matrix: - version: - - 'min' - - '1' - os: - - ubuntu-latest + runner: + - version: '1' + os: 'ubuntu-latest' + - version: '1' + os: 'macos-latest' + - version: '1' + os: 'windows-latest' + - version: 'min' + os: 'ubuntu-latest' group: - 'Interface' - 'AD' @@ -31,7 +35,7 @@ jobs: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 with: - version: ${{ matrix.version }} + version: ${{ matrix.runner.version }} - uses: julia-actions/cache@v2 - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 From d67f0c057ac130e13899088e63c82f838e5fdd82 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 15 Aug 2025 01:09:56 +0100 Subject: [PATCH 16/28] format --- test/ad/corr.jl | 7 +++++-- test/ad/mooncake.jl | 18 ++---------------- test/runtests.jl | 13 ++++++++++--- 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/test/ad/corr.jl b/test/ad/corr.jl index f498ebc1..88937b0b 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -31,13 +31,16 @@ end x = rand(dist) y = b(x) - cholesky_to_triangular = uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower + cholesky_to_triangular = + uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower if adtype isa AutoEnzyme{<:ForwardMode} && d == 1 # For d == 1, y has length 0, and DI doesn't handle this well # https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/802 @test_throws DivideError test_ad(y -> sum(transform(b, binv(y))), adtype, y) - @test_throws DivideError test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) + @test_throws DivideError test_ad( + y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y + ) else # roundtrip test_ad(y -> sum(transform(b, binv(y))), adtype, y) diff --git a/test/ad/mooncake.jl b/test/ad/mooncake.jl index 2431e520..1e9259be 100644 --- a/test/ad/mooncake.jl +++ b/test/ad/mooncake.jl @@ -13,24 +13,10 @@ rng = Xoshiro(123456) # TODO: Enable Mooncake.ForwardMode as well. @testset "Mooncake $mode: find_alpha" for mode in (Mooncake.ReverseMode,) Mooncake.TestUtils.test_rule( - rng, - Bijectors.find_alpha, - x, - y, - z; - is_primitive=true, - perf_flag=:none, - mode=mode, + rng, Bijectors.find_alpha, x, y, z; is_primitive=true, perf_flag=:none, mode=mode ) Mooncake.TestUtils.test_rule( - rng, - Bijectors.find_alpha, - x, - y, - 3; - is_primitive=true, - perf_flag=:none, - mode=mode, + rng, Bijectors.find_alpha, x, y, 3; is_primitive=true, perf_flag=:none, mode=mode ) Mooncake.TestUtils.test_rule( rng, diff --git a/test/runtests.jl b/test/runtests.jl index 697a8e22..4e4541f1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,10 +42,17 @@ TEST_ADTYPES = [ if !IS_PRERELEASE # Mooncake and Enzyme tend to be unstable on prerelease, so only # run these on stable Julia releases - TEST_ADTYPES = [TEST_ADTYPES..., + TEST_ADTYPES = [ + TEST_ADTYPES..., ("Mooncake", AutoMooncake()), - ("EnzymeForward", AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const)), - ("EnzymeReverse", AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const)), + ( + "EnzymeForward", + AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), + ), + ( + "EnzymeReverse", + AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), + ), ] end From 28d0ca8656729e1b747570ad56d8de3f8b513e2f Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 15 Aug 2025 01:24:11 +0100 Subject: [PATCH 17/28] flows only fail on 1.11 apparently --- test/ad/flows.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/test/ad/flows.jl b/test/ad/flows.jl index f984e72b..117d025f 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -1,6 +1,8 @@ using Enzyme: ForwardMode @testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES + ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:ForwardMode} + # logpdf of a flow with a planar layer and two-dimensional inputs function f(θ) layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5]) @@ -8,7 +10,7 @@ using Enzyme: ForwardMode x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - if adtype isa AutoEnzyme{<:ForwardMode} + if ENZYME_FWD_AND_1p11 @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) else test_ad(f, adtype, randn(7)) @@ -20,7 +22,7 @@ using Enzyme: ForwardMode x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end - if adtype isa AutoEnzyme{<:ForwardMode} + if ENZYME_FWD_AND_1p11 @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) else test_ad(g, adtype, randn(11)) @@ -33,7 +35,7 @@ using Enzyme: ForwardMode x = θ[6:7] return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end - if adtype isa AutoEnzyme{<:ForwardMode} + if ENZYME_FWD_AND_1p11 @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) else test_ad(f, adtype, randn(7)) @@ -45,7 +47,7 @@ using Enzyme: ForwardMode x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end - if adtype isa AutoEnzyme{<:ForwardMode} + if ENZYME_FWD_AND_1p11 @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) else test_ad(g, adtype, randn(11)) From a7852b6f268f8fa5945049be5fa555f9513e8eac Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 15 Aug 2025 01:24:42 +0100 Subject: [PATCH 18/28] comment --- test/ad/flows.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/ad/flows.jl b/test/ad/flows.jl index 117d025f..1c565570 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -11,6 +11,7 @@ using Enzyme: ForwardMode return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end if ENZYME_FWD_AND_1p11 + # TODO: Report this upstream (or check if it's already been reported) @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) else test_ad(f, adtype, randn(7)) From a86b7160dd317b4c1b7e0c5608cd5cd25534e89a Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 22 Aug 2025 10:19:11 +0100 Subject: [PATCH 19/28] Use new DI patch --- test/Project.toml | 1 + test/ad/corr.jl | 21 ++++++--------------- 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 979b70a8..e1df2f40 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -33,6 +33,7 @@ AdvancedHMC = "0.6, 0.7, 0.8" ChainRulesTestUtils = "0.7, 1" ChangesOfVariables = "0.1" Combinatorics = "1.0.2" +DifferentiationInterface = "0.7.7" DistributionsAD = "0.6.3" Documenter = "1" Enzyme = "0.13.12" diff --git a/test/ad/corr.jl b/test/ad/corr.jl index 88937b0b..d80962e6 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -34,20 +34,11 @@ end cholesky_to_triangular = uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower - if adtype isa AutoEnzyme{<:ForwardMode} && d == 1 - # For d == 1, y has length 0, and DI doesn't handle this well - # https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/802 - @test_throws DivideError test_ad(y -> sum(transform(b, binv(y))), adtype, y) - @test_throws DivideError test_ad( - y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y - ) - else - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, - # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't - # give a scalar) - test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) - end + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), adtype, y) + # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, + # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't + # give a scalar) + test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) end end From ce870f760e40b01e42325b3981a094c47c8eb406 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 22 Aug 2025 11:57:18 +0100 Subject: [PATCH 20/28] Fix Enzyme tests again --- test/ad/corr.jl | 15 +++------ test/ad/enzyme.jl | 86 ++++++++++++++++++++--------------------------- 2 files changed, 41 insertions(+), 60 deletions(-) diff --git a/test/ad/corr.jl b/test/ad/corr.jl index d80962e6..7e69e219 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -9,17 +9,10 @@ using Enzyme: ForwardMode x = rand(dist) y = b(x) - if adtype isa AutoEnzyme{<:ForwardMode} && d == 1 - # For d == 1, y has length 0, and DI doesn't handle this well - # https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/802 - @test_throws DivideError test_ad(y -> sum(transform(b, binv(y))), adtype, y) - @test_throws DivideError test_ad(y -> sum(transform(binv, y)), adtype, y) - else - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse only - test_ad(y -> sum(transform(binv, y)), adtype, y) - end + # roundtrip + test_ad(y -> sum(transform(b, binv(y))), adtype, y) + # inverse only + test_ad(y -> sum(transform(binv, y)), adtype, y) end end diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl index c6a87863..f4d350dc 100644 --- a/test/ad/enzyme.jl +++ b/test/ad/enzyme.jl @@ -5,67 +5,55 @@ using Enzyme using EnzymeTestUtils: test_forward, test_reverse using Test -@testset "Enzyme: Bijectors.find_alpha" begin - x = randn() - y = expm1(randn()) - z = randn() +# This entire test suite is broken on 1.11. +# +# https://github.com/EnzymeAD/Enzyme.jl/issues/2121 +# https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 +# +# Ideally we'd use `@test_throws`. However, that doesn't work because +# `test_forward` itself calls `@test`, and the error is captured by that +# `@test`, not our `@test_throws`. Consequently `@test_throws` doesn't actually +# see any error. Weird Julia behaviour. + +@static if VERSION < v"1.11" + @testset "Enzyme: Bijectors.find_alpha" begin + x = randn() + y = expm1(randn()) + z = randn() + + @testset "forward" begin + # No batches + @testset for RT in (Const, Duplicated, DuplicatedNoNeed), + Tx in (Const, Duplicated), + Ty in (Const, Duplicated), + Tz in (Const, Duplicated) - @testset "forward" begin - # No batches - @testset for RT in (Const, Duplicated, DuplicatedNoNeed), - Tx in (Const, Duplicated), - Ty in (Const, Duplicated), - Tz in (Const, Duplicated) - - if VERSION >= v"1.11" && - (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) - # https://github.com/EnzymeAD/Enzyme.jl/issues/2121 - # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2470766968 - # - # Ideally we'd use `@test_throws`. However, that doesn't work - # because `test_forward` itself calls `@test`, and the error is - # captured by that `@test`, not our `@test_throws`. - # Consequently `@test_throws` doesn't actually see any error. - # Weird Julia behaviour. - continue - else test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) end - end - # Batches - @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed), - Tx in (Const, BatchDuplicated), - Ty in (Const, BatchDuplicated), - Tz in (Const, BatchDuplicated) + # Batches + @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed), + Tx in (Const, BatchDuplicated), + Ty in (Const, BatchDuplicated), + Tz in (Const, BatchDuplicated) - if VERSION >= v"1.11" && - (!(RT <: Const) || (Tx <: Const && Ty <: Const && Tz <: Const)) - # See above - continue - else test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) end end - end - @testset "reverse" begin - # No batches - @testset for RT in (Const, Active), - Tx in (Const, Active), - Ty in (Const, Active), - Tz in (Const, Active) + @testset "reverse" begin + # No batches + @testset for RT in (Const, Active), + Tx in (Const, Active), + Ty in (Const, Active), + Tz in (Const, Active) - if VERSION >= v"1.11" - # See above - continue - else test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) end - end - # TODO: Test batch mode - # This is a bit problematic since Enzyme does not support all combinations of activities currently - # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728 + # TODO: Test batch mode + # This is a bit problematic since Enzyme does not support all combinations of activities currently + # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728 + end end end From 1d8e22470650816595d3d4ba5df94a9534936fd3 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Thu, 11 Sep 2025 10:18:34 +0100 Subject: [PATCH 21/28] fix merge --- .github/workflows/Interface.yml | 36 --------------------------------- 1 file changed, 36 deletions(-) delete mode 100644 .github/workflows/Interface.yml diff --git a/.github/workflows/Interface.yml b/.github/workflows/Interface.yml deleted file mode 100644 index e971e538..00000000 --- a/.github/workflows/Interface.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: Interface tests - -on: - push: - branches: - - main - pull_request: - -concurrency: - # Skip intermediate builds: always. - # Cancel intermediate builds: only if it is a pull request build. - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} - -jobs: - test: - runs-on: ${{ matrix.os }} - strategy: - fail-fast: false - matrix: - version: - - 'min' - - 'lts' - - '1' - os: - - ubuntu-latest - - macOS-latest - steps: - - uses: actions/checkout@v5 - - uses: julia-actions/setup-julia@v2 - with: - version: ${{ matrix.version }} - - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-runtest@v1 - env: - GROUP: Interface From 9971c97798afd0c0b92ad175c7792ba791c5d197 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 12 Sep 2025 17:18:18 +0100 Subject: [PATCH 22/28] Add link to issue --- test/ad/flows.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/ad/flows.jl b/test/ad/flows.jl index 1c565570..5cccdb55 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -1,6 +1,7 @@ using Enzyme: ForwardMode @testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES + # https://github.com/TuringLang/Bijectors.jl/issues/415 ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:ForwardMode} # logpdf of a flow with a planar layer and two-dimensional inputs @@ -11,7 +12,6 @@ using Enzyme: ForwardMode return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end if ENZYME_FWD_AND_1p11 - # TODO: Report this upstream (or check if it's already been reported) @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) else test_ad(f, adtype, randn(7)) From af877e89bfc8006b02d03414f4267353b9cd964e Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Mon, 15 Sep 2025 16:40:08 +0100 Subject: [PATCH 23/28] Remove warnonly=true from docs build --- docs/make.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index c59fed7f..80ad43bf 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -16,6 +16,4 @@ makedocs(; ], checkdocs=:exports, doctest=false, - # TODO: Remove this - warnonly=true, ) From 4b2cbbafba0916e77f5eaf549457bf8284a9808f Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Mon, 15 Sep 2025 16:47:22 +0100 Subject: [PATCH 24/28] Fix wrong function in test --- test/ad/flows.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/ad/flows.jl b/test/ad/flows.jl index 5cccdb55..069dbeed 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -37,9 +37,9 @@ using Enzyme: ForwardMode return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end if ENZYME_FWD_AND_1p11 - @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(f, adtype, randn(7)) + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(finv, adtype, randn(7)) else - test_ad(f, adtype, randn(7)) + test_ad(finv, adtype, randn(7)) end function ginv(θ) @@ -49,8 +49,8 @@ using Enzyme: ForwardMode return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end if ENZYME_FWD_AND_1p11 - @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) + @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(ginv, adtype, randn(11)) else - test_ad(g, adtype, randn(11)) + test_ad(ginv, adtype, randn(11)) end end From fac6c00bd29052b0944ad3d37cdb9ee928e5de9a Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 16 Sep 2025 13:23:01 +0100 Subject: [PATCH 25/28] update Enzyme failures --- test/ad/flows.jl | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/test/ad/flows.jl b/test/ad/flows.jl index 069dbeed..8cfd4cea 100644 --- a/test/ad/flows.jl +++ b/test/ad/flows.jl @@ -1,8 +1,9 @@ -using Enzyme: ForwardMode +using Enzyme: Enzyme @testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES # https://github.com/TuringLang/Bijectors.jl/issues/415 - ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:ForwardMode} + ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:Enzyme.ForwardMode} + ENZYME_RVS_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:Enzyme.ReverseMode} # logpdf of a flow with a planar layer and two-dimensional inputs function f(θ) @@ -24,7 +25,8 @@ using Enzyme: ForwardMode return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end if ENZYME_FWD_AND_1p11 - @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) + @warn "Skipping forward-mode Enzyme for `g` on 1.11 due to segfault" + # @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(g, adtype, randn(11)) else test_ad(g, adtype, randn(11)) end @@ -37,7 +39,9 @@ using Enzyme: ForwardMode return logpdf(flow.dist, x) - logabsdetjac(flow.transform, x) end if ENZYME_FWD_AND_1p11 - @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(finv, adtype, randn(7)) + @warn "Skipping forward-mode Enzyme for `finv` on 1.11 due to segfault" + elseif ENZYME_RVS_AND_1p11 + @test_throws Enzyme.LLVM.LLVMException test_ad(finv, adtype, randn(7)) else test_ad(finv, adtype, randn(7)) end @@ -48,8 +52,10 @@ using Enzyme: ForwardMode x = reshape(θ[6:end], 2, :) return sum(logpdf(flow.dist, x) - logabsdetjac(flow.transform, x)) end - if ENZYME_FWD_AND_1p11 - @test_throws Enzyme.Compiler.EnzymeInternalError test_ad(ginv, adtype, randn(11)) + if ENZYME_FWD_AND_1p11 || ENZYME_RVS_AND_1p11 + @warn "Skipping forward-mode Enzyme for `ginv` on 1.11 due to segfault" + elseif ENZYME_RVS_AND_1p11 + @test_throws Enzyme.LLVM.LLVMException test_ad(finv, adtype, randn(7)) else test_ad(ginv, adtype, randn(11)) end From 7382f7b74b05299947b025b3c6fcbc2456b3a8f4 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Tue, 16 Sep 2025 14:07:59 +0100 Subject: [PATCH 26/28] fix more stuff --- test/ad/corr.jl | 32 ++++++++++++++++++++++---------- test/ad/pd.jl | 30 ++++++++++++++++++++++++++---- test/runtests.jl | 1 + 3 files changed, 49 insertions(+), 14 deletions(-) diff --git a/test/ad/corr.jl b/test/ad/corr.jl index 7e69e219..7a8107b7 100644 --- a/test/ad/corr.jl +++ b/test/ad/corr.jl @@ -1,6 +1,8 @@ using Enzyme: ForwardMode @testset "VecCorrBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES + ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:Enzyme.ForwardMode} + @testset "d = $d" for d in (1, 2, 4) dist = LKJ(d, 2.0) b = bijector(dist) @@ -9,10 +11,19 @@ using Enzyme: ForwardMode x = rand(dist) y = b(x) - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse only - test_ad(y -> sum(transform(binv, y)), adtype, y) + roundtrip(y) = sum(transform(b, binv(y))) + inverse_only(y) = sum(transform(binv, y)) + if d == 4 && ENZYME_FWD_AND_1p11 + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + roundtrip, adtype, y + ) + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + inverse_only, adtype, y + ) + else + test_ad(roundtrip, adtype, y) + test_ad(inverse_only, adtype, y) + end end end @@ -27,11 +38,12 @@ end cholesky_to_triangular = uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower - # roundtrip - test_ad(y -> sum(transform(b, binv(y))), adtype, y) - # inverse (we need to tack on `cholesky_upper`/`cholesky_lower`, - # because directly calling `sum` on a LinearAlgebra.Cholesky doesn't - # give a scalar) - test_ad(y -> sum(cholesky_to_triangular(transform(binv, y))), adtype, y) + roundtrip(y) = sum(transform(b, binv(y))) + test_ad(roundtrip, adtype, y) + + # we need to tack on `cholesky_upper`/`cholesky_lower`, because directly calling + # `sum` on a LinearAlgebra.Cholesky doesn't give a scalar + inverse_only(y) = sum(cholesky_to_triangular(transform(binv, y))) + test_ad(inverse_only, adtype, y) end end diff --git a/test/ad/pd.jl b/test/ad/pd.jl index acfe42bd..923d9ac0 100644 --- a/test/ad/pd.jl +++ b/test/ad/pd.jl @@ -1,6 +1,8 @@ _topd(x) = x * x' + I @testset "PDVecBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES + ENZYME_FWD_AND_1p11 = VERSION >= v"1.11" && adtype isa AutoEnzyme{<:Enzyme.ForwardMode} + d = 4 b = Bijectors.PDVecBijector() binv = inverse(b) @@ -9,8 +11,28 @@ _topd(x) = x * x' + I x = _topd(z) y = b(x) - test_ad(x -> sum(transform(b, _topd(reshape(x, d, d)))), adtype, vec(z)) - test_ad(y -> sum(transform(binv, y)), adtype, y) - test_ad(y -> sum(Bijectors.cholesky_lower(transform(binv, y))), adtype, y) - test_ad(y -> sum(Bijectors.cholesky_upper(transform(binv, y))), adtype, y) + forward_only(x) = sum(transform(b, _topd(reshape(x, d, d)))) + inverse_only(y) = sum(transform(binv, y)) + inverse_chol_lower(y) = sum(Bijectors.cholesky_lower(transform(binv, y))) + inverse_chol_upper(y) = sum(Bijectors.cholesky_upper(transform(binv, y))) + + if ENZYME_FWD_AND_1p11 + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + forward_only, adtype, vec(z) + ) + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + inverse_only, adtype, vec(z) + ) + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + inverse_chol_lower, adtype, y + ) + @test_throws Enzyme.Compiler.EnzymeNoDerivativeError test_ad( + inverse_chol_upper, adtype, y + ) + else + test_ad(forward_only, adtype, vec(z)) + test_ad(inverse_only, adtype, vec(z)) + test_ad(inverse_chol_lower, adtype, y) + test_ad(inverse_chol_upper, adtype, y) + end end diff --git a/test/runtests.jl b/test/runtests.jl index 55a15d48..da6a19dd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -87,6 +87,7 @@ if GROUP == "All" || GROUP == "AD" # functionality without explicit rules. const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) function test_ad(f, backend, x; rtol=1e-6, atol=1e-6) + @info "testing AD for function $f with $backend" ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x) gradient = DifferentiationInterface.gradient(f, backend, x) @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol) From ebe3de3f34559bc67fd1ee0023d5ac28c28015fb Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 26 Sep 2025 18:35:04 +0100 Subject: [PATCH 27/28] That should fix ReverseDiff --- test/ad/pd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/ad/pd.jl b/test/ad/pd.jl index 923d9ac0..b2e8b707 100644 --- a/test/ad/pd.jl +++ b/test/ad/pd.jl @@ -31,7 +31,7 @@ _topd(x) = x * x' + I ) else test_ad(forward_only, adtype, vec(z)) - test_ad(inverse_only, adtype, vec(z)) + test_ad(inverse_only, adtype, z) test_ad(inverse_chol_lower, adtype, y) test_ad(inverse_chol_upper, adtype, y) end From 99cf5863a050380f14826c5269b5de27090a83f2 Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Fri, 26 Sep 2025 18:35:40 +0100 Subject: [PATCH 28/28] I meant this --- test/ad/pd.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/ad/pd.jl b/test/ad/pd.jl index b2e8b707..6c84ab34 100644 --- a/test/ad/pd.jl +++ b/test/ad/pd.jl @@ -31,7 +31,7 @@ _topd(x) = x * x' + I ) else test_ad(forward_only, adtype, vec(z)) - test_ad(inverse_only, adtype, z) + test_ad(inverse_only, adtype, y) test_ad(inverse_chol_lower, adtype, y) test_ad(inverse_chol_upper, adtype, y) end