diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 54d04837b..1c05480d6 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -6,6 +6,11 @@ on: - main pull_request: +# needed to allow julia-actions/cache to delete old caches that it has created +permissions: + actions: write + contents: read + concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. @@ -14,6 +19,7 @@ concurrency: jobs: test: + name: ${{ matrix.group }} / Julia ${{ matrix.runner.version }} / ${{ matrix.runner.os }} runs-on: ${{ matrix.runner.os }} strategy: fail-fast: false @@ -25,8 +31,6 @@ jobs: os: 'macos-latest' - version: '1' os: 'windows-latest' - - version: '1.11' - os: 'ubuntu-latest' - version: 'min' os: 'ubuntu-latest' group: diff --git a/.github/workflows/DocTests.yml b/.github/workflows/DocTests.yml deleted file mode 100644 index ca09ec819..000000000 --- a/.github/workflows/DocTests.yml +++ /dev/null @@ -1,41 +0,0 @@ -# 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@v6 - - - uses: julia-actions/setup-julia@v3 - with: - version: '1' - - - uses: julia-actions/cache@v3 - - - uses: julia-actions/julia-buildpkg@v1 - - - uses: julia-actions/julia-runtest@v1 - env: - GROUP: Doctests diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml index 8d67f381a..926e15a0f 100644 --- a/.github/workflows/Enzyme.yml +++ b/.github/workflows/Enzyme.yml @@ -1,4 +1,4 @@ -name: Enzyme AD tests +name: Enzyme on: push: @@ -17,19 +17,33 @@ concurrency: cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} jobs: - enzyme: + test: + name: Julia ${{ matrix.version }} runs-on: ubuntu-latest + # Enzyme is an integration test isolated from the main suite. It reports failures for + # visibility but does not block PRs: Enzyme regressions surface as red here without + # making the rest of the workflow fail. + continue-on-error: true + strategy: + fail-fast: false + matrix: + version: + - "min" + - "1" steps: - uses: actions/checkout@v6 - uses: julia-actions/setup-julia@v3 with: - version: "1.11" + version: ${{ matrix.version }} - uses: julia-actions/cache@v3 - name: Run Enzyme tests - working-directory: test/integration/enzyme + working-directory: test/integration_tests/enzyme + # `Pkg.develop` makes Julia 1.10 use the working tree's Bijectors. The + # `[sources]` section in `Project.toml` only takes effect on Julia 1.11+, + # so on 1.10 the registry version would be resolved instead. run: | - julia --project=. --color=yes -e 'using Pkg; Pkg.instantiate()' + julia --project=. --color=yes -e 'using Pkg; Pkg.develop(PackageSpec(path="../../..")); Pkg.instantiate()' julia --project=. --color=yes main.jl diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 259e056c7..e9c0966ff 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -15,7 +15,8 @@ concurrency: cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} jobs: - format: + check: + name: runic runs-on: ubuntu-latest steps: diff --git a/.github/workflows/Mooncake.yml b/.github/workflows/Mooncake.yml new file mode 100644 index 000000000..ec0c38115 --- /dev/null +++ b/.github/workflows/Mooncake.yml @@ -0,0 +1,45 @@ +name: Mooncake + +on: + push: + branches: + - main + pull_request: + +# 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: + name: Julia ${{ matrix.version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "min" + - "1" + steps: + - uses: actions/checkout@v6 + + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + + - uses: julia-actions/cache@v3 + + - name: Run Mooncake tests + working-directory: test/integration_tests/mooncake + # `Pkg.develop` makes Julia 1.10 use the working tree's Bijectors. The + # `[sources]` section in `Project.toml` only takes effect on Julia 1.11+, + # so on 1.10 the registry version would be resolved instead. + run: | + julia --project=. --color=yes -e 'using Pkg; Pkg.develop(PackageSpec(path="../../..")); Pkg.instantiate()' + julia --project=. --color=yes main.jl diff --git a/.github/workflows/ReverseDiff.yml b/.github/workflows/ReverseDiff.yml new file mode 100644 index 000000000..b44ad5454 --- /dev/null +++ b/.github/workflows/ReverseDiff.yml @@ -0,0 +1,45 @@ +name: ReverseDiff + +on: + push: + branches: + - main + pull_request: + +# 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: + name: Julia ${{ matrix.version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "min" + - "1" + steps: + - uses: actions/checkout@v6 + + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + + - uses: julia-actions/cache@v3 + + - name: Run ReverseDiff tests + working-directory: test/integration_tests/reversediff + # `Pkg.develop` makes Julia 1.10 use the working tree's Bijectors. The + # `[sources]` section in `Project.toml` only takes effect on Julia 1.11+, + # so on 1.10 the registry version would be resolved instead. + run: | + julia --project=. --color=yes -e 'using Pkg; Pkg.develop(PackageSpec(path="../../..")); Pkg.instantiate()' + julia --project=. --color=yes main.jl diff --git a/Project.toml b/Project.toml index b658fdebf..ea949b413 100644 --- a/Project.toml +++ b/Project.toml @@ -27,14 +27,12 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [weakdeps] ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" -DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] -BijectorsDistributionsADExt = "DistributionsAD" BijectorsForwardDiffExt = "ForwardDiff" BijectorsLazyArraysExt = "LazyArrays" BijectorsMooncakeExt = "Mooncake" @@ -49,7 +47,6 @@ ChainRulesCore = "0.10.11, 1" ChangesOfVariables = "0.1" DifferentiationInterface = "0.7.14" Distributions = "0.25.33" -DistributionsAD = "0.6" DocStringExtensions = "0.9" EnzymeCore = "0.8.15" FillArrays = "1" @@ -69,7 +66,6 @@ Test = "1" julia = "1.10.8" [extras] -DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" diff --git a/docs/make.jl b/docs/make.jl index ab5f2b1c8..781902b1d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -20,5 +20,10 @@ makedocs(; "vector.md", ], checkdocs=:exports, - doctest=false, + doctest=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]+", + ], ) diff --git a/docs/src/vector.md b/docs/src/vector.md index 3d3120f17..975f74cce 100644 --- a/docs/src/vector.md +++ b/docs/src/vector.md @@ -117,6 +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. -The default AD backends tested are ForwardDiff, ReverseDiff, Mooncake, and Enzyme. +`test_all` accepts an `adtypes` keyword argument and defaults to `[AutoForwardDiff()]`, so the main test suite exercises ForwardDiff as the reference backend. +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. -However 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. +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/ext/BijectorsDistributionsADExt.jl b/ext/BijectorsDistributionsADExt.jl deleted file mode 100644 index 4be6b4897..000000000 --- a/ext/BijectorsDistributionsADExt.jl +++ /dev/null @@ -1,92 +0,0 @@ -module BijectorsDistributionsADExt - -using Bijectors -using Bijectors: LinearAlgebra -using Bijectors.Distributions: AbstractMvLogNormal -using DistributionsAD: - TuringDirichlet, - TuringWishart, - TuringInverseWishart, - FillVectorOfUnivariate, - FillMatrixOfUnivariate, - MatrixOfUnivariate, - FillVectorOfMultivariate, - VectorOfMultivariate, - TuringScalMvNormal, - TuringDiagMvNormal, - TuringDenseMvNormal - -# Bijectors - -Bijectors.bijector(::TuringDirichlet) = Bijectors.SimplexBijector() -Bijectors.bijector(::TuringWishart) = Bijectors.PDVecBijector() -Bijectors.bijector(::TuringInverseWishart) = Bijectors.PDVecBijector() -Bijectors.bijector(::TuringScalMvNormal) = identity -Bijectors.bijector(::TuringDiagMvNormal) = identity -Bijectors.bijector(::TuringDenseMvNormal) = identity - -function Bijectors.bijector(d::FillVectorOfUnivariate{Continuous}) - return elementwise(Bijectors.bijector(d.v.value)) -end -function Bijectors.bijector(d::FillMatrixOfUnivariate{Continuous}) - return elementwise(Bijectors.bijector(d.dists.value)) -end -Bijectors.bijector(d::MatrixOfUnivariate{Discrete}) = identity -function Bijectors.bijector(d::MatrixOfUnivariate{Continuous}) - return Bijectors.TruncatedBijector(Bijectors._minmax(d.dists)...) -end -Bijectors.bijector(d::VectorOfMultivariate{Discrete}) = identity -for T in (:VectorOfMultivariate, :FillVectorOfMultivariate) - @eval begin - Bijectors.bijector(d::$T{Continuous,<:MvNormal}) = identity - Bijectors.bijector(d::$T{Continuous,<:TuringScalMvNormal}) = identity - Bijectors.bijector(d::$T{Continuous,<:TuringDiagMvNormal}) = identity - Bijectors.bijector(d::$T{Continuous,<:TuringDenseMvNormal}) = identity - Bijectors.bijector(d::$T{Continuous,<:MvNormalCanon}) = identity - Bijectors.bijector(d::$T{Continuous,<:AbstractMvLogNormal}) = Log() - function Bijectors.bijector(d::$T{Continuous,<:SimplexDistribution}) - return Bijectors.SimplexBijector() - end - function Bijectors.bijector(d::$T{Continuous,<:TuringDirichlet}) - return Bijectors.SimplexBijector() - end - end -end -function Bijectors.bijector(d::FillVectorOfMultivariate{Continuous}) - return Bijectors.columnwise(Bijectors.bijector(d.dists.value)) -end - -Bijectors.isdirichlet(::VectorOfMultivariate{Continuous,<:Dirichlet}) = true -Bijectors.isdirichlet(::VectorOfMultivariate{Continuous,<:TuringDirichlet}) = true -Bijectors.isdirichlet(::TuringDirichlet) = true - -function Bijectors.link(d::TuringDirichlet, x::AbstractVecOrMat{<:Real}) - return Bijectors.SimplexBijector()(x) -end - -function Bijectors.link_jacobian(d::TuringDirichlet, x::AbstractVector{<:Real}) - return jacobian(Bijectors.SimplexBijector(), x) -end - -function Bijectors.invlink(d::TuringDirichlet, y::AbstractVecOrMat{<:Real}) - return inverse(Bijectors.SimplexBijector())(y) -end -function Bijectors.invlink_jacobian(d::TuringDirichlet, y::AbstractVector{<:Real}) - return jacobian(inverse(Bijectors.SimplexBijector()), y) -end - -Bijectors.ispd(::TuringWishart) = true -Bijectors.ispd(::TuringInverseWishart) = true -function Bijectors.getlogp(d::TuringWishart, Xcf, X) - return ( - (d.df - (size(d, 1) + 1)) * LinearAlgebra.logdet(Xcf) - LinearAlgebra.tr(d.chol \ X) - ) / 2 + d.logc0 -end -function Bijectors.getlogp(d::TuringInverseWishart, Xcf, X) - Ψ = d.S - return -( - (d.df + size(d, 1) + 1) * LinearAlgebra.logdet(Xcf) + LinearAlgebra.tr(Xcf \ Ψ) - ) / 2 + d.logc0 -end - -end diff --git a/src/vector/test_utils.jl b/src/vector/test_utils.jl index bb5003f49..fe73053f9 100644 --- a/src/vector/test_utils.jl +++ b/src/vector/test_utils.jl @@ -1,20 +1,14 @@ using Test using LinearAlgebra: logabsdet, Cholesky, UpperTriangular, LowerTriangular import DifferentiationInterface as DI -import EnzymeCore as EC # 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() -const default_adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=EC.Forward, function_annotation=EC.Const), - DI.AutoEnzyme(; mode=EC.Reverse, function_annotation=EC.Const), -] +# 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] _get_value_support(::D.Distribution{<:Any,VS}) where {VS<:D.ValueSupport} = VS @@ -234,12 +228,14 @@ function test_all( d::D.Distribution; expected_zero_allocs=(), adtypes=default_adtypes, + broken_adtypes=DI.AbstractADType[], ad_atol=1e-10, ad_rtol=sqrt(eps()), roundtrip_atol=1e-10, roundtrip_rtol=sqrt(eps()), test_in_support=(_get_value_support(d) <: D.Continuous), test_construction_type_stable=true, + skip::Bool=false, ) @info "Testing $(_name(d))" @testset "$(_name(d))" begin @@ -250,7 +246,14 @@ function test_all( test_optics(d) test_allocations(d, expected_zero_allocs) test_logjac(d, ad_atol, ad_rtol) - test_ad(d, adtypes, ad_atol, ad_rtol) + if !skip + # Any adtype in `broken_adtypes` is wrapped via `_test_ad_broken`; the rest + # of `adtypes` run normally under `test_ad`. + normal = filter(a -> !(a in broken_adtypes), adtypes) + isempty(normal) || test_ad(d, normal, ad_atol, ad_rtol) + isempty(broken_adtypes) || + _test_ad_broken(d, collect(broken_adtypes), ad_atol, ad_rtol) + end end end @@ -339,14 +342,11 @@ function test_roundtrip_inverse(d::D.Distribution, test_in_support, atol, rtol) end ynew = ffwd(x) - if d isa D.JointOrderStatistics && ( - any(isnan, x) || - !all(isfinite, x) || - any(isnan, ynew) || - !all(isfinite, ynew) - ) - @warn "NaNs or Inf produced in roundtrip test for $(_name(d)), skipping isapprox test" - else + # JointOrderStatistics: extreme random `y` pushes `x` to the support + # boundary, where the y → x → ynew roundtrip loses precision (or produces + # NaN/Inf). Structural correctness is covered by `test_roundtrip`. See + # https://github.com/TuringLang/Bijectors.jl/issues/441. + if !(d isa D.JointOrderStatistics) @test _isapprox_safe(y, ynew; atol=atol, rtol=rtol) end end @@ -591,23 +591,8 @@ Test that various AD backends can differentiate the conversions to and from vect linked vector forms for the given distribution `d`. """ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol) - # If `d` is a discrete distribution, Mooncake refuses to differentiate through the - # transforms (which are just identity transforms). Likewise, Enzyme will throw an - # error saying that the output is Const but was not marked as such. - # - # Arguably, the other AD backends probably should also refuse to differentiate it, but - # they do actually return the right 'gradients' so we can test them. - adtypes = if d isa D.Distribution{<:Any,D.Discrete} - filter(adtypes) do adtype - !( - adtype isa DI.AutoMooncake || - adtype isa DI.AutoMooncakeForward || - adtype isa DI.AutoEnzyme - ) - end - else - adtypes - end + # Discrete-distribution link transforms are identities; nothing AD-meaningful to check. + d isa D.Distribution{<:Any,D.Discrete} && return nothing @testset "AD forward: $(_name(d))" begin x = _rand_safe_ad(d) @@ -652,3 +637,48 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, end end 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 isa D.Distribution{<:Any,D.Discrete} && return nothing + @testset "AD (broken): $(_name(d))" begin + x = _rand_safe_ad(d) + xvec = to_vec(d)(x) + yvec = to_linked_vec(d)(x) + ffwd = to_linked_vec(d) ∘ from_vec(d) + 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) + 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 + ) && + isapprox( + DI.gradient(ladj_fwd, adtype, xvec), + ref_grad_ladj_fwd; + atol=atol, + rtol=rtol, + ) && + isapprox( + DI.jacobian(frvs, adtype, yvec), ref_jac_rev; atol=atol, rtol=rtol + ) && + isapprox( + DI.gradient(ladj_rev, adtype, yvec), + ref_grad_ladj_rev; + atol=atol, + rtol=rtol, + ) + ) + end + end + end +end diff --git a/test/Project.toml b/test/Project.toml index 570cc7521..2a5e13a6d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,9 +9,6 @@ ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" -Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -22,11 +19,9 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0" -Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" @@ -42,9 +37,6 @@ ChangesOfVariables = "0.1" Combinatorics = "1.0.2" DifferentiationInterface = "0.7.7" Distributions = "0.25" -DistributionsAD = "0.6.3" -Documenter = "1" -Enzyme = "0.13" FillArrays = "1" FiniteDifferences = "0.11, 0.12" ForwardDiff = "0.10, 1.0.1" @@ -54,9 +46,7 @@ LazyArrays = "1, 2" LogDensityProblems = "2" LogExpFunctions = "0.3.1" MCMCDiagnosticTools = "0.3" -Mooncake = "0.4, 0.5" PDMats = "0.11" -ReverseDiff = "1.4.2" StableRNGs = "1" Tracker = "0.2.11" julia = "1.10" diff --git a/test/ad/corr.jl b/test/ad/corr.jl deleted file mode 100644 index 7d236eddb..000000000 --- a/test/ad/corr.jl +++ /dev/null @@ -1,42 +0,0 @@ -@testset "VecCorrBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES - # Enzyme is tested separately as these tests are flaky - # TODO(penelopeysm): Fix upstream and re-enable. - if adtype isa AutoEnzyme - continue - end - - @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) - - roundtrip(y) = sum(transform(b, binv(y))) - inverse_only(y) = sum(transform(binv, y)) - test_ad(roundtrip, adtype, y) - test_ad(inverse_only, adtype, y) - end -end - -@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) - binv = inverse(b) - - x = rand(dist) - y = b(x) - cholesky_to_triangular = - uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower - - 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/enzyme.jl b/test/ad/enzyme.jl deleted file mode 100644 index 39ef5b177..000000000 --- a/test/ad/enzyme.jl +++ /dev/null @@ -1,68 +0,0 @@ -module BijectorsEnzymeRulesTests - -using Bijectors -using Enzyme -using EnzymeTestUtils: test_forward, test_reverse -using Test - -# 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 -# -# The fix to this needs to be made in Julia itself: it seems that this has already been done -# in https://github.com/JuliaLang/llvm-project/pull/49 although whether this will be -# incorporated into the built Julia version itself seems unclear. See -# https://github.com/JuliaLang/julia/pull/59521#issuecomment-3300480633. -# -# If this does not end up being backported to 1.11, then we may have to permanently skip -# these tests. -# -# On another note: 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) - - 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) - - 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 -end - -end diff --git a/test/ad/flows.jl b/test/ad/flows.jl deleted file mode 100644 index ad649f379..000000000 --- a/test/ad/flows.jl +++ /dev/null @@ -1,41 +0,0 @@ -@testset "PlanarLayer: $backend_name" for (backend_name, adtype) in TEST_ADTYPES - # Enzyme is tested separately as these tests are flaky - # TODO(penelopeysm): Fix upstream and re-enable. - if adtype isa AutoEnzyme - continue - end - - # logpdf of a flow with a planar layer and two-dimensional inputs - 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(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 - 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(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/mooncake.jl b/test/ad/mooncake.jl deleted file mode 100644 index 1e9259be7..000000000 --- a/test/ad/mooncake.jl +++ /dev/null @@ -1,33 +0,0 @@ -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/ad/pd.jl b/test/ad/pd.jl deleted file mode 100644 index 16f715f59..000000000 --- a/test/ad/pd.jl +++ /dev/null @@ -1,27 +0,0 @@ -_topd(x) = x * x' + I - -@testset "PDVecBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES - # Enzyme is tested separately as these tests are flaky - # TODO(penelopeysm): Fix upstream and re-enable. - if adtype isa AutoEnzyme - continue - end - - d = 4 - b = Bijectors.PDVecBijector() - binv = inverse(b) - - z = randn(d, d) - x = _topd(z) - y = b(x) - - 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))) - - test_ad(forward_only, adtype, vec(z)) - test_ad(inverse_only, adtype, y) - test_ad(inverse_chol_lower, adtype, y) - test_ad(inverse_chol_upper, adtype, y) -end diff --git a/test/ad/stacked.jl b/test/ad/stacked.jl deleted file mode 100644 index fe86fcda7..000000000 --- a/test/ad/stacked.jl +++ /dev/null @@ -1,28 +0,0 @@ -@testset "StackedBijector: $backend_name" for (backend_name, adtype) in TEST_ADTYPES - dist1 = Dirichlet(4, 1.0) - b1 = bijector(dist1) - - dist2 = LogNormal(0.0, 1.0) - b2 = bijector(dist2) - - x1 = rand(dist1) - x2 = rand(dist2) - - y1 = b1(x1) - y2 = b2(x2) - - b = Stacked((b1, b2), (1:4, 5:5)) - binv = inverse(b) - - y = vcat(y1, [y2]) - x = binv(y) - - 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 -> sum(transform(bvec, binv(y))), adtype, y) - test_ad(y -> sum(transform(bvec_inv, y)), adtype, y) -end diff --git a/test/ad/chainrules.jl b/test/bijectors/chainrules.jl similarity index 100% rename from test/ad/chainrules.jl rename to test/bijectors/chainrules.jl diff --git a/test/bijectors/corr.jl b/test/bijectors/corr.jl index 6e153fab9..c48349dab 100644 --- a/test/bijectors/corr.jl +++ b/test/bijectors/corr.jl @@ -1,4 +1,4 @@ -using Bijectors, DistributionsAD, LinearAlgebra, Test +using Bijectors, LinearAlgebra, Test using Bijectors: VecCorrBijector, VecCholeskyBijector, CorrBijector @testset "CorrBijector & VecCorrBijector" begin diff --git a/test/bijectors/equality.jl b/test/bijectors/equality.jl new file mode 100644 index 000000000..1a6a3992f --- /dev/null +++ b/test/bijectors/equality.jl @@ -0,0 +1,46 @@ +@testset "Equality" begin + bs = [ + identity, + elementwise(exp), + elementwise(log), + Scale(2.0), + Scale(3.0), + Scale(rand(2, 2)), + Scale(rand(2, 2)), + Shift(2.0), + Shift(3.0), + Shift(rand(2)), + Shift(rand(2)), + Logit(1.0, 2.0), + Logit(1.0, 3.0), + Logit(2.0, 3.0), + Logit(0.0, 2.0), + InvertibleBatchNorm(2), + InvertibleBatchNorm(3), + PDBijector(), + Permute([1.0, 2.0, 3.0]), + Permute([2.0, 3.0, 4.0]), + PlanarLayer(2), + PlanarLayer(3), + RadialLayer(2), + RadialLayer(3), + SimplexBijector(), + Stacked((elementwise(exp), elementwise(log))), + Stacked((elementwise(log), elementwise(exp))), + Stacked([elementwise(exp), elementwise(log)]), + Stacked([elementwise(log), elementwise(exp)]), + elementwise(exp) ∘ elementwise(log), + elementwise(log) ∘ elementwise(exp), + TruncatedBijector(1.0, 2.0), + TruncatedBijector(1.0, 3.0), + TruncatedBijector(0.0, 2.0), + ] + for i in eachindex(bs), j in eachindex(bs) + if i == j + @test bs[i] == deepcopy(bs[j]) + @test inverse(bs[i]) == inverse(deepcopy(bs[j])) + else + @test bs[i] != bs[j] + end + end +end diff --git a/test/bijectors/pd.jl b/test/bijectors/pd.jl index 2375d27a9..21c6aaa7b 100644 --- a/test/bijectors/pd.jl +++ b/test/bijectors/pd.jl @@ -1,4 +1,4 @@ -using Bijectors, DistributionsAD, LinearAlgebra, Test +using Bijectors, LinearAlgebra, Test using Bijectors: PDBijector, PDVecBijector @testset "PDBijector" begin diff --git a/test/bijectors/scale.jl b/test/bijectors/scale.jl new file mode 100644 index 000000000..abedd5582 --- /dev/null +++ b/test/bijectors/scale.jl @@ -0,0 +1,7 @@ +@testset "test_inverse and test_with_logabsdet_jacobian" begin + b = Bijectors.Scale{Float64}(4.2) + x = 0.3 + + InverseFunctions.test_inverse(b, x) + ChangesOfVariables.test_with_logabsdet_jacobian(b, x, (f::Bijectors.Scale, x) -> f.a) +end diff --git a/test/bijectors/simplex.jl b/test/bijectors/simplex.jl new file mode 100644 index 000000000..c2cea1273 --- /dev/null +++ b/test/bijectors/simplex.jl @@ -0,0 +1,11 @@ +@testset "Jacobians of SimplexBijector" begin + b = SimplexBijector() + ib = inverse(b) + + d_x = 10 + x = ib(randn(d_x - 1)) + y = b(x) + + @test Bijectors.jacobian(b, x) ≈ ForwardDiff.jacobian(b, x) + @test Bijectors.jacobian(ib, y) ≈ ForwardDiff.jacobian(ib, y) +end diff --git a/test/bijectors/stacked.jl b/test/bijectors/stacked.jl index 1f221b398..9a2d87c8b 100644 --- a/test/bijectors/stacked.jl +++ b/test/bijectors/stacked.jl @@ -75,3 +75,166 @@ end @test binv(y) == [1.0, 2.0, 0.0, 0.0] end end + +@testset "Stacked <: Bijector" begin + # `logabsdetjac` withOUT AD + d = Beta() + b = bijector(d) + x = rand(d) + y = b(x) + + sb1 = @inferred Stacked(b, b, inverse(b), inverse(b)) # <= Tuple + res1 = with_logabsdet_jacobian(sb1, [x, x, y, y]) + + @test sb1([x, x, y, y]) ≈ res1[1] + @test logabsdetjac(sb1, [x, x, y, y]) ≈ 0 atol = 1e-6 + @test res1[2] ≈ 0 atol = 1e-6 + + sb2 = Stacked([b, b, inverse(b), inverse(b)]) # <= Array + res2 = with_logabsdet_jacobian(sb2, [x, x, y, y]) + + @test sb2([x, x, y, y]) ≈ res2[1] + @test logabsdetjac(sb2, [x, x, y, y]) ≈ 0.0 atol = 1e-12 + @test res2[2] ≈ 0.0 atol = 1e-12 + + # value-test + x = ones(3) + sb = @inferred Stacked(elementwise(exp), elementwise(log), Shift(5.0)) + res = with_logabsdet_jacobian(sb, x) + @test sb(x) == [exp(x[1]), log(x[2]), x[3] + 5.0] + @test res[1] == [exp(x[1]), log(x[2]), x[3] + 5.0] + @test logabsdetjac(sb, x) == + sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:3]) + @test res[2] == logabsdetjac(sb, x) + + # TODO: change when we have dimensionality in the type + sb = @inferred Stacked((elementwise(exp), SimplexBijector()), (1:1, 2:3)) + x = ones(3) ./ 3.0 + res = @inferred with_logabsdet_jacobian(sb, x) + @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] + @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] + @test logabsdetjac(sb, x) == + sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) + @test res[2] == logabsdetjac(sb, x) + + x = ones(4) ./ 4.0 + @test_throws ErrorException sb(x) + + # Array-version + sb = Stacked([elementwise(exp), SimplexBijector()], [1:1, 2:3]) + x = ones(3) ./ 3.0 + res = with_logabsdet_jacobian(sb, x) + @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] + @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] + @test logabsdetjac(sb, x) == + sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) + @test res[2] == logabsdetjac(sb, x) + + x = ones(4) ./ 4.0 + @test_throws ErrorException sb(x) + + # Mixed versions + # Tuple, Array + sb = Stacked([elementwise(exp), SimplexBijector()], (1:1, 2:3)) + x = ones(3) ./ 3.0 + res = with_logabsdet_jacobian(sb, x) + @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] + @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] + @test logabsdetjac(sb, x) == + sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) + @test res[2] == logabsdetjac(sb, x) + + x = ones(4) ./ 4.0 + @test_throws ErrorException sb(x) + + # Array, Tuple + sb = Stacked((elementwise(exp), SimplexBijector()), [1:1, 2:3]) + x = ones(3) ./ 3.0 + res = with_logabsdet_jacobian(sb, x) + @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] + @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] + @test logabsdetjac(sb, x) == + sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) + @test res[2] == logabsdetjac(sb, x) + + x = ones(4) ./ 4.0 + @test_throws ErrorException sb(x) + + @testset "Stacked: ADVI with MvNormal" begin + # MvNormal test + dists = [ + Beta(), + Beta(), + Beta(), + InverseGamma(), + InverseGamma(), + Gamma(), + Gamma(), + InverseGamma(), + Cauchy(), + Gamma(), + MvNormal(zeros(2), I), + ] + + ranges = [] + idx = 1 + for i in 1:length(dists) + d = dists[i] + push!(ranges, idx:(idx + length(d) - 1)) + idx += length(d) + end + ranges = tuple(ranges...) + + num_params = ranges[end][end] + d = MvNormal(zeros(num_params), I) + + # Stacked{<:Array} + bs = bijector.(dists) # constrained-to-unconstrained bijectors for dists + ibs = inverse.(bs) # invert, so we get unconstrained-to-constrained + sb = Stacked(ibs, ranges) # => Stacked <: Bijector + x = rand(d) + + @test sb isa Stacked + + td = transformed(d, sb) # => MultivariateTransformed <: Distribution{Multivariate, Continuous} + @test td isa Distribution{Multivariate,Continuous} + + # check that wrong ranges fails + sb = Stacked(ibs) + x = rand(d) + @test_throws ErrorException sb(x) + + # Stacked{<:Tuple} + bs = bijector.(tuple(dists...)) + ibs = inverse.(bs) + sb = @inferred Stacked(ibs, ranges) + isb = @inferred inverse(sb) + @test sb isa Stacked{<:Tuple} + + # inverse + td = @inferred transformed(d, sb) + y = @inferred rand(td) + x = @inferred isb(y) + @test sb(x) ≈ y + + # verification of computation + x = rand(d) + y = sb(x) + y_ = vcat([ibs[i](x[ranges[i]]) for i in 1:length(dists)]...) + x_ = vcat([bs[i](y[ranges[i]]) for i in 1:length(dists)]...) + @test x ≈ x_ + @test y ≈ y_ + + # AD verification + @test log(abs(det(ForwardDiff.jacobian(sb, x)))) ≈ logabsdetjac(sb, x) + @test log(abs(det(ForwardDiff.jacobian(isb, y)))) ≈ logabsdetjac(isb, y) + + # Ensure `Stacked` works for a single bijector + d = (MvNormal(zeros(2), I),) + sb = Stacked(bijector.(d), (1:2,)) + x = [0.5, 1.0] + @test sb(x) == x + @test logabsdetjac(sb, x) == 0 + @test with_logabsdet_jacobian(sb, x) == (x, zero(eltype(x))) + end +end diff --git a/test/distributionsad.jl b/test/distributionsad.jl deleted file mode 100644 index d30fb0dcf..000000000 --- a/test/distributionsad.jl +++ /dev/null @@ -1,7 +0,0 @@ -@testset "DistributionsAD" begin - # https://github.com/TuringLang/Bijectors.jl/issues/298 - @testset "#298" begin - dists = arraydist(fill(InverseGamma(), 2, 2)) - @test bijector(dists) isa Bijectors.TruncatedBijector - end -end diff --git a/test/integration/enzyme/main.jl b/test/integration/enzyme/main.jl deleted file mode 100644 index 45cb5c5be..000000000 --- a/test/integration/enzyme/main.jl +++ /dev/null @@ -1,105 +0,0 @@ -using ADTypes -using Bijectors -using DifferentiationInterface -using Enzyme -using FiniteDifferences -using LinearAlgebra -using Test - -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) -end - -BACKENDS = [ - ("EnzymeForward", AutoEnzyme(; mode=Forward)), - ("EnzymeReverse", AutoEnzyme(; mode=Reverse)), -] - -@testset "$backend" for (backend, adtype) in BACKENDS - @info "Testing with backend: $backend" - - @testset "VecCorrBijector" begin - @info " - Testing VecCorrBijector" - - @testset "d = $d" for d in (1, 2, 4) - @info " - Dimension: $d" - - dist = LKJ(d, 2.0) - b = bijector(dist) - binv = inverse(b) - - x = rand(dist) - y = b(x) - - roundtrip(y) = sum(transform(b, binv(y))) - inverse_only(y) = sum(transform(binv, y)) - test_ad(roundtrip, adtype, y) - test_ad(inverse_only, adtype, y) - end - end - - @testset "PlanarLayer" begin - @info " - Testing PlanarLayer" - # logpdf of a flow with a planar layer and two-dimensional inputs - 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(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 - 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(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 - - @testset "PDVecBijector" begin - @info " - Testing PDVecBijector" - _topd(x) = x * x' + I - - d = 4 - b = Bijectors.PDVecBijector() - binv = inverse(b) - - z = randn(d, d) - x = _topd(z) - y = b(x) - - 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))) - - 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/integration/enzyme/Project.toml b/test/integration_tests/enzyme/Project.toml similarity index 54% rename from test/integration/enzyme/Project.toml rename to test/integration_tests/enzyme/Project.toml index 3ce74b7ab..36e6c82dd 100644 --- a/test/integration/enzyme/Project.toml +++ b/test/integration_tests/enzyme/Project.toml @@ -2,13 +2,19 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" 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" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +[sources] +Bijectors = {path = "../../.."} + [compat] Enzyme = "0.13.111" - -[sources] -Bijectors = {path = "../../../"} diff --git a/test/integration_tests/enzyme/main.jl b/test/integration_tests/enzyme/main.jl new file mode 100644 index 000000000..bdd7495d3 --- /dev/null +++ b/test/integration_tests/enzyme/main.jl @@ -0,0 +1,114 @@ +using ADTypes +using Bijectors +using 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 +using Test + +include(joinpath(@__DIR__, "..", "..", "test_resources.jl")) + +# `set_runtime_activity` is empirically required for VecCholesky / Stacked / LKJCholesky / +# products. Applying it uniformly (rather than per-tag) is the conservative form that works +# everywhere — we lose the assertion "this case happens to work without runtime_activity" +# in exchange for a single adtype list across all cases. `function_annotation=Const` pins +# DI's default to match what these tests were originally validated against. +const ENZYME_FORWARD = AutoEnzyme(; + mode=set_runtime_activity(Forward), function_annotation=Const +) +const ENZYME_REVERSE = AutoEnzyme(; + mode=set_runtime_activity(Reverse), function_annotation=Const +) +const adtypes = [ENZYME_FORWARD, ENZYME_REVERSE] + +# Enzyme cannot differentiate through triple-nested tuple-of-products (e.g. +# `product_distribution(p1t, p1t, p1t)`); identify them structurally. +function _enzyme_failing_product(d) + d isa Distributions.ProductDistribution || return false + d.dists isa Tuple || return false + return first(d.dists) isa Union{Distributions.Product,Distributions.ProductDistribution} +end + +# Return the subset of `adtypes` that is known-broken for this case. Only Reverse mode +# hits https://github.com/EnzymeAD/Enzyme.jl/issues/2987 on `:reshaped_beta_special` on +# Julia 1.10 — Forward mode passes. Triple-nested products in `:type_unstable_products` +# defeat activity inference for both modes. +function vector_broken_adtypes(c::VectorTestCase) + c.tag === :reshaped_beta_special && VERSION < v"1.11-" && return [ENZYME_REVERSE] + c.tag === :type_unstable_products && _enzyme_failing_product(c.dist) && return adtypes + return DI.AbstractADType[] +end + +# 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 +# +# The fix to this needs to be made in Julia itself: it seems that this has already been done +# in https://github.com/JuliaLang/llvm-project/pull/49 although whether this will be +# incorporated into the built Julia version itself seems unclear. See +# https://github.com/JuliaLang/julia/pull/59521#issuecomment-3300480633. +# +# If this does not end up being backported to 1.11, then we may have to permanently skip +# these tests. +# +# On another note: 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 + @testset for RT in (Const, Enzyme.Duplicated, Enzyme.DuplicatedNoNeed), + Tx in (Const, Enzyme.Duplicated), + Ty in (Const, Enzyme.Duplicated), + Tz in (Const, Enzyme.Duplicated) + + test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + + @testset for RT in + (Const, Enzyme.BatchDuplicated, Enzyme.BatchDuplicatedNoNeed), + Tx in (Const, Enzyme.BatchDuplicated), + Ty in (Const, Enzyme.BatchDuplicated), + Tz in (Const, Enzyme.BatchDuplicated) + + test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + end + @testset "reverse" begin + @testset for RT in (Const, Enzyme.Active), + Tx in (Const, Enzyme.Active), + Ty in (Const, Enzyme.Active), + Tz in (Const, Enzyme.Active) + + test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz)) + end + + # TODO: Test batch mode. Enzyme does not support all combinations of activities + # currently: + # https://github.com/TuringLang/Bijectors.jl/pull/350#issuecomment-2480468728 + end + end +end + +@testset "Enzyme bijector AD" begin + for c in generate_ad_testcases(), adtype in adtypes + run_ad_case(c, adtype) + end +end + +@testset "Enzyme vector test_all" begin + for c in generate_vector_testcases() + run_vector_case(c, adtypes; broken_adtypes=vector_broken_adtypes(c)) + end +end diff --git a/test/integration_tests/mooncake/Project.toml b/test/integration_tests/mooncake/Project.toml new file mode 100644 index 000000000..a914acb95 --- /dev/null +++ b/test/integration_tests/mooncake/Project.toml @@ -0,0 +1,20 @@ +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +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" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Bijectors = {path = "../../.."} + +[compat] +Mooncake = "0.4, 0.5" diff --git a/test/integration_tests/mooncake/main.jl b/test/integration_tests/mooncake/main.jl new file mode 100644 index 000000000..77b416bc2 --- /dev/null +++ b/test/integration_tests/mooncake/main.jl @@ -0,0 +1,53 @@ +using ADTypes +using Bijectors +using DifferentiationInterface +using Distributions +using FillArrays: Fill +using FiniteDifferences +using ForwardDiff: ForwardDiff +using LinearAlgebra +using Mooncake: Mooncake +using PDMats +using Random: Xoshiro +using Test + +include(joinpath(@__DIR__, "..", "..", "test_resources.jl")) + +const adtypes = [AutoMooncake(), AutoMooncakeForward()] + +# ===== Mooncake rule for `find_alpha` ===== + +@testset "Mooncake $mode: find_alpha" for mode in (Mooncake.ReverseMode,) + rng = Xoshiro(123456) + x = randn() + y = expm1(randn()) + z = randn() + 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 + +@testset "Mooncake bijector AD" begin + for c in generate_ad_testcases(), adtype in adtypes + run_ad_case(c, adtype) + end +end + +@testset "Mooncake vector test_all" begin + for c in generate_vector_testcases() + run_vector_case(c, adtypes) + end +end diff --git a/test/integration_tests/reversediff/Project.toml b/test/integration_tests/reversediff/Project.toml new file mode 100644 index 000000000..2178ac33d --- /dev/null +++ b/test/integration_tests/reversediff/Project.toml @@ -0,0 +1,20 @@ +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +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" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +Bijectors = {path = "../../.."} + +[compat] +ReverseDiff = "1.4.2" diff --git a/test/integration_tests/reversediff/main.jl b/test/integration_tests/reversediff/main.jl new file mode 100644 index 000000000..f92298948 --- /dev/null +++ b/test/integration_tests/reversediff/main.jl @@ -0,0 +1,37 @@ +using ADTypes +using Bijectors +using ChainRules: ChainRules +using DifferentiationInterface +using Distributions +using FillArrays: Fill +using FiniteDifferences +using ForwardDiff: ForwardDiff +using LinearAlgebra +using PDMats +using ReverseDiff: ReverseDiff +using Test + +include(joinpath(@__DIR__, "..", "..", "test_resources.jl")) + +const adtypes = [AutoReverseDiff(), AutoReverseDiff(; compile=true)] + +# LKJ: https://github.com/TuringLang/Bijectors.jl/issues/434 (wrong forward jacobian +# through VecCorrBijector). JointOrderStatistics: https://github.com/JuliaDiff/ReverseDiff.jl/issues/43 +# (setindex!). +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[] +end + +@testset "ReverseDiff bijector AD" begin + for c in generate_ad_testcases(), adtype in adtypes + run_ad_case(c, adtype) + end +end + +@testset "ReverseDiff vector test_all" begin + for c in generate_vector_testcases() + run_vector_case(c, adtypes; broken_adtypes=vector_broken_adtypes(c)) + end +end diff --git a/test/interface.jl b/test/interface.jl index 084d23acc..bb40407b5 100644 --- a/test/interface.jl +++ b/test/interface.jl @@ -2,30 +2,10 @@ using Test using Random using LinearAlgebra using ForwardDiff -using ReverseDiff -using DistributionsAD using Bijectors -using Bijectors: - Shift, - Scale, - Logit, - SimplexBijector, - PDBijector, - Permute, - PlanarLayer, - RadialLayer, - Stacked, - TruncatedBijector, - RationalQuadraticSpline, - LeakyReLU - Random.seed!(123) -contains(predicate::Function, b::Bijector) = predicate(b) -contains(predicate::Function, b::ComposedFunction) = any(contains.(predicate, b.ts)) -contains(predicate::Function, b::Stacked) = any(contains.(predicate, b.bs)) - @testset "Univariate" begin # Tests with scalar-valued distributions. uni_dists = [ @@ -165,7 +145,7 @@ end @test logpdf(td, y) ≈ logpdf_with_trans(dist, x, true) # verify against AD - # similar to what we do in test/transform.jl for Dirichlet + # similar to what we do in test/legacy_interface.jl for Dirichlet if dist isa Dirichlet b = Bijectors.SimplexBijector() x = rand(dist) @@ -197,8 +177,6 @@ end matrix_dists = [ Wishart(v, S), InverseWishart(v, S), - TuringWishart(v, S), - TuringInverseWishart(v, S), LKJ(3, 1.0), reshape(MvNormal(zeros(6), I), 2, 3), product_distribution(fill(InverseGamma(2, 3), 6)), @@ -230,270 +208,3 @@ end @test logpdf_with_trans(d, x, false) == logpdf(d, x) @test logpdf_with_trans(d, x, true) == logpdf(d, x) - logabsdetjac(b, x) end - -@testset "DistributionsAD" begin - @testset "$dist" for dist in [ - filldist(Normal(), 2), - filldist(Normal(), 2, 3), - filldist(Exponential(), 2), - filldist(Exponential(), 2, 3), - filldist(filldist(Exponential(), 2), 3), - filldist(Dirichlet(ones(2)), 3), - ] - x = rand(dist) - b = bijector(dist) - y = b(x) - td = transformed(dist) - @test logpdf(dist, x) - logabsdetjac(b, x) ≈ logpdf(td, y) - end -end - -@testset "Stacked <: Bijector" begin - # `logabsdetjac` withOUT AD - d = Beta() - b = bijector(d) - x = rand(d) - y = b(x) - - sb1 = @inferred Stacked(b, b, inverse(b), inverse(b)) # <= Tuple - res1 = with_logabsdet_jacobian(sb1, [x, x, y, y]) - - @test sb1([x, x, y, y]) ≈ res1[1] - @test logabsdetjac(sb1, [x, x, y, y]) ≈ 0 atol = 1e-6 - @test res1[2] ≈ 0 atol = 1e-6 - - sb2 = Stacked([b, b, inverse(b), inverse(b)]) # <= Array - res2 = with_logabsdet_jacobian(sb2, [x, x, y, y]) - - @test sb2([x, x, y, y]) ≈ res2[1] - @test logabsdetjac(sb2, [x, x, y, y]) ≈ 0.0 atol = 1e-12 - @test res2[2] ≈ 0.0 atol = 1e-12 - - # value-test - x = ones(3) - sb = @inferred Stacked(elementwise(exp), elementwise(log), Shift(5.0)) - res = with_logabsdet_jacobian(sb, x) - @test sb(x) == [exp(x[1]), log(x[2]), x[3] + 5.0] - @test res[1] == [exp(x[1]), log(x[2]), x[3] + 5.0] - @test logabsdetjac(sb, x) == - sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:3]) - @test res[2] == logabsdetjac(sb, x) - - # TODO: change when we have dimensionality in the type - sb = @inferred Stacked((elementwise(exp), SimplexBijector()), (1:1, 2:3)) - x = ones(3) ./ 3.0 - res = @inferred with_logabsdet_jacobian(sb, x) - @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] - @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] - @test logabsdetjac(sb, x) == - sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) - @test res[2] == logabsdetjac(sb, x) - - x = ones(4) ./ 4.0 - @test_throws ErrorException sb(x) - - # Array-version - sb = Stacked([elementwise(exp), SimplexBijector()], [1:1, 2:3]) - x = ones(3) ./ 3.0 - res = with_logabsdet_jacobian(sb, x) - @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] - @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] - @test logabsdetjac(sb, x) == - sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) - @test res[2] == logabsdetjac(sb, x) - - x = ones(4) ./ 4.0 - @test_throws ErrorException sb(x) - - # Mixed versions - # Tuple, Array - sb = Stacked([elementwise(exp), SimplexBijector()], (1:1, 2:3)) - x = ones(3) ./ 3.0 - res = with_logabsdet_jacobian(sb, x) - @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] - @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] - @test logabsdetjac(sb, x) == - sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) - @test res[2] == logabsdetjac(sb, x) - - x = ones(4) ./ 4.0 - @test_throws ErrorException sb(x) - - # Array, Tuple - sb = Stacked((elementwise(exp), SimplexBijector()), [1:1, 2:3]) - x = ones(3) ./ 3.0 - res = with_logabsdet_jacobian(sb, x) - @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...] - @test res[1] == [exp(x[1]), sb.bs[2](x[2:3])...] - @test logabsdetjac(sb, x) == - sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges_in[i]])) for i in 1:2]) - @test res[2] == logabsdetjac(sb, x) - - x = ones(4) ./ 4.0 - @test_throws ErrorException sb(x) - - @testset "Stacked: ADVI with MvNormal" begin - # MvNormal test - dists = [ - Beta(), - Beta(), - Beta(), - InverseGamma(), - InverseGamma(), - Gamma(), - Gamma(), - InverseGamma(), - Cauchy(), - Gamma(), - MvNormal(zeros(2), I), - ] - - ranges = [] - idx = 1 - for i in 1:length(dists) - d = dists[i] - push!(ranges, idx:(idx + length(d) - 1)) - idx += length(d) - end - ranges = tuple(ranges...) - - num_params = ranges[end][end] - d = MvNormal(zeros(num_params), I) - - # Stacked{<:Array} - bs = bijector.(dists) # constrained-to-unconstrained bijectors for dists - ibs = inverse.(bs) # invert, so we get unconstrained-to-constrained - sb = Stacked(ibs, ranges) # => Stacked <: Bijector - x = rand(d) - - @test sb isa Stacked - - td = transformed(d, sb) # => MultivariateTransformed <: Distribution{Multivariate, Continuous} - @test td isa Distribution{Multivariate,Continuous} - - # check that wrong ranges fails - sb = Stacked(ibs) - x = rand(d) - @test_throws ErrorException sb(x) - - # Stacked{<:Tuple} - bs = bijector.(tuple(dists...)) - ibs = inverse.(bs) - sb = @inferred Stacked(ibs, ranges) - isb = @inferred inverse(sb) - @test sb isa Stacked{<:Tuple} - - # inverse - td = @inferred transformed(d, sb) - y = @inferred rand(td) - x = @inferred isb(y) - @test sb(x) ≈ y - - # verification of computation - x = rand(d) - y = sb(x) - y_ = vcat([ibs[i](x[ranges[i]]) for i in 1:length(dists)]...) - x_ = vcat([bs[i](y[ranges[i]]) for i in 1:length(dists)]...) - @test x ≈ x_ - @test y ≈ y_ - - # AD verification - @test log(abs(det(ForwardDiff.jacobian(sb, x)))) ≈ logabsdetjac(sb, x) - @test log(abs(det(ForwardDiff.jacobian(isb, y)))) ≈ logabsdetjac(isb, y) - - # Ensure `Stacked` works for a single bijector - d = (MvNormal(zeros(2), I),) - sb = Stacked(bijector.(d), (1:2,)) - x = [0.5, 1.0] - @test sb(x) == x - @test logabsdetjac(sb, x) == 0 - @test with_logabsdet_jacobian(sb, x) == (x, zero(eltype(x))) - end -end - -@testset "Example: ADVI single" begin - # Usage in ADVI - d = Beta() - b = bijector(d) # [0, 1] → ℝ - ib = inverse(b) # ℝ → [0, 1] - td = transformed(Normal(), ib) # x ∼ 𝓝(0, 1) then f(x) ∈ [0, 1] - x = rand(td) # ∈ [0, 1] - @test 0 ≤ x ≤ 1 -end - -@testset "Jacobians of SimplexBijector" begin - b = SimplexBijector() - ib = inverse(b) - - d_x = 10 - x = ib(randn(d_x - 1)) - y = b(x) - - @test Bijectors.jacobian(b, x) ≈ ForwardDiff.jacobian(b, x) - @test Bijectors.jacobian(ib, y) ≈ ForwardDiff.jacobian(ib, y) - - # Just some additional computation so we also ensure the pullbacks are the same - weights_x = randn(d_x) - weights_y = randn(d_x - 1) - - # ForwardDiff.jl - Δ_forwarddiff = ForwardDiff.gradient(z -> sum(weights_y .* b(z)), x) - - # ForwardDiff.jl - Δ_forwarddiff_inv = ForwardDiff.gradient(z -> sum(weights_x .* ib(z)), y) -end - -@testset "Equality" begin - bs = [ - identity, - elementwise(exp), - elementwise(log), - Scale(2.0), - Scale(3.0), - Scale(rand(2, 2)), - Scale(rand(2, 2)), - Shift(2.0), - Shift(3.0), - Shift(rand(2)), - Shift(rand(2)), - Logit(1.0, 2.0), - Logit(1.0, 3.0), - Logit(2.0, 3.0), - Logit(0.0, 2.0), - InvertibleBatchNorm(2), - InvertibleBatchNorm(3), - PDBijector(), - Permute([1.0, 2.0, 3.0]), - Permute([2.0, 3.0, 4.0]), - PlanarLayer(2), - PlanarLayer(3), - RadialLayer(2), - RadialLayer(3), - SimplexBijector(), - Stacked((elementwise(exp), elementwise(log))), - Stacked((elementwise(log), elementwise(exp))), - Stacked([elementwise(exp), elementwise(log)]), - Stacked([elementwise(log), elementwise(exp)]), - elementwise(exp) ∘ elementwise(log), - elementwise(log) ∘ elementwise(exp), - TruncatedBijector(1.0, 2.0), - TruncatedBijector(1.0, 3.0), - TruncatedBijector(0.0, 2.0), - ] - for i in 1:length(bs), j in 1:length(bs) - if i == j - @test bs[i] == deepcopy(bs[j]) - @test inverse(bs[i]) == inverse(deepcopy(bs[j])) - else - @test bs[i] != bs[j] - end - end -end - -@testset "test_inverse and test_with_logabsdet_jacobian" begin - b = Bijectors.Scale{Float64}(4.2) - x = 0.3 - - InverseFunctions.test_inverse(b, x) - ChangesOfVariables.test_with_logabsdet_jacobian(b, x, (f::Bijectors.Scale, x) -> f.a) -end diff --git a/test/transform.jl b/test/legacy_interface.jl similarity index 100% rename from test/transform.jl rename to test/legacy_interface.jl diff --git a/test/norm_flows.jl b/test/normalising_flows.jl similarity index 100% rename from test/norm_flows.jl rename to test/normalising_flows.jl diff --git a/test/runtests.jl b/test/runtests.jl index 54d237220..522069ea4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,15 +3,10 @@ using Bijectors using ChainRulesTestUtils using Combinatorics using DifferentiationInterface -using DistributionsAD -using Documenter: Documenter using FiniteDifferences using ForwardDiff using Functors using LogExpFunctions -using Mooncake -using ReverseDiff -using Pkg using Random, LinearAlgebra, Test @@ -32,39 +27,10 @@ using InverseFunctions: InverseFunctions using LazyArrays: LazyArrays const GROUP = get(ENV, "GROUP", "All") -# Enzyme doesn't work on 1.12 yet -const TEST_ENZYME = VERSION < v"1.12.0" -TEST_ADTYPES = [ - ("ForwardDiff", AutoForwardDiff()), - ("ReverseDiff", AutoReverseDiff(; compile=false)), - ("ReverseDiffCompiled", AutoReverseDiff(; compile=true)), - ("Mooncake", AutoMooncake()), -] -if TEST_ENZYME - Pkg.add("Enzyme") - Pkg.add("EnzymeTestUtils") - using Enzyme: Enzyme, set_runtime_activity, Forward, Reverse, Const - TEST_ADTYPES = [ - TEST_ADTYPES..., - ( - "EnzymeForward", - AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), - ), - ( - "EnzymeReverse", - AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), - ), - ] -end -const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1)) +# Non-ForwardDiff backends run under `test/integration_tests/`. -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) -end +include("test_resources.jl") # Always include this since it can be useful for other tests. include("bijectors/utils.jl") @@ -72,8 +38,8 @@ include("bijectors/utils.jl") @testset "Bijectors.jl" begin if GROUP == "All" || GROUP == "Classic" include("interface.jl") - include("transform.jl") - include("norm_flows.jl") + include("legacy_interface.jl") + include("normalising_flows.jl") include("bijectors/permute.jl") include("bijectors/rational_quadratic_spline.jl") include("bijectors/named_bijector.jl") @@ -83,51 +49,42 @@ include("bijectors/utils.jl") include("bijectors/pd.jl") include("bijectors/reshape.jl") include("bijectors/corr.jl") + include("bijectors/chainrules.jl") include("bijectors/product_bijector.jl") include("bijectors/named_stacked.jl") - include("distributionsad.jl") + include("bijectors/stacked.jl") + include("bijectors/simplex.jl") + include("bijectors/equality.jl") + include("bijectors/scale.jl") - # These tests specifically check the implementation of AD backend rules. - include("ad/chainrules.jl") - include("ad/mooncake.jl") - if TEST_ENZYME - include("ad/enzyme.jl") + @testset "ForwardDiff bijector AD" begin + for c in generate_ad_testcases() + run_ad_case(c, AutoForwardDiff()) + end end - - # These tests check that AD can differentiate through Bijectors # functionality without explicit rules. - include("ad/flows.jl") - include("ad/pd.jl") - include("ad/corr.jl") - include("ad/stacked.jl") - end - - if GROUP == "All" || GROUP == "Vector" - # VectorBijectors module. - include("vector/univariate.jl") - include("vector/multivariate.jl") - include("vector/matrix.jl") - include("vector/reshaped.jl") - include("vector/cholesky.jl") - include("vector/order.jl") - include("vector/transformed.jl") - end - - if GROUP == "All" || GROUP == "VectorProduct" - # VectorBijectors module, part 2 - include("vector/product.jl") end - if GROUP == "All" || GROUP == "Doctests" - @testset "doctests" begin - Documenter.DocMeta.setdocmeta!( - Bijectors, :DocTestSetup, :(using Bijectors); recursive=true + if GROUP == "All" || GROUP == "Vector" || GROUP == "VectorProduct" + # Main-suite non-AD `VectorBijectors.test_all` coverage. `Vector` covers everything + # except product distributions, `VectorProduct` covers products only, `All` runs both. + let + product_only_tags = ( + :products, :nested_product_namedtuple, :type_unstable_products ) - 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) + selected_tags = if GROUP == "Vector" + Tuple(t for t in _VECTOR_TAGS if t ∉ product_only_tags) + elseif GROUP == "VectorProduct" + product_only_tags + else + _VECTOR_TAGS + end + + @testset "VectorBijectors test_all" begin + for c in generate_vector_testcases() + c.tag in selected_tags || continue + run_vector_case(c) + end + end end end end diff --git a/test/test_resources.jl b/test/test_resources.jl new file mode 100644 index 000000000..5409364d5 --- /dev/null +++ b/test/test_resources.jl @@ -0,0 +1,346 @@ +# Shared test resources used by `test/runtests.jl` (main suite, including the inline +# vector-loop) and the AD integration tests +# (`test/integration_tests/{enzyme,mooncake,reversediff}/main.jl`). +# +# Each test case carries a `tag::Symbol` identifying its source generator. Aggregate +# entry points return flat lists: +# - `generate_ad_testcases()` — every bijector-level AD case (ADTestCase). +# - `generate_vector_testcases()` — every distribution-level vector case (VectorTestCase). +# 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 +# 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()]`). +# +# 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. + +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 Distributions +const _D = Distributions +using FillArrays: Fill +using FiniteDifferences: central_fdm +using LinearAlgebra +using PDMats +using StableRNGs: StableRNG +using Test + +# Seed used to construct randomised AD test inputs. Stable across CI runs so that a +# backend failure can be reproduced with the same `c.arg` and so all backends see +# identical inputs. +const TESTCASE_SEED = 23 + +_testcase_rng() = StableRNG(TESTCASE_SEED) + +# Compact, readable name for a distribution. `Bijectors.VectorBijectors._name` already +# handles wrappers (Truncated, Censored, ReshapedDistribution, OrderStatistic, etc.); fall +# back to a stripped `nameof` for everything else so `VectorTestCase.name` doesn't carry +# the 1000-character `string(d)` form for nested products. +_case_name(d) = string(VectorBijectors._name(d)) + +# ===== Test case dispatch ===== + +# Per-tag generators live as methods of `_gen_testcases(::Val{:tag})`. The public +# `generate_testcases(Val(:tag))` wraps them and stamps `tag` onto every case so callers +# can filter a flat list of cases by source. +function _gen_testcases end +generate_testcases(t::Val{T}) where {T} = [_settag(c, T) for c in _gen_testcases(t)] + +# ===== ADTestCase ===== + +struct ADTestCase + name::String + func::Function + arg::Any + tag::Symbol +end +ADTestCase(name::String, func, arg) = ADTestCase(name, func, arg, :_default) + +_settag(c::ADTestCase, tag::Symbol) = ADTestCase(c.name, c.func, c.arg, tag) + +const _AD_TAGS = ( + :veccorrbijector, :veccholeskybijector, :planarlayer, :pdvecbijector, :stackedbijector +) + +""" +Return every bijector-level AD test case as a single flat list. Each case's `tag` field +identifies its source generator so integration suites can filter via `is_broken(c)`. +""" +generate_ad_testcases() = reduce(vcat, generate_testcases(Val(t)) for t in _AD_TAGS) + +# ===== ADTestCase runner ===== + +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 = gradient(f, REF_BACKEND, x) + ad_gradient = gradient(f, backend, x) + @test isapprox(ad_gradient, ref_gradient; rtol=rtol, atol=atol) +end + +function run_ad_case(c::ADTestCase, adtype; broken::Bool=false, rtol=1e-6, atol=1e-6) + @testset "$(c.name)" begin + if broken + # Evaluate the comparison anyway under @test_broken: if `gradient` throws or + # returns the wrong result the test stays broken; if the upstream bug is fixed + # 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); + rtol=rtol, + atol=atol, + ) + else + test_ad(c.func, adtype, c.arg; rtol=rtol, atol=atol) + end + end +end + +# ===== Bijector-specific AD test cases ===== + +function _gen_testcases(::Val{:veccorrbijector}) + rng = _testcase_rng() + cases = ADTestCase[] + for d in (1, 2, 4) + dist = LKJ(d, 2.0) + b = bijector(dist) + binv = inverse(b) + x = rand(rng, dist) + y = b(x) + roundtrip = let b = b, binv = binv + y_ -> sum(transform(b, binv(y_))) + end + inverse_only = let binv = binv + y_ -> sum(transform(binv, y_)) + end + push!(cases, ADTestCase("VecCorrBijector d=$d roundtrip", roundtrip, y)) + push!(cases, ADTestCase("VecCorrBijector d=$d inverse", inverse_only, y)) + end + return cases +end + +function _gen_testcases(::Val{:veccholeskybijector}) + rng = _testcase_rng() + cases = ADTestCase[] + for d in (1, 2, 4), uplo in ('U', 'L') + dist = LKJCholesky(d, 2.0, uplo) + b = bijector(dist) + binv = inverse(b) + x = rand(rng, dist) + y = b(x) + cholesky_to_triangular = + uplo == 'U' ? Bijectors.cholesky_upper : Bijectors.cholesky_lower + roundtrip = let b = b, binv = binv + y_ -> sum(transform(b, binv(y_))) + end + # `cholesky_upper`/`cholesky_lower` is needed because `sum` on a + # LinearAlgebra.Cholesky doesn't return a scalar + inverse_only = let binv = binv, f = cholesky_to_triangular + y_ -> sum(f(transform(binv, y_))) + end + push!( + cases, ADTestCase("VecCholeskyBijector d=$d uplo=$uplo roundtrip", roundtrip, y) + ) + push!( + cases, + ADTestCase("VecCholeskyBijector d=$d uplo=$uplo inverse", inverse_only, y), + ) + end + return cases +end + +function _gen_testcases(::Val{:planarlayer}) + rng = _testcase_rng() + # logpdf of a flow with a planar layer and two-dimensional inputs + f = function (θ) + 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 + g = function (θ) + 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 + # logpdf of a flow with the inverse of a planar layer and two-dimensional inputs + finv = function (θ) + 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 + ginv = function (θ) + 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 + return [ + ADTestCase("PlanarLayer logpdf vector input", f, randn(rng, 7)), + ADTestCase("PlanarLayer logpdf matrix input", g, randn(rng, 11)), + ADTestCase("PlanarLayer inverse logpdf vector input", finv, randn(rng, 7)), + ADTestCase("PlanarLayer inverse logpdf matrix input", ginv, randn(rng, 11)), + ] +end + +function _gen_testcases(::Val{:pdvecbijector}) + rng = _testcase_rng() + _topd(x) = x * x' + I + d = 4 + b = Bijectors.PDVecBijector() + binv = inverse(b) + z = randn(rng, d, d) + x = _topd(z) + y = b(x) + forward_only = let b = b, _topd = _topd, d = d + x_ -> sum(transform(b, _topd(reshape(x_, d, d)))) + end + inverse_only = let binv = binv + y_ -> sum(transform(binv, y_)) + end + inverse_chol_lower = let binv = binv + y_ -> sum(Bijectors.cholesky_lower(transform(binv, y_))) + end + inverse_chol_upper = let binv = binv + y_ -> sum(Bijectors.cholesky_upper(transform(binv, y_))) + end + return [ + ADTestCase("PDVecBijector forward", forward_only, vec(z)), + ADTestCase("PDVecBijector inverse", inverse_only, y), + ADTestCase("PDVecBijector inverse + cholesky_lower", inverse_chol_lower, y), + ADTestCase("PDVecBijector inverse + cholesky_upper", inverse_chol_upper, y), + ] +end + +function _gen_testcases(::Val{:stackedbijector}) + rng = _testcase_rng() + dist1 = Dirichlet(4, 1.0) + b1 = bijector(dist1) + dist2 = LogNormal(0.0, 1.0) + b2 = bijector(dist2) + x1 = rand(rng, dist1) + x2 = rand(rng, dist2) + y1 = b1(x1) + y2 = b2(x2) + + b_tuple = Stacked((b1, b2), (1:4, 5:5)) + binv_tuple = inverse(b_tuple) + b_vec = Stacked([b1, b2], [1:4, 5:5]) + binv_vec = inverse(b_vec) + + y = vcat(y1, [y2]) + + return [ + ADTestCase( + "StackedBijector tuple roundtrip", + let b = b_tuple, binv = binv_tuple + y_ -> sum(transform(b, binv(y_))) + end, + y, + ), + ADTestCase( + "StackedBijector tuple inverse", + let binv = binv_tuple + y_ -> sum(transform(binv, y_)) + end, + y, + ), + ADTestCase( + "StackedBijector vector roundtrip", + # Note: matches `main` — uses the tuple-form `binv` inside, not `binv_vec`. + let bvec = b_vec, binv = binv_tuple + y_ -> sum(transform(bvec, binv(y_))) + end, + y, + ), + ADTestCase( + "StackedBijector vector inverse", + let bvec_inv = binv_vec + y_ -> sum(transform(bvec_inv, y_)) + end, + y, + ), + ] +end + +# ===== VectorTestCase ===== + +struct VectorTestCase{D<:Distributions.Distribution} + name::String + dist::D + test_kwargs::NamedTuple + tag::Symbol +end + +function VectorTestCase(name::String, dist::Distributions.Distribution; kwargs...) + return VectorTestCase(name, dist, NamedTuple(kwargs), :_default) +end +function VectorTestCase(dist::Distributions.Distribution; kwargs...) + return VectorTestCase(_case_name(dist), dist; kwargs...) +end + +_settag(c::VectorTestCase, tag::Symbol) = VectorTestCase(c.name, c.dist, c.test_kwargs, tag) + +const _VECTOR_TAGS = ( + :univariates, + :multivariates, + :matrix_dists, + :lkj_matrix_dists, + :cholesky_dists, + :reshaped_dists, + :reshaped_beta_special, + :transformed_dists, + :order_orderstatistic, + :order_joint, + :order_ordered, + :products, + :nested_product_namedtuple, + :type_unstable_products, +) + +""" +Return every distribution-level vector test case as a single flat list. Each case's `tag` +field identifies its source generator so integration suites can filter via `is_broken(c)`. +""" +generate_vector_testcases() = reduce(vcat, generate_testcases(Val(t)) for t in _VECTOR_TAGS) + +# ===== VectorTestCase runner ===== + +function run_vector_case( + c::VectorTestCase, adtypes=nothing; broken_adtypes=DI.AbstractADType[], skip::Bool=false +) + if adtypes === nothing + VectorBijectors.test_all( + c.dist; c.test_kwargs..., broken_adtypes=broken_adtypes, skip=skip + ) + else + VectorBijectors.test_all( + c.dist; c.test_kwargs..., adtypes, broken_adtypes=broken_adtypes, skip=skip + ) + end + return nothing +end + +# ===== Per-category vector generators ===== + +include("vector/univariate.jl") +include("vector/multivariate.jl") +include("vector/matrix.jl") +include("vector/cholesky.jl") +include("vector/reshaped.jl") +include("vector/transformed.jl") +include("vector/order.jl") +include("vector/product.jl") diff --git a/test/vector/cholesky.jl b/test/vector/cholesky.jl index 74ea1fd41..879ae8d21 100644 --- a/test/vector/cholesky.jl +++ b/test/vector/cholesky.jl @@ -1,39 +1,12 @@ -module VBCholeskyTests - -using Distributions -using LinearAlgebra -using Test -using Bijectors.VectorBijectors -import DifferentiationInterface as DI -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake -using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse - -# Need runtime activity for some reason. -# TODO(penelopeysm): Report upstream -const adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), - DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), -] - -dists = [ - # Note: can't test LKJCholesky(1, ...) because its linked vector is length-zero and - # DifferentiationInterface trips up with empty vectors. +# Can't test LKJCholesky(1, ...) because its linked vector is length-zero and +# DifferentiationInterface trips up with empty vectors. +const cholesky_dists = [ LKJCholesky(3, 1.0, 'U'), LKJCholesky(3, 1.0, 'L'), LKJCholesky(5, 1.0, 'U'), LKJCholesky(5, 1.0, 'L'), ] -@testset "Cholesky" begin - for d in dists - VectorBijectors.test_all(d; adtypes=adtypes, expected_zero_allocs=()) - end +function _gen_testcases(::Val{:cholesky_dists}) + return [VectorTestCase(d; expected_zero_allocs=()) for d in cholesky_dists] end - -end # module VBCholeskyTests diff --git a/test/vector/matrix.jl b/test/vector/matrix.jl index 53c5c8109..1aeb3f2ec 100644 --- a/test/vector/matrix.jl +++ b/test/vector/matrix.jl @@ -1,57 +1,32 @@ -module VBMatrixTests +const _matrix_ν = 5 +const _matrix_M = [1 2 3; 4 5 6] +const _matrix_Σ = PDMats.PDMat([1 0.5; 0.5 1]) +const _matrix_Ω = PDMats.PDMat([1 0.3 0.2; 0.3 1 0.4; 0.2 0.4 1]) -using Distributions -using LinearAlgebra -using Test -using PDMats -import DifferentiationInterface as DI -using Bijectors.VectorBijectors -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -ν = 5 -M = [1 2 3; 4 5 6] -Σ = PDMats.PDMat([1 0.5; 0.5 1]) -Ω = PDMats.PDMat([1 0.3 0.2; 0.3 1 0.4; 0.2 0.4 1]) - -# TODO(penelopeysm): ReverseDiff gives wrong results when differentiating -# through VecCorrBijector. Correctness tests are disabled for now. -# https://github.com/TuringLang/Bijectors.jl/issues/434 -lkj_test_adtypes = [DI.AutoMooncake(), DI.AutoMooncakeForward()] - -# Don't check that from_linked_vec(d)(randn(...)) is in support for LKJ, -# The reason is because the inverse bijector for LKJ causes the diagonal -# entries to be not exactly 1 due to numerical precision issues. This -# should in principle be fixed, but for now we skip the test. -# https://github.com/TuringLang/Bijectors.jl/issues/435 -test_in_support(d) = !(d isa LKJ) - -matrix_dists = [ +const matrix_dists = [ MatrixNormal(2, 4), MatrixNormal(3, 5), - MatrixTDist(ν, M, Σ, Ω), + MatrixTDist(_matrix_ν, _matrix_M, _matrix_Σ, _matrix_Ω), Wishart(7, Matrix{Float64}(I, 2, 2)), Wishart(7, Matrix{Float64}(I, 4, 4)), InverseWishart(7, Matrix{Float64}(I, 2, 2)), InverseWishart(7, Matrix{Float64}(I, 4, 4)), - # See note in matrix/posdef.jl - # MatrixBeta(3, 3, 1000), - # MatrixBeta(5, 8, 1000), - LKJ(3, 1.0), - LKJ(7, 1.0), ] -@testset "Matrix distributions" begin - for d in matrix_dists - kwargs = (expected_zero_allocs=(), test_in_support=test_in_support(d)) - if d isa LKJ - VectorBijectors.test_all(d; adtypes=lkj_test_adtypes, kwargs...) - else - VectorBijectors.test_all(d; kwargs...) - end - end +const lkj_matrix_dists = [LKJ(3, 1.0), LKJ(7, 1.0)] + +function _gen_testcases(::Val{:matrix_dists}) + return [VectorTestCase(d; expected_zero_allocs=()) for d in matrix_dists] end -end # module VBMatrixTests +# LKJ is split into its own tag because ReverseDiff gives wrong results when differentiating +# through VecCorrBijector (https://github.com/TuringLang/Bijectors.jl/issues/434), so it +# only runs in the Mooncake integration suite. Don't check `from_linked_vec(d)(randn(...))` +# support — numerical precision in the inverse bijector means diagonal entries are not +# exactly 1 (https://github.com/TuringLang/Bijectors.jl/issues/435). +function _gen_testcases(::Val{:lkj_matrix_dists}) + return [ + VectorTestCase(d; expected_zero_allocs=(), test_in_support=false) for + d in lkj_matrix_dists + ] +end diff --git a/test/vector/multivariate.jl b/test/vector/multivariate.jl index b217c44d3..de3a097d2 100644 --- a/test/vector/multivariate.jl +++ b/test/vector/multivariate.jl @@ -1,38 +1,23 @@ -module VBMultivariateTests - -using Distributions -using LinearAlgebra -using Test -using Bijectors.VectorBijectors -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -multivariates = [ - # identity transforms (discrete multivariate) +const multivariates = [ Multinomial(10, [0.2, 0.5, 0.3]), - # identity transforms (continuous multivariate) MvNormal([0.0, 0.0], I), MvNormalCanon([1.0, 2.0, 3.0], [4.0 -2.0 -1.0; -2.0 5.0 -1.0; -1.0 -1.0 6.0]), MvTDist(5.0, zeros(2), Matrix(1.0I, 2, 2)), MvTDist(1.0, [1.0, -1.0, 0.5], [2.0 0.5 0.0; 0.5 3.0 0.5; 0.0 0.5 1.5]), - # broadcast exp/log MvLogNormal([0.0, 0.0], I), - # simplex distribution MvLogitNormal([1.0, 2.0], Diagonal([4.0, 5.0])), Dirichlet([2.0, 3.0, 5.0]), ] -@testset "Multivariates" begin +function _gen_testcases(::Val{:multivariates}) + cases = VectorTestCase[] for d in multivariates expected_zero_allocs = if d isa Union{Dirichlet,MvLogitNormal,MvLogNormal} (to_vec, from_vec) else (to_vec, from_vec, to_linked_vec, from_linked_vec) end - VectorBijectors.test_all(d; expected_zero_allocs=expected_zero_allocs) + push!(cases, VectorTestCase(d; expected_zero_allocs=expected_zero_allocs)) end + return cases end - -end # module VBMultivariateTests diff --git a/test/vector/order.jl b/test/vector/order.jl index 682bb1f0d..14b416bca 100644 --- a/test/vector/order.jl +++ b/test/vector/order.jl @@ -1,17 +1,4 @@ -module VBOrderTests - -using Distributions -using LinearAlgebra -using Test -using Bijectors.VectorBijectors -using Bijectors: ordered -import DifferentiationInterface as DI -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -base_dists = [ +const order_base_dists = [ Normal(), InverseGamma(2, 3), InverseGamma(2, 3) * -2, @@ -20,49 +7,61 @@ base_dists = [ DiscreteUniform(10), ] -# TODO(penelopeysm): ReverseDiff can't differentiate through JointOrderStatistics transform -# because of the heavy setindex! usage. -# https://github.com/JuliaDiff/ReverseDiff.jl/issues/43 We just avoid testing it for now. -joint_test_adtypes = [ - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=Enzyme.Forward), - DI.AutoEnzyme(; mode=Enzyme.Reverse), -] - -@testset "Order statistics" begin - for d in base_dists +function _gen_testcases(::Val{:order_orderstatistic}) + cases = VectorTestCase[] + for d in order_base_dists unvec_only = (from_vec, from_linked_vec) - VectorBijectors.test_all(OrderStatistic(d, 10, 1); expected_zero_allocs=unvec_only) - VectorBijectors.test_all(OrderStatistic(d, 10, 10); expected_zero_allocs=unvec_only) - # JointOrderStatistics is only defined for continuous distributions (technically, it - # *should* work for discrete distributions whose support is some set which has a - # total order, but Distributions.jl doesn't actually implement that). - if d isa ContinuousUnivariateDistribution - # In the unlinked case, the transform is identity. - # - # See https://github.com/TuringLang/Bijectors.jl/issues/441 for details about - # the unusually large atol. - unlinked_only = (from_vec, to_vec) - VectorBijectors.test_all( + push!( + cases, + VectorTestCase( + "order statistic $(_case_name(d)) i=1 of n=10", + OrderStatistic(d, 10, 1); + expected_zero_allocs=unvec_only, + ), + ) + push!( + cases, + VectorTestCase( + "order statistic $(_case_name(d)) i=10 of n=10", + OrderStatistic(d, 10, 10); + expected_zero_allocs=unvec_only, + ), + ) + end + return cases +end + +# JointOrderStatistics is only defined for continuous distributions. In the unlinked case +# the transform is identity. https://github.com/TuringLang/Bijectors.jl/issues/441 explains +# the unusually large `roundtrip_atol`. +function _gen_testcases(::Val{:order_joint}) + cases = VectorTestCase[] + unlinked_only = (from_vec, to_vec) + for d in order_base_dists + d isa ContinuousUnivariateDistribution || continue + push!( + cases, + VectorTestCase( + "joint order statistic $(_case_name(d)) n=4 (all ranks)", JointOrderStatistics(d, 4); expected_zero_allocs=unlinked_only, - adtypes=joint_test_adtypes, roundtrip_atol=1e-1, - ) - VectorBijectors.test_all( + ), + ) + push!( + cases, + VectorTestCase( + "joint order statistic $(_case_name(d)) n=10 ranks=2:5", JointOrderStatistics(d, 10, 2:5); expected_zero_allocs=unlinked_only, - adtypes=joint_test_adtypes, roundtrip_atol=1e-1, - ) - end + ), + ) end - - # Bijectors.ordered - VectorBijectors.test_all( - ordered(MvNormal([0.0, 1.0, 2.0], I)); expected_zero_allocs=(from_vec, to_vec) - ) + return cases end -end # module VBOrderTests +function _gen_testcases(::Val{:order_ordered}) + d = ordered(MvNormal([0.0, 1.0, 2.0], I)) + return [VectorTestCase("ordered MvNormal", d; expected_zero_allocs=(from_vec, to_vec))] +end diff --git a/test/vector/product.jl b/test/vector/product.jl index ecf5f680e..8e29d12c8 100644 --- a/test/vector/product.jl +++ b/test/vector/product.jl @@ -1,130 +1,73 @@ -module VBProductTests +const _m2 = MvNormal(zeros(2), I) +const _d2 = Dirichlet(ones(2)) +const _p1t = product_distribution(Normal(), Beta(2, 2)) +const _p2t = product_distribution(_m2, _d2) +const _p1a = product_distribution(fill(Beta(2, 2), 2)) +const _p2a = product_distribution(fill(_d2, 2)) -using Distributions -using LinearAlgebra -using FillArrays: Fill -using Test -using Bijectors.VectorBijectors -import DifferentiationInterface as DI -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake -using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse - -adtypes = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - # Need runtime activity for some reason. - # TODO(penelopeysm): Report upstream - DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const), - DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const), -] - -# Enzyme segfaults on 1.12 + Windows. -# https://github.com/EnzymeAD/Enzyme.jl/issues/2986 -if VERSION >= v"1.12-" && Sys.iswindows() - filter!(a -> !(a isa DI.AutoEnzyme), adtypes) -end - -# These are purposely chosen because the vec_length output is the same but -# linked_vec_length differs. -m2 = MvNormal(zeros(2), I) -d2 = Dirichlet(ones(2)) - -p1t = product_distribution(Normal(), Beta(2, 2)) -p2t = product_distribution(m2, d2) -p1a = product_distribution(fill(Beta(2, 2), 2)) -p2a = product_distribution(fill(d2, 2)) - -products = [ - # Tuples +# Purposely chosen so that `vec_length` agrees but `linked_vec_length` differs. +const products = [ product_distribution(Normal()), product_distribution(Normal(), Normal()), product_distribution(Normal(), Beta(2, 2)), product_distribution(Beta(2, 2), Exponential()), - product_distribution(m2, d2), - product_distribution(m2, d2, m2, d2), - # Vectors of univariate (Distributions.Product) - product_distribution(fill(Normal(), 2)), # This is actually an MvNormal in disguise + product_distribution(_m2, _d2), + product_distribution(_m2, _d2, _m2, _d2), + product_distribution(fill(Normal(), 2)), product_distribution(fill(Beta(2, 2), 2)), product_distribution([Uniform(0, 1), Uniform(1, 2), Uniform(2, 3)]), product_distribution(Fill(Uniform(1, 2), 2)), - # >1D arrays, or vectors of >1D distributions (Distributions.ProductDistribution) product_distribution(fill(Normal(), 2, 2)), product_distribution(Fill(Uniform(1, 2), 2, 2)), - product_distribution(fill(m2, 2, 2)), - product_distribution(Fill(m2, 2, 2)), - product_distribution(fill(d2, 2, 2)), - # NamedTuples + product_distribution(fill(_m2, 2, 2)), + product_distribution(Fill(_m2, 2, 2)), + product_distribution(fill(_d2, 2, 2)), product_distribution((a=Normal(), b=Beta(2, 2))), product_distribution((a=Normal(), b=Dirichlet(ones(2)))), product_distribution((a=Normal(), b=product_distribution(fill(Beta(2, 2), 2)))), - # Nested - product_distribution(fill(p1t, 2)), - product_distribution(fill(p1t, 2, 2)), - product_distribution(p2t, p2t, p2t), - product_distribution(fill(p2t, 2)), - product_distribution(fill(p2t, 2, 2)), - product_distribution(fill(p1a, 2)), - product_distribution(fill(p1a, 2, 2)), - product_distribution(p2a, p2a, p2a), - product_distribution(fill(p2a, 2)), - product_distribution(fill(p2a, 2, 2)), + product_distribution(fill(_p1t, 2)), + product_distribution(fill(_p1t, 2, 2)), + product_distribution(_p2t, _p2t, _p2t), + product_distribution(fill(_p2t, 2)), + product_distribution(fill(_p2t, 2, 2)), + product_distribution(fill(_p1a, 2)), + product_distribution(fill(_p1a, 2, 2)), + product_distribution(_p2a, _p2a, _p2a), + product_distribution(fill(_p2a, 2)), + product_distribution(fill(_p2a, 2, 2)), ] -nested_product_namedtuple = [ - # This one is very weird: on Julia 1.10 (and only on Julia 1.10), @inferred to_vec(d) - # fails, but @code_warntype to_vec(d) is completely type stable. Almost certainly a - # Julia bug or at least a weird quirk - product_distribution((a=Normal(), b=product_distribution((c=Normal(), d=Beta(2, 2))))), +# On Julia 1.10 (and only 1.10), `@inferred to_vec(d)` fails for this case even though +# `@code_warntype to_vec(d)` is type stable. Almost certainly a Julia bug. +const nested_product_namedtuple = [ + product_distribution((a=Normal(), b=product_distribution((c=Normal(), d=Beta(2, 2))))) ] -heterogeneous_products = [ - # These contain heterogeneous arrays, which means that the construction of the bijector - # is type unstable. I don't think it's possible to fix this, but someone should probably at - # least try. +# Heterogeneous arrays make bijector construction type unstable. The triple-nested tuple +# products (last two) hit Enzyme activity-inference limits; the Enzyme suite filters them. +const type_unstable_products = [ product_distribution([Normal(), Beta(2, 2), Exponential()]), product_distribution([Normal() Beta(2, 2); Exponential() Uniform(-1, 1)]), - product_distribution([m2 d2; m2 d2]), + product_distribution([_m2 _d2; _m2 _d2]), + product_distribution(_p1t, _p1t, _p1t), + product_distribution(_p1a, _p1a, _p1a), ] -enzyme_failures = [ - # These work generally but fail with Enzyme -- should probably be reported upstream - product_distribution(p1t, p1t, p1t), - product_distribution(p1a, p1a, p1a), -] - -@testset "Product distributions" begin - for d in products - VectorBijectors.test_all(d; adtypes=adtypes, expected_zero_allocs=()) - end - - for d in nested_product_namedtuple - VectorBijectors.test_all( - d; - adtypes=adtypes, - expected_zero_allocs=(), - test_construction_type_stable=(VERSION >= v"1.11-"), - ) - end - - for d in heterogeneous_products - VectorBijectors.test_all( - d; adtypes=adtypes, expected_zero_allocs=(), test_construction_type_stable=false - ) - end +function _gen_testcases(::Val{:products}) + return [VectorTestCase(d; expected_zero_allocs=()) for d in products] +end - no_enzyme_adtypes = filter(adtype -> !(adtype isa DI.AutoEnzyme), adtypes) - for d in enzyme_failures - VectorBijectors.test_all( - d; - adtypes=no_enzyme_adtypes, - expected_zero_allocs=(), - test_construction_type_stable=false, - ) - end +function _gen_testcases(::Val{:nested_product_namedtuple}) + return [ + VectorTestCase( + d; expected_zero_allocs=(), test_construction_type_stable=(VERSION >= v"1.11-") + ) for d in nested_product_namedtuple + ] end -end # module VBProductTests +function _gen_testcases(::Val{:type_unstable_products}) + return [ + VectorTestCase(d; expected_zero_allocs=(), test_construction_type_stable=false) for + d in type_unstable_products + ] +end diff --git a/test/vector/reshaped.jl b/test/vector/reshaped.jl index 64d07ce99..5dc5489e2 100644 --- a/test/vector/reshaped.jl +++ b/test/vector/reshaped.jl @@ -1,17 +1,5 @@ -module VBReshapedTests - -using Distributions -using LinearAlgebra -using Test -using Bijectors.VectorBijectors -import DifferentiationInterface as DI -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -reshaped = [ - # 0-dim array output: doesn't work because +const reshaped_default_dists = [ + # 0-dim array output is blocked by # https://github.com/JuliaStats/Distributions.jl/issues/2025 # reshape(Normal(), ()), vec(Normal()), @@ -28,28 +16,14 @@ reshaped = [ reshape(Wishart(7, Matrix{Float64}(I, 4, 4)), 1, 1, 4, 1, 4), ] -# Fails on 1.10: https://github.com/EnzymeAD/Enzyme.jl/issues/2987 -adtypes_no_enz_rvs = [ - DI.AutoReverseDiff(), - DI.AutoReverseDiff(; compile=true), - DI.AutoMooncake(), - DI.AutoMooncakeForward(), - DI.AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const), -] -reshaped_no_enzyme = [reshape(Beta(2, 2), (1, 1, 1, 1, 1))] +function _gen_testcases(::Val{:reshaped_dists}) + return [VectorTestCase(d; expected_zero_allocs=()) for d in reshaped_default_dists] +end -@testset "Reshaped distributions" begin - for d in reshaped - VectorBijectors.test_all(d; expected_zero_allocs=()) - end +# Isolated under its own tag because Enzyme Reverse fails on Julia 1.10 +# (https://github.com/EnzymeAD/Enzyme.jl/issues/2987). +const reshaped_beta_dist = reshape(Beta(2, 2), (1, 1, 1, 1, 1)) - for d in reshaped_no_enzyme - @static if VERSION >= v"1.11-" - VectorBijectors.test_all(d; expected_zero_allocs=()) - else - VectorBijectors.test_all(d; adtypes=adtypes_no_enz_rvs, expected_zero_allocs=()) - end - end +function _gen_testcases(::Val{:reshaped_beta_special}) + return [VectorTestCase(reshaped_beta_dist; expected_zero_allocs=())] end - -end # module VBReshapedTests diff --git a/test/vector/transformed.jl b/test/vector/transformed.jl index a4ee85c06..6c9ae5337 100644 --- a/test/vector/transformed.jl +++ b/test/vector/transformed.jl @@ -1,33 +1,14 @@ -module VBTransformedTests - -using Bijectors -using Distributions -using LinearAlgebra -using Test -using Bijectors.VectorBijectors -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -transformed_dists = [ - # Univariate +const transformed_dists = [ transformed(Normal(), exp), transformed(Beta(2, 3), Bijectors.Logit(0.0, 1.0)), transformed(Gamma(2, 1), elementwise(log)), - # Multivariate transformed(product_distribution(fill(Beta(2, 2), 4)), elementwise(exp)), transformed(MvNormal(zeros(3), I), Bijectors.Scale(2.0)), transformed(Dirichlet([1.0, 2.0, 3.0])), transformed(MvLogNormal(zeros(2), I), elementwise(log)), - # Matrix transformed(MatrixNormal(zeros(2, 3), I(2), I(3)), elementwise(exp)), ] -@testset "TransformedDistributions" begin - for d in transformed_dists - VectorBijectors.test_all(d; test_in_support=false) - end +function _gen_testcases(::Val{:transformed_dists}) + return [VectorTestCase(d; test_in_support=false) for d in transformed_dists] end - -end # module VBTransformedTests diff --git a/test/vector/univariate.jl b/test/vector/univariate.jl index ebc9ed40e..8ebb5e5b0 100644 --- a/test/vector/univariate.jl +++ b/test/vector/univariate.jl @@ -1,79 +1,61 @@ -module VBUnivariateTests - -using Distributions -using Test -using Bijectors.VectorBijectors -using Enzyme: Enzyme -using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff -using Mooncake: Mooncake - -univariates = [ - # Continuous - Arcsine(0, 1), # trunc - Beta(2, 2), # trunc - BetaPrime(1, 2), # pos +const univariates = [ + Arcsine(0, 1), + Beta(2, 2), + BetaPrime(1, 2), Biweight(1, 2), - Cauchy(-2, 1), # iden - # sampling from Chernoff errors randomly: - # https://github.com/JuliaStats/Distributions.jl/issues/1999 - # https://github.com/JuliaStats/Distributions.jl/pull/2000 - # Chernoff(), # iden - Chi(1), # pos - Chisq(3), # pos - Cosine(0, 1), # trunc + Cauchy(-2, 1), + Chi(1), + Chisq(3), + Cosine(0, 1), Epanechnikov(0, 1), - Erlang(7, 0.5), # pos - Exponential(0.5), # pos - FDist(10, 1), # pos - Frechet(1, 1), # trunc - Gamma(7.5, 1), # pos - GeneralizedExtremeValue(0, 1, 1), # trunc - GeneralizedPareto(0, 1, 1), # trunc - Gumbel(0, 1), # iden - InverseGamma(3, 0.5), # pos - InverseGaussian(1, 1), # pos - JohnsonSU(0.0, 1.0, 0.0, 1.0), # iden - Kolmogorov(), # pos - # KSDist(5), # can't rand - # KSOneSided(5), # can't rand - Kumaraswamy(2, 5), # trunc - Laplace(0, 4), # iden - Levy(0, 1), # trunc - Lindley(1.5), # pos - Logistic(2, 1), # iden - LogitNormal(0, 1), # trunc - LogNormal(0, 1), # pos - LogUniform(1, 10), # trunc - NoncentralBeta(2, 3, 1), # trunc - NoncentralChisq(2, 3), # pos - NoncentralF(2, 3, 1), # pos - NoncentralT(2, 3), # iden - Normal(0, 1), # iden - NormalCanon(0, 1), # iden - NormalInverseGaussian(0, 0.5, 0.2, 0.1), # iden - Pareto(1, 1), # trunc - PGeneralizedGaussian(0.2), # iden - Rayleigh(0.5), # pos - Rician(0.5, 1), # pos - Semicircle(1), # trunc - SkewedExponentialPower(0, 1, 0.7, 0.7), # iden - SkewNormal(0, 1, -1), # iden - StudentizedRange(2, 2), # pos - SymTriangularDist(0, 1), # trunc - TDist(5), # iden - TriangularDist(0, 1.5, 0.5), # trunc - Triweight(1, 1), # trunc - Uniform(0, 1), # trunc - VonMises(0.5), # trunc - Weibull(0.5, 1), # pos - truncated(Normal(); lower=0.0), # trunc - truncated(Normal(); upper=0.0), # trunc - truncated(Normal(); lower=0.0, upper=1.0), # trunc - censored(Normal(); lower=0.0), # trunc - censored(Normal(); upper=0.0), # trunc - censored(Normal(); lower=0.0, upper=1.0), # trunc - # Discrete + Erlang(7, 0.5), + Exponential(0.5), + FDist(10, 1), + Frechet(1, 1), + Gamma(7.5, 1), + GeneralizedExtremeValue(0, 1, 1), + GeneralizedPareto(0, 1, 1), + Gumbel(0, 1), + InverseGamma(3, 0.5), + InverseGaussian(1, 1), + JohnsonSU(0.0, 1.0, 0.0, 1.0), + Kolmogorov(), + Kumaraswamy(2, 5), + Laplace(0, 4), + Levy(0, 1), + Lindley(1.5), + Logistic(2, 1), + LogitNormal(0, 1), + LogNormal(0, 1), + LogUniform(1, 10), + NoncentralBeta(2, 3, 1), + NoncentralChisq(2, 3), + NoncentralF(2, 3, 1), + NoncentralT(2, 3), + Normal(0, 1), + NormalCanon(0, 1), + NormalInverseGaussian(0, 0.5, 0.2, 0.1), + Pareto(1, 1), + PGeneralizedGaussian(0.2), + Rayleigh(0.5), + Rician(0.5, 1), + Semicircle(1), + SkewedExponentialPower(0, 1, 0.7, 0.7), + SkewNormal(0, 1, -1), + StudentizedRange(2, 2), + SymTriangularDist(0, 1), + TDist(5), + TriangularDist(0, 1.5, 0.5), + Triweight(1, 1), + Uniform(0, 1), + VonMises(0.5), + Weibull(0.5, 1), + truncated(Normal(); lower=0.0), + truncated(Normal(); upper=0.0), + truncated(Normal(); lower=0.0, upper=1.0), + censored(Normal(); lower=0.0), + censored(Normal(); upper=0.0), + censored(Normal(); lower=0.0, upper=1.0), Bernoulli(0.5), BernoulliLogit(0.0), BetaBinomial(5, 2, 2), @@ -89,51 +71,44 @@ univariates = [ PoissonBinomial([0.2, 0.5, 0.3]), Skellam(2.0, 3.0), Soliton(100, 60, 0.2), - # Univariate mixture models MixtureModel([Normal(-2.0, 1.2), Normal(0.0, 1.0), Normal(3.0, 2.5)], [0.2, 0.5, 0.3]), MixtureModel([Normal(0, 1)], [1.0]), MixtureModel([Beta(2, 2), Beta(5, 1)], [0.5, 0.5]), - # Shifted / scaled distributions - # Identity Logistic() + 2, Logistic() - 2, Logistic() * 3, Logistic() * -3, - # Positive Gamma(2, 3) + 2, Gamma(2, 3) - 2, Gamma(2, 3) * 3, Gamma(2, 3) * -3, - # Truncated Beta(2, 5) + 2, Beta(2, 5) - 2, Beta(2, 5) * 3, Beta(2, 5) * -3, - # Some extra tests just for good measure truncated(Beta(2, 5); lower=0.2, upper=0.8), truncated(Beta(2, 5) * -4; lower=-3.0, upper=-1.0), ] -heterogeneous_mixtures = [ - # Because of the abstract eltype, there's some loss of inference somewhere that causes - # allocations on < 1.12, so for these we skip the zero-allocation tests. +# Abstract eltype causes a loss of inference that allocates on < 1.12; skip the +# zero-allocation tests there. +const heterogeneous_mixtures = [ MixtureModel(Union{Normal,Exponential}[Normal(0, 1), Exponential(1)], [0.4, 0.6]), MixtureModel(Union{Gamma,Exponential}[Gamma(2, 1), Exponential(3)], [0.5, 0.5]), ] -@testset "Univariates" begin +function _gen_testcases(::Val{:univariates}) + cases = VectorTestCase[] for d in univariates - VectorBijectors.test_all(d; expected_zero_allocs=(from_vec, from_linked_vec)) + push!(cases, VectorTestCase(d; expected_zero_allocs=(from_vec, from_linked_vec))) + end + expected_alloc_for_hm = @static if VERSION >= v"1.12-" + (from_vec, from_linked_vec) + else + () end - for d in heterogeneous_mixtures - expected_zero_allocs = @static if VERSION >= v"1.12-" - (from_vec, from_linked_vec) - else - () - end - VectorBijectors.test_all(d; expected_zero_allocs=expected_zero_allocs) + push!(cases, VectorTestCase(d; expected_zero_allocs=expected_alloc_for_hm)) end + return cases end - -end # module VBUnivariateTests