From 133f5b7c93fa07a08b896e04fd858076370ef71d Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Sat, 1 Feb 2025 17:32:09 +0000 Subject: [PATCH 1/3] Add Mooncake to list of tested AD backends --- Project.toml | 12 ++++++++---- src/TuringBenchmarking.jl | 2 ++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index b377cb0..341829d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,17 +1,19 @@ name = "TuringBenchmarking" uuid = "0db1332d-5c25-4deb-809f-459bc696f94f" authors = ["Tor Erlend Fjelde and contributors"] -version = "0.5.8" +version = "0.5.9" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" 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" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" Requires = "ae029012-a4dd-5104-9daa-d747884805df" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" @@ -24,19 +26,21 @@ BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" TuringBenchmarkingBridgeStanExt = "BridgeStan" [compat] -ADTypes = "0.2, 1" +ADTypes = "1" AbstractMCMC = "5" BenchmarkTools = "1" BridgeStan = "2" -DynamicPPL = "0.24.7, 0.25, 0.27, 0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34" +DifferentiationInterface = "0.6.32" +DynamicPPL = "0.30, 0.31, 0.32, 0.33, 0.34" ForwardDiff = "0.10.3" LogDensityProblems = "2" LogDensityProblemsAD = "1" +Mooncake = "0.4.79" PrettyTables = "2" Requires = "1" ReverseDiff = "1.14" Zygote = "0.6" -julia = "1.6" +julia = "1.10" [extras] BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a" diff --git a/src/TuringBenchmarking.jl b/src/TuringBenchmarking.jl index 33b8713..88bb1b9 100644 --- a/src/TuringBenchmarking.jl +++ b/src/TuringBenchmarking.jl @@ -18,6 +18,7 @@ using DynamicPPL: DynamicPPL using ForwardDiff: ForwardDiff using ReverseDiff: ReverseDiff using Zygote: Zygote +using Mooncake: Mooncake if !isdefined(Base, :get_extension) using Requires @@ -31,6 +32,7 @@ const DEFAULT_ADBACKENDS = [ AutoReverseDiff(compile=false), AutoReverseDiff(compile=true), AutoZygote(), + AutoMooncake(; config=nothing), ] backend_label(x) = "$x" From 88a5926123cf6ed963fe3e76530d7f675bba832a Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Sat, 1 Feb 2025 17:37:47 +0000 Subject: [PATCH 2/3] Fix other dependencies / compats --- docs/Project.toml | 1 - test/Project.toml | 2 ++ test/runtests.jl | 2 ++ 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/Project.toml b/docs/Project.toml index 8af898b..a816c55 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -4,4 +4,3 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [compat] Documenter = "1" -Turing = "0.30" \ No newline at end of file diff --git a/test/Project.toml b/test/Project.toml index a4d1f6b..dbbd530 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,10 +1,12 @@ [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/test/runtests.jl b/test/runtests.jl index f763c9e..24b19e8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,6 +6,8 @@ using ADTypes using Zygote: Zygote using ReverseDiff: ReverseDiff +using Mooncake: Mooncake +import DifferentiationInterface # Just make things run a bit faster. BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1 From 8de31b316c4e936f5d22df5de8b0c650aab1d8ea Mon Sep 17 00:00:00 2001 From: Penelope Yong Date: Sat, 1 Feb 2025 17:42:20 +0000 Subject: [PATCH 3/3] Fix more errors --- src/TuringBenchmarking.jl | 6 ++++-- test/runtests.jl | 2 +- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/TuringBenchmarking.jl b/src/TuringBenchmarking.jl index 88bb1b9..b8cf6fd 100644 --- a/src/TuringBenchmarking.jl +++ b/src/TuringBenchmarking.jl @@ -37,12 +37,13 @@ const DEFAULT_ADBACKENDS = [ backend_label(x) = "$x" backend_label(::AutoForwardDiff) = "ForwardDiff" -function backend_label(ad::AutoReverseDiff) - "ReverseDiff" * (ad.compile ? " [compiled]" : "") +function backend_label(::AutoReverseDiff{compiled}) where compiled + "ReverseDiff" * (compiled ? " [compiled]" : "") end backend_label(::AutoZygote) = "Zygote" backend_label(::AutoTracker) = "Tracker" backend_label(::AutoEnzyme) = "Enzyme" +backend_label(::AutoMooncake) = "Mooncake" const SYMBOL_TO_BACKEND = Dict( :forwarddiff => AutoForwardDiff(chunksize=0), @@ -50,6 +51,7 @@ const SYMBOL_TO_BACKEND = Dict( :reversediff_compiled => AutoReverseDiff(compile=true), :zygote => AutoZygote(), :tracker => AutoTracker(), + :mooncake => AutoMooncake(; config=nothing), ) to_backend(x) = error("Unknown backend: $x") diff --git a/test/runtests.jl b/test/runtests.jl index 24b19e8..c4d77f8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -82,7 +82,7 @@ ADBACKENDS = TuringBenchmarking.DEFAULT_ADBACKENDS varinfo = DynamicPPL.VarInfo(model) suite = TuringBenchmarking.make_turing_suite( model; - adbackends=[:forwarddiff, :reversediff, :reversediff_compiled, :zygote], + adbackends=[:forwarddiff, :reversediff, :reversediff_compiled, :zygote, :mooncake], varinfo=varinfo, ) results = run(suite, verbose=true)