TuringLang · devmotion · Sep 28, 2022 · Nov 7, 2022 · Nov 12, 2022 · Dec 23, 2022
diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
@@ -8,7 +8,4 @@ ignore = [
     # https://github.com/TuringLang/Turing.jl/pull/2328/files
     "src/experimental/gibbs.jl",
     "test/experimental/gibbs.jl",
-    # https://github.com/TuringLang/Turing.jl/pull/1887 # Enzyme PR
-    "test/mcmc/hmc.jl",
-    "test/mcmc/sghmc.jl",
 ]
diff --git a/.github/workflows/DocsNav.yml b/.github/workflows/DocsNav.yml
@@ -32,13 +32,13 @@ jobs:
 
           # Define the URL of the navbar to be used
           NAVBAR_URL="https://raw.githubusercontent.com/TuringLang/turinglang.github.io/main/assets/scripts/TuringNavbar.html"
-         
+
           # Update all HTML files in the current directory (gh-pages root)
           ./insert_navbar.sh . $NAVBAR_URL
-         
+
           # Remove the insert_navbar.sh file
           rm insert_navbar.sh
-         
+
           # Check if there are any changes
           if [[ -n $(git status -s) ]]; then
             git add .

diff --git a/src/Turing.jl b/src/Turing.jl
@@ -106,6 +106,7 @@ export @model,                 # modelling
     AutoForwardDiff,        # ADTypes
     AutoReverseDiff,
     AutoZygote,
+    AutoEnzyme,
     AutoMooncake,
     setprogress!,           # debugging
     Flat,

diff --git a/src/essential/Essential.jl b/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,

diff --git a/src/mcmc/mh.jl b/src/mcmc/mh.jl
@@ -54,7 +54,7 @@ Specifying a single distribution implies the use of static MH:
 
 ```julia
 # Use a static proposal for s² (which happens to be the same
-# as the prior) and a static proposal for m (note that this 
+# as the prior) and a static proposal for m (note that this
 # isn't a random walk proposal).
 chain = sample(
     gdemo(1.5, 2.0),

diff --git a/test/Project.toml b/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"
diff --git a/test/mcmc/Inference.jl b/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
@@ -411,51 +412,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 +482,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