Updated to use ADTypes.jl (#17)

* updated to use ADTypes

* compat entry for ADTypes

* fixed tests

* Zygote failing sometimes results in a Zygote.CompileError

* bumped Turing.jl bound in docs

Project.toml

@@ -1,11 +1,13 @@
 name = "TuringBenchmarking"
 uuid = "0db1332d-5c25-4deb-809f-459bc696f94f"
 authors = ["Tor Erlend Fjelde <tor.erlend95@gmail.com> and contributors"]
-version = "0.3.4"
+version = "0.4.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -22,6 +24,7 @@ BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a"
 TuringBenchmarkingBridgeStanExt = "BridgeStan"
 
 [compat]
+ADTypes = "0.2"
 BenchmarkTools = "1"
 BridgeStan = "2"
 DynamicPPL = "0.23, 0.24"
@@ -30,7 +33,7 @@ LogDensityProblemsAD = "1"
 PrettyTables = "2"
 Requires = "1"
 ReverseDiff = "1.14"
-Turing = "0.29"
+Turing = "0.30"
 Zygote = "0.6"
 julia = "1.6"
 

docs/Project.toml

@@ -4,4 +4,4 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
 Documenter = "1"
-Turing = "0.29"
+Turing = "0.30"

src/TuringBenchmarking.jl

@@ -6,12 +6,15 @@ using BenchmarkTools
 using LogDensityProblems
 using LogDensityProblemsAD
 
+using Turing
+using ADTypes
+
 using PrettyTables: PrettyTables
 
-using Turing
-using Turing.Essential: ForwardDiffAD, TrackerAD, ReverseDiffAD, ZygoteAD
 using DynamicPPL: DynamicPPL
 
+# Load some the default backends to trigger conditional loading.
+using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
 using Zygote: Zygote
 
@@ -23,28 +26,29 @@ export benchmark_model, make_turing_suite, BenchmarkTools, @tagged
 
 # Don't include `TrackerAD` because it's never going to win.
 const DEFAULT_ADBACKENDS = [
-    ForwardDiffAD{Turing.Essential.CHUNKSIZE[]}(), # chunksize=40
-    ReverseDiffAD{false}(), # rdcache=false
-    ReverseDiffAD{true}(),  # rdcache=false
-    ZygoteAD(),
+    AutoForwardDiff(chunksize=0),
+    AutoReverseDiff(false),
+    AutoReverseDiff(true),
+    AutoZygote(),
 ]
 
-backend_label(::ForwardDiffAD) = "ForwardDiff"
-backend_label(::ReverseDiffAD{false}) = "ReverseDiff"
-backend_label(::ReverseDiffAD{true}) = "ReverseDiff[compiled]"
-backend_label(::ZygoteAD) = "Zygote"
-backend_label(::TrackerAD) = "Tracker"
+backend_label(::AutoForwardDiff) = "ForwardDiff"
+function backend_label(ad::AutoReverseDiff)
+    "ReverseDiff" * (ad.compile ? " [compiled]" : "")
+end
+backend_label(::AutoZygote) = "Zygote"
+backend_label(::AutoTracker) = "Tracker"
 
 const SYMBOL_TO_BACKEND = Dict(
-    :forwarddiff => ForwardDiffAD{Turing.Essential.CHUNKSIZE[]}(),
-    :reversediff => ReverseDiffAD{false}(),
-    :reversediff_compiled => ReverseDiffAD{true}(),
-    :zygote => ZygoteAD(),
-    :tracker => TrackerAD(),
+    :forwarddiff => AutoForwardDiff(chunksize=0),
+    :reversediff => AutoReverseDiff(false),
+    :reversediff_compiled => AutoReverseDiff(true),
+    :zygote => AutoZygote(),
+    :tracker => AutoTracker(),
 )
 
 to_backend(x) = error("Unknown backend: $x")
-to_backend(x::Turing.Essential.ADBackend) = x
+to_backend(x::ADTypes.AbstractADType) = x
 function to_backend(x::Union{AbstractString,Symbol})
     k = Symbol(lowercase(string(x)))
     haskey(SYMBOL_TO_BACKEND, k) || error("Unknown backend: $x")
@@ -106,8 +110,8 @@ end
 Create default benchmark suite for `model`.
 
 # Keyword arguments
-- `adbackends`: a collection of adbackends to use, specified either as a
- `Turing.Essential.ADBackend` or using a `Symbol`. Defaults to `$(DEFAULT_ADBACKENDS)`.
+- `adbackends`: a collection of adbackends to use, specified either as a type from
+ ADTypes.jl or using a `Symbol`. Defaults to `$(DEFAULT_ADBACKENDS)`.
 - `run_once=true`: if `true`, the body of each benchmark will be run once to avoid
   compilation to be included in the timings (this may occur if compilation runs
   longer than the allowed time limit).

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

test/runtests.jl

@@ -2,6 +2,7 @@ using TuringBenchmarking
 using BenchmarkTools
 using Turing
 using Test
+using ADTypes
 
 using Zygote: Zygote
 using ReverseDiff: ReverseDiff
@@ -67,12 +68,12 @@ ADBACKENDS = TuringBenchmarking.DEFAULT_ADBACKENDS
                 adbackend_string = "$(adbackend)"
                 results_backend = results[@tagged adbackend_string]
                 # Each AD backend should have two results.
-                @test length(leaves(results_backend)) == 2
+                @test length(BenchmarkTools.leaves(results_backend)) == 2
                 # It should be under the "gradient" section.
                 @test haskey(results_backend, "gradient")
                 # It should have one tagged "linked" and one "standard"
-                @test length(leaves(results_backend[@tagged "linked"])) == 1
-                @test length(leaves(results_backend[@tagged "standard"])) == 1
+                @test length(BenchmarkTools.leaves(results_backend[@tagged "linked"])) == 1
+                @test length(BenchmarkTools.leaves(results_backend[@tagged "standard"])) == 1
             end
         end
 
@@ -89,12 +90,12 @@ ADBACKENDS = TuringBenchmarking.DEFAULT_ADBACKENDS
                 adbackend_string = "$(adbackend)"
                 results_backend = results[@tagged adbackend_string]
                 # Each AD backend should have two results.
-                @test length(leaves(results_backend)) == 2
+                @test length(BenchmarkTools.leaves(results_backend)) == 2
                 # It should be under the "gradient" section.
                 @test haskey(results_backend, "gradient")
                 # It should have one tagged "linked" and one "standard"
-                @test length(leaves(results_backend[@tagged "linked"])) == 1
-                @test length(leaves(results_backend[@tagged "standard"])) == 1
+                @test length(BenchmarkTools.leaves(results_backend[@tagged "linked"])) == 1
+                @test length(BenchmarkTools.leaves(results_backend[@tagged "standard"])) == 1
             end
         end
     end
@@ -116,7 +117,7 @@ ADBACKENDS = TuringBenchmarking.DEFAULT_ADBACKENDS
             DynamicPPL.SimpleVarInfo(DynamicPPL.OrderedDict(@varname(x) => randn(2))),
         ]
             # Zygote will fail.
-            @test_throws Union{MethodError,ErrorException} TuringBenchmarking.make_turing_suite(
+            @test_throws Union{MethodError,ErrorException,Zygote.CompileError} TuringBenchmarking.make_turing_suite(
                 model;
                 adbackends=ADBACKENDS,
                 varinfo=varinfo,
@@ -134,17 +135,17 @@ ADBACKENDS = TuringBenchmarking.DEFAULT_ADBACKENDS
             @testset "$adbackend" for (i, adbackend) in enumerate(ADBACKENDS)
                 adbackend_string = "$(adbackend)"
                 results_backend = results[@tagged adbackend_string]
-                if adbackend isa TuringBenchmarking.ZygoteAD
+                if adbackend isa AutoZygote
                     # Zygote.jl should fail, i.e. return an empty suite.
-                    @test length(leaves(results_backend)) == 0
+                    @test length(BenchmarkTools.leaves(results_backend)) == 0
                 else
                     # Each AD backend should have two results.
-                    @test length(leaves(results_backend)) == 2
+                    @test length(BenchmarkTools.leaves(results_backend)) == 2
                     # It should be under the "gradient" section.
                     @test haskey(results_backend, "gradient")
                     # It should have one tagged "linked" and one "standard"
-                    @test length(leaves(results_backend[@tagged "linked"])) == 1
-                    @test length(leaves(results_backend[@tagged "standard"])) == 1
+                    @test length(BenchmarkTools.leaves(results_backend[@tagged "linked"])) == 1
+                    @test length(BenchmarkTools.leaves(results_backend[@tagged "standard"])) == 1
                 end
             end
         end