Test Enzyme and reexport ADTypes.AutoEnzyme

src/Turing.jl

@@ -106,6 +106,7 @@ export @model,                 # modelling
     AutoForwardDiff,        # ADTypes
     AutoReverseDiff,
     AutoZygote,
+    AutoEnzyme,
     AutoMooncake,
     setprogress!,           # debugging
     Flat,

src/essential/Essential.jl

@@ -11,14 +11,16 @@ using Bijectors: PDMatDistribution
 using AdvancedVI
 using StatsFuns: logsumexp, softmax
 @reexport using DynamicPPL
-using ADTypes: ADTypes, AutoForwardDiff, AutoReverseDiff, AutoZygote, AutoMooncake
+using ADTypes:
+    ADTypes, AutoForwardDiff, AutoEnzyme, AutoReverseDiff, AutoZygote, AutoMooncake
 
 using AdvancedPS: AdvancedPS
 
 include("container.jl")
 
 export @model,
     @varname,
+    AutoEnzyme,
     AutoForwardDiff,
     AutoZygote,
     AutoReverseDiff,

test/Project.toml

@@ -9,6 +9,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
@@ -45,6 +46,7 @@ Clustering = "0.14, 0.15"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 DynamicHMC = "2.1.6, 3.0"
+Enzyme = "0.12, 0.13"
 DynamicPPL = "0.29"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10"
@@ -70,4 +72,3 @@ StatsBase = "0.33, 0.34"
 StatsFuns = "0.9.5, 1"
 TimerOutputs = "0.5"
 Zygote = "0.5.4, 0.6"
-julia = "1.3"

test/mcmc/Inference.jl

@@ -7,6 +7,7 @@ using Distributions: Bernoulli, Beta, InverseGamma, Normal
 using Distributions: sample
 import DynamicPPL
 using DynamicPPL: Sampler, getlogp
+import Enzyme
 import ForwardDiff
 using LinearAlgebra: I
 import MCMCChains
@@ -17,70 +18,55 @@ using Test: @test, @test_throws, @testset
 using Turing
 
 @testset "Testing inference.jl with $adbackend" for adbackend in ADUtils.adbackends
-    # Only test threading if 1.3+.
-    if VERSION > v"1.2"
-        @testset "threaded sampling" begin
-            # Test that chains with the same seed will sample identically.
-            @testset "rng" begin
-                model = gdemo_default
-
-                # multithreaded sampling with PG causes segfaults on Julia 1.5.4
-                # https://github.com/TuringLang/Turing.jl/issues/1571
-                samplers = @static if VERSION <= v"1.5.3" || VERSION >= v"1.6.0"
-                    (
-                        HMC(0.1, 7; adtype=adbackend),
-                        PG(10),
-                        IS(),
-                        MH(),
-                        Gibbs(PG(3, :s), HMC(0.4, 8, :m; adtype=adbackend)),
-                        Gibbs(HMC(0.1, 5, :s; adtype=adbackend), ESS(:m)),
-                    )
-                else
-                    (
-                        HMC(0.1, 7; adtype=adbackend),
-                        IS(),
-                        MH(),
-                        Gibbs(HMC(0.1, 5, :s; adtype=adbackend), ESS(:m)),
-                    )
-                end
-                for sampler in samplers
-                    Random.seed!(5)
-                    chain1 = sample(model, sampler, MCMCThreads(), 1000, 4)
+    @testset "threaded sampling" begin
+        # Test that chains with the same seed will sample identically.
+        @testset "rng" begin
+            model = gdemo_default
+
+            samplers = (
+                HMC(0.1, 7; adtype=adbackend),
+                PG(10),
+                IS(),
+                MH(),
+                Gibbs(PG(3, :s), HMC(0.4, 8, :m; adtype=adbackend)),
+                Gibbs(HMC(0.1, 5, :s; adtype=adbackend), ESS(:m)),
+            )
+            for sampler in samplers
+                Random.seed!(5)
+                chain1 = sample(model, sampler, MCMCThreads(), 1000, 4)
 
-                    Random.seed!(5)
-                    chain2 = sample(model, sampler, MCMCThreads(), 1000, 4)
+                Random.seed!(5)
+                chain2 = sample(model, sampler, MCMCThreads(), 1000, 4)
 
-                    @test chain1.value == chain2.value
-                end
+                @test chain1.value == chain2.value
+            end
 
-                # Should also be stable with am explicit RNG
-                seed = 5
-                rng = Random.MersenneTwister(seed)
-                for sampler in samplers
-                    Random.seed!(rng, seed)
-                    chain1 = sample(rng, model, sampler, MCMCThreads(), 1000, 4)
+            # Should also be stable with am explicit RNG
+            seed = 5
+            rng = Random.MersenneTwister(seed)
+            for sampler in samplers
+                Random.seed!(rng, seed)
+                chain1 = sample(rng, model, sampler, MCMCThreads(), 1000, 4)
 
-                    Random.seed!(rng, seed)
-                    chain2 = sample(rng, model, sampler, MCMCThreads(), 1000, 4)
+                Random.seed!(rng, seed)
+                chain2 = sample(rng, model, sampler, MCMCThreads(), 1000, 4)
 
-                    @test chain1.value == chain2.value
-                end
+                @test chain1.value == chain2.value
             end
+        end
 
-            # Smoke test for default sample call.
-            Random.seed!(100)
-            chain = sample(
-                gdemo_default, HMC(0.1, 7; adtype=adbackend), MCMCThreads(), 1000, 4
-            )
-            check_gdemo(chain)
+        # Smoke test for default sample call.
+        Random.seed!(100)
+        chain = sample(gdemo_default, HMC(0.1, 7; adtype=adbackend), MCMCThreads(), 1000, 4)
+        check_gdemo(chain)
 
-            # run sampler: progress logging should be disabled and
-            # it should return a Chains object
-            sampler = Sampler(HMC(0.1, 7; adtype=adbackend), gdemo_default)
-            chains = sample(gdemo_default, sampler, MCMCThreads(), 1000, 4)
-            @test chains isa MCMCChains.Chains
-        end
+        # run sampler: progress logging should be disabled and
+        # it should return a Chains object
+        sampler = Sampler(HMC(0.1, 7; adtype=adbackend), gdemo_default)
+        chains = sample(gdemo_default, sampler, MCMCThreads(), 1000, 4)
+        @test chains isa MCMCChains.Chains
     end
+
     @testset "chain save/resume" begin
         Random.seed!(1234)
 
@@ -411,51 +397,54 @@ using Turing
         alg = HMC(0.01, 5; adtype=adbackend)
         res = sample(vdemo2(randn(D, 100)), alg, 250)
 
-        # Vector assumptions
-        N = 10
-        alg = HMC(0.2, 4; adtype=adbackend)
+        # TODO(mhauru) Type unstable getfield of tuple not supported in Enzyme yet
+        if !(adbackend isa AutoEnzyme)
+            # Vector assumptions
+            N = 10
+            alg = HMC(0.2, 4; adtype=adbackend)
 
-        @model function vdemo3()
-            x = Vector{Real}(undef, N)
-            for i in 1:N
-                x[i] ~ Normal(0, sqrt(4))
+            @model function vdemo3()
+                x = Vector{Real}(undef, N)
+                for i in 1:N
+                    x[i] ~ Normal(0, sqrt(4))
+                end
             end
-        end
 
-        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
+            t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
 
-        # Test for vectorize UnivariateDistribution
-        @model function vdemo4()
-            x = Vector{Real}(undef, N)
-            @. x ~ Normal(0, 2)
-        end
+            # Test for vectorize UnivariateDistribution
+            @model function vdemo4()
+                x = Vector{Real}(undef, N)
+                @. x ~ Normal(0, 2)
+            end
 
-        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
+            t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
 
-        @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)
+            @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)
 
-        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
+            t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
 
-        println("Time for")
-        println("  Loop : ", t_loop)
-        println("  Vec  : ", t_vec)
-        println("  Mv   : ", t_mv)
+            println("Time for")
+            println("  Loop : ", t_loop)
+            println("  Vec  : ", t_vec)
+            println("  Mv   : ", t_mv)
 
-        # Transformed test
-        @model function vdemo6()
-            x = Vector{Real}(undef, N)
-            @. x ~ InverseGamma(2, 3)
-        end
+            # Transformed test
+            @model function vdemo6()
+                x = Vector{Real}(undef, N)
+                @. x ~ InverseGamma(2, 3)
+            end
 
-        sample(vdemo6(), alg, 1000)
+            sample(vdemo6(), alg, 1000)
 
-        N = 3
-        @model function vdemo7()
-            x = Array{Real}(undef, N, N)
-            @. x ~ [InverseGamma(2, 3) for i in 1:N]
-        end
+            N = 3
+            @model function vdemo7()
+                x = Array{Real}(undef, N, N)
+                @. x ~ [InverseGamma(2, 3) for i in 1:N]
+            end
 
-        sample(vdemo7(), alg, 1000)
+            sample(vdemo7(), alg, 1000)
+        end
     end
     @testset "vectorization .~" begin
         @model function vdemo1(x)
@@ -478,50 +467,53 @@ using Turing
         alg = HMC(0.01, 5; adtype=adbackend)
         res = sample(vdemo2(randn(D, 100)), alg, 250)
 
-        # Vector assumptions
-        N = 10
-        alg = HMC(0.2, 4; adtype=adbackend)
+        # TODO(mhauru) Type unstable getfield of tuple not supported in Enzyme yet
+        if !(adbackend isa AutoEnzyme)
+            # Vector assumptions
+            N = 10
+            alg = HMC(0.2, 4; adtype=adbackend)
 
-        @model function vdemo3()
-            x = Vector{Real}(undef, N)
-            for i in 1:N
-                x[i] ~ Normal(0, sqrt(4))
+            @model function vdemo3()
+                x = Vector{Real}(undef, N)
+                for i in 1:N
+                    x[i] ~ Normal(0, sqrt(4))
+                end
             end
-        end
 
-        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
+            t_loop = @elapsed res = sample(vdemo3(), alg, 1000)
 
-        # Test for vectorize UnivariateDistribution
-        @model function vdemo4()
-            x = Vector{Real}(undef, N)
-            return x .~ Normal(0, 2)
-        end
+            # Test for vectorize UnivariateDistribution
+            @model function vdemo4()
+                x = Vector{Real}(undef, N)
+                return x .~ Normal(0, 2)
+            end
 
-        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
+            t_vec = @elapsed res = sample(vdemo4(), alg, 1000)
 
-        @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)
+            @model vdemo5() = x ~ MvNormal(zeros(N), 4 * I)
 
-        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
+            t_mv = @elapsed res = sample(vdemo5(), alg, 1000)
 
-        println("Time for")
-        println("  Loop : ", t_loop)
-        println("  Vec  : ", t_vec)
-        println("  Mv   : ", t_mv)
+            println("Time for")
+            println("  Loop : ", t_loop)
+            println("  Vec  : ", t_vec)
+            println("  Mv   : ", t_mv)
 
-        # Transformed test
-        @model function vdemo6()
-            x = Vector{Real}(undef, N)
-            return x .~ InverseGamma(2, 3)
-        end
+            # Transformed test
+            @model function vdemo6()
+                x = Vector{Real}(undef, N)
+                return x .~ InverseGamma(2, 3)
+            end
 
-        sample(vdemo6(), alg, 1000)
+            sample(vdemo6(), alg, 1000)
 
-        @model function vdemo7()
-            x = Array{Real}(undef, N, N)
-            return x .~ [InverseGamma(2, 3) for i in 1:N]
-        end
+            @model function vdemo7()
+                x = Array{Real}(undef, N, N)
+                return x .~ [InverseGamma(2, 3) for i in 1:N]
+            end
 
-        sample(vdemo7(), alg, 1000)
+            sample(vdemo7(), alg, 1000)
+        end
     end
     @testset "Type parameters" begin
         N = 10

test/mcmc/abstractmcmc.jl

@@ -5,6 +5,7 @@ using AdvancedMH: AdvancedMH
 using Distributions: sample
 using Distributions.FillArrays: Zeros
 using DynamicPPL: DynamicPPL
+import Enzyme
 using ForwardDiff: ForwardDiff
 using LinearAlgebra: I
 using LogDensityProblems: LogDensityProblems

test/mcmc/gibbs.jl

@@ -5,6 +5,7 @@ using ..NumericalTests: check_MoGtest_default, check_gdemo, check_numerical
 import ..ADUtils
 using Distributions: InverseGamma, Normal
 using Distributions: sample
+import Enzyme
 using ForwardDiff: ForwardDiff
 using Random: Random
 using ReverseDiff: ReverseDiff

test/mcmc/gibbs_conditional.jl

@@ -5,6 +5,7 @@ using ..NumericalTests: check_gdemo, check_numerical
 import ..ADUtils
 using Clustering: Clustering
 using Distributions: Categorical, InverseGamma, Normal, sample
+import Enzyme
 using ForwardDiff: ForwardDiff
 using LinearAlgebra: Diagonal, I
 using Random: Random