diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 5093cf3a5a..522e59165d 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -239,6 +239,7 @@ jobs: - ubuntu-latest test: - DynamicExpressions + - Turing steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 @@ -248,8 +249,8 @@ jobs: - uses: julia-actions/julia-buildpkg@v1 - name: "Run tests" run: | - julia --color=yes --project=test/integration -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")]); Pkg.instantiate()' - julia --color=yes --project=test/integration --threads=auto --check-bounds=yes test/integration/${{ matrix.test }}.jl + julia --color=yes --project=test/integration/${{ matrix.test }} -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "lib/EnzymeCore")]); Pkg.instantiate()' + julia --color=yes --project=test/integration/${{ matrix.test }} --threads=auto --check-bounds=yes test/integration/${{ matrix.test }}/runtests.jl shell: bash docs: name: Documentation diff --git a/test/integration/Project.toml b/test/integration/DynamicExpressions/Project.toml similarity index 100% rename from test/integration/Project.toml rename to test/integration/DynamicExpressions/Project.toml diff --git a/test/integration/DynamicExpressions.jl b/test/integration/DynamicExpressions/runtests.jl similarity index 100% rename from test/integration/DynamicExpressions.jl rename to test/integration/DynamicExpressions/runtests.jl diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml new file mode 100644 index 0000000000..e9ef3697f4 --- /dev/null +++ b/test/integration/Turing/Project.toml @@ -0,0 +1,21 @@ +[deps] +AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" +AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" +Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" +FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1" + +[compat] +AbstractMCMC = "=5.3.0" +AbstractPPL = "=0.8.4" +Bijectors = "=0.13.18" +Distributions = "=0.25.111" +DynamicPPL = "=0.28.4" +FiniteDifferences = "0.12.32" +LogDensityProblems = "=2.1.1" +LogDensityProblemsAD = "=1.10.0" diff --git a/test/integration/Turing/runtests.jl b/test/integration/Turing/runtests.jl new file mode 100644 index 0000000000..7435968961 --- /dev/null +++ b/test/integration/Turing/runtests.jl @@ -0,0 +1,99 @@ +module TuringIntegrationTests + +using Distributions: Distributions +using DynamicPPL: DynamicPPL +using Enzyme: Enzyme +using FiniteDifferences: FiniteDifferences +using LinearAlgebra: LinearAlgebra +using Random: randn +using Test: @test, @testset + +# TODO(mhauru) Could we at some point make do without this? +Enzyme.API.runtimeActivity!(true) + +""" +Turn a Turing model, possibly with given example values, into a log density function and a +random value that it can be evaluated at. +""" +function build_turing_problem(model) + ctx = DynamicPPL.DefaultContext() + vi = DynamicPPL.VarInfo(model) + vi_linked = DynamicPPL.link(vi, model) + ldp = DynamicPPL.LogDensityFunction(vi_linked, model, ctx) + test_function = Base.Fix1(DynamicPPL.LogDensityProblems.logdensity, ldp) + d = DynamicPPL.LogDensityProblems.dimension(ldp) + return test_function, randn(d) +end + +""" +Test Enzyme.gradient, both Forward and Reverse mode, against FiniteDifferences.grad, for a +given function f and argument x. +""" +function test_grad(f, x; rtol=1e-6, atol=1e-6) + finitediff = FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), f, x)[1] + # TODO(mhauru) The Val(1) works around https://github.com/EnzymeAD/Enzyme.jl/issues/1807 + @test( + collect(Enzyme.gradient(Enzyme.Forward, Enzyme.Const(f), x, Val(1))) ≈ finitediff, + rtol = rtol, + atol = atol + ) + @test( + Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(f), x) ≈ finitediff, + rtol = rtol, atol = atol + ) + return nothing +end + +# Turing models to test with. These come from Turing's test suite. +models = collect(DynamicPPL.TestUtils.DEMO_MODELS) + +# Add some other models that use features that have previously been problematic for Enzyme. + +DynamicPPL.@model function MvDirichletWithManualAccumulation(w, doc) + β ~ DynamicPPL.filldist(Distributions.Dirichlet([1.0, 1.0]), 2) + log_product = log.(β) + DynamicPPL.@addlogprob! sum(log_product[CartesianIndex.(w, doc)]) +end + +push!(models, MvDirichletWithManualAccumulation([1, 1, 1, 1], [1, 1, 2, 2])) + +DynamicPPL.@model function demo_lkjchol(d::Int=2) + x ~ DynamicPPL.LKJCholesky(d, 1.0) + return (x=x,) +end + +push!(models, demo_lkjchol()) + +DynamicPPL.@model function hmcmatrixsup() + return v ~ Distributions.Wishart(7, [1 0.5; 0.5 1]) +end + +push!(models, hmcmatrixsup()) + +DynamicPPL.@model function mvnormal_with_transpose(x=transpose([1.5 2.0;])) + m ~ Distributions.MvNormal(LinearAlgebra.Diagonal([1.0, 1.0])) + x .~ Distributions.MvNormal(m, LinearAlgebra.Diagonal([1.0, 1.0])) + return nothing +end + +push!(models, mvnormal_with_transpose()) + +DynamicPPL.@model function mvnorm_with_argtype(::Type{TV}=Matrix{Float64}) where {TV} + P0 = vcat([0.1 0.0], [0.0 0.1]) + x = TV(undef, 2, 2) + fill!(x, zero(eltype(x))) + x[:, 2] ~ Distributions.MvNormal(x[:, 1], P0) + return nothing +end + +push!(models, mvnorm_with_argtype()) + +# Test each model in turn, checking Enzyme's gradient against FiniteDifferences. +@testset "Turing integration tests" begin + @testset "$(typeof(model.f))" for model in models + f, x = build_turing_problem(model) + test_grad(f, x) + end +end + +end