Test Enzyme and reexport ADTypes.AutoEnzyme

.github/workflows/Tests.yml

@@ -26,15 +26,15 @@ jobs:
           - "mcmc/ess.jl"
           - "--skip essential/ad.jl mcmc/gibbs.jl mcmc/hmc.jl mcmc/abstractmcmc.jl mcmc/Inference.jl experimental/gibbs.jl mcmc/ess.jl"
         version:
-          - '1.7'
+          #- '1.7'  TODO(mhauru): Temporarily disabled for Enzyme
           - '1'
         os:
           - ubuntu-latest
-          - windows-latest
-          - macOS-latest
+          #- windows-latest  TODO(mhauru): Temporarily disabled for Enzyme
+          #- macOS-latest  TODO(mhauru): Temporarily disabled for Enzyme
         arch:
           - x64
-          - x86
+          #- x86  TODO(mhauru): Temporarily disabled for Enzyme
         num_threads:
           - 1
           - 2

src/Turing.jl

@@ -106,6 +106,7 @@ export @model,                 # modelling
     AutoForwardDiff,        # ADTypes
     AutoReverseDiff,
     AutoZygote,
+    AutoEnzyme,
     AutoTracker,
     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, AutoTracker, AutoReverseDiff, AutoZygote
+using ADTypes:
+    ADTypes, AutoForwardDiff, AutoEnzyme, AutoTracker, AutoReverseDiff, AutoZygote
 
 using AdvancedPS: AdvancedPS
 
 include("container.jl")
 
 export @model,
     @varname,
+    AutoEnzyme,
     AutoForwardDiff,
     AutoTracker,
     AutoZygote,

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"
 DynamicPPL = "0.28"
 FiniteDifferences = "0.10.8, 0.11, 0.12"
 ForwardDiff = "0.10.12 - 0.10.32, 0.10"
@@ -70,4 +72,3 @@ StatsFuns = "0.9.5, 1"
 TimerOutputs = "0.5"
 Tracker = "0.2.11"
 Zygote = "0.5.4, 0.6"
-julia = "1.3"

test/mcmc/Inference.jl

@@ -6,6 +6,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
@@ -14,7 +15,13 @@ import ReverseDiff
 using Test: @test, @test_throws, @testset
 using Turing
 
-@testset "Testing inference.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+Enzyme.API.typeWarning!(false)
+
+# Enable runtime activity (workaround)
+Enzyme.API.runtimeActivity!(true)
+
+# @testset "Testing inference.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing inference.jl with $adbackend" for adbackend in (AutoEnzyme(),)
     # Only test threading if 1.3+.
     if VERSION > v"1.2"
         @testset "threaded sampling" begin
@@ -367,6 +374,8 @@ using Turing
         alg = Gibbs(HMC(0.2, 3, :m; adtype=adbackend), PG(10, :s))
         chn = sample(gdemo_default, alg, 1000)
     end
+    # Type unstable getfield of tuple not supported in Enzyme yet
+    if adbackend != AutoEnzyme()
     @testset "vectorization @." begin
         # https://github.com/FluxML/Tracker.jl/issues/119
         @model function vdemo1(x)
@@ -549,6 +558,7 @@ using Turing
         vdemo3kw(; T) = vdemo3(T)
         sample(vdemo3kw(; T=DynamicPPL.TypeWrap{Vector{Float64}}()), alg, 250)
     end
+    end
 
     @testset "names_values" begin
         ks, xs = Turing.Inference.names_values([

test/mcmc/hmc.jl

@@ -6,6 +6,7 @@ using ..NumericalTests: check_gdemo, check_numerical
 using Distributions: Bernoulli, Beta, Categorical, Dirichlet, Normal, Wishart, sample
 import DynamicPPL
 using DynamicPPL: Sampler
+import Enzyme
 import ForwardDiff
 using HypothesisTests: ApproximateTwoSampleKSTest, pvalue
 import ReverseDiff
@@ -16,7 +17,14 @@ using StatsFuns: logistic
 using Test: @test, @test_logs, @testset
 using Turing
 
-@testset "Testing hmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+# Disable Enzyme warnings
+Enzyme.API.typeWarning!(false)
+
+# Enable runtime activity (workaround)
+Enzyme.API.runtimeActivity!(true)
+
+# @testset "Testing hmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing hmc.jl with $adbackend" for adbackend in (AutoEnzyme(),)
     # Set a seed
     rng = StableRNG(123)
     @testset "constrained bounded" begin
@@ -103,7 +111,7 @@ using Turing
         alpha = 0.16                  # regularizatin term
         var_prior = sqrt(1.0 / alpha) # variance of the Gaussian prior
 
-        @model function bnn(ts)
+        @model function bnn(ts, var_prior)
             b1 ~ MvNormal([0. ;0.; 0.],
                 [var_prior 0. 0.; 0. var_prior 0.; 0. 0. var_prior])
             w11 ~ MvNormal([0.; 0.], [var_prior 0.; 0. var_prior])
@@ -121,7 +129,7 @@ using Turing
         end
 
         # Sampling
-        chain = sample(rng, bnn(ts), HMC(0.1, 5; adtype=adbackend), 10)
+        chain = sample(rng, bnn(ts, var_prior), HMC(0.1, 5; adtype=adbackend), 10)
     end
 
     @testset "hmcda inference" begin

test/mcmc/sghmc.jl

@@ -3,14 +3,22 @@ module SGHMCTests
 using ..Models: gdemo_default
 using ..NumericalTests: check_gdemo
 using Distributions: sample
+import Enzyme
 import ForwardDiff
 using LinearAlgebra: dot
 import ReverseDiff
 using StableRNGs: StableRNG
 using Test: @test, @testset
 using Turing
 
-@testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+# Disable Enzyme warnings
+Enzyme.API.typeWarning!(false)
+
+# Enable runtime activity (workaround)
+Enzyme.API.runtimeActivity!(true)
+
+# @testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing sghmc.jl with $adbackend" for adbackend in (AutoEnzyme(),)
     @testset "sghmc constructor" begin
         alg = SGHMC(; learning_rate=0.01, momentum_decay=0.1, adtype=adbackend)
         @test alg isa SGHMC
@@ -36,7 +44,8 @@ using Turing
     end
 end
 
-@testset "Testing sgld.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+# @testset "Testing sgld.jl with $adbackend" for adbackend in (AutoForwardDiff(; chunksize=0), AutoReverseDiff(; compile=false))
+@testset "Testing sgld.jl with $adbackend" for adbackend in (AutoEnzyme(),)
     @testset "sgld constructor" begin
         alg = SGLD(; stepsize=PolynomialStepsize(0.25), adtype=adbackend)
         @test alg isa SGLD

test/runtests.jl

@@ -1,3 +1,6 @@
+import Pkg
+Pkg.add(Pkg.PackageSpec(; url="https://github.com/simsurace/Enzyme.jl.git", rev="fix-cholesky"))
+
 include("test_utils/SelectiveTests.jl")
 using .SelectiveTests: isincluded, parse_args
 using Pkg