diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 7e7b9f7cb..97afe3532 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -7,6 +7,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" @@ -23,6 +24,7 @@ ADTypes = "1.14.0" BenchmarkTools = "1.6.0" Distributions = "0.25.117" DynamicPPL = "0.37" +Enzyme = "0.13" ForwardDiff = "0.10.38, 1" LogDensityProblems = "2.1.2" Mooncake = "0.4" diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index dfcdcc271..dd1140dba 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -43,6 +43,7 @@ chosen_combinations = [ ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true), ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true), ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true), + ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true), ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true), ("Multivariate 1k", multivariate1k, :typed, :mooncake, true), ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 8c5032ace..dcb6ac9be 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -13,7 +13,7 @@ using StableRNGs: StableRNG include("./Models.jl") using .Models: Models - +using Enzyme: Enzyme export Models, make_suite, model_dimension """ @@ -37,6 +37,10 @@ const SYMBOL_TO_BACKEND = Dict( :reversediff => ADTypes.AutoReverseDiff(; compile=false), :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true), :mooncake => ADTypes.AutoMooncake(; config=nothing), + :enzyme => ADTypes.AutoEnzyme(; + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), + function_annotation=Enzyme.Const, + ), ) to_backend(x) = error("Unknown backend: $x") diff --git a/docs/Project.toml b/docs/Project.toml index cc0be339d..47563d9a7 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -24,5 +24,6 @@ FillArrays = "0.13, 1" ForwardDiff = "0.10, 1" JET = "0.9, 0.10" LogDensityProblems = "2" +MarginalLogDensities = "0.4" MCMCChains = "5, 6, 7" StableRNGs = "1" diff --git a/test/Project.toml b/test/Project.toml index 589b150f4..0317977d4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -44,6 +44,7 @@ EnzymeCore = "0.6 - 0.8" ForwardDiff = "0.10.12, 1" JET = "0.9 - 0.10.6" LogDensityProblems = "2" +MarginalLogDensities = "0.4" MCMCChains = "7.2.1" MacroTools = "0.5.6" OrderedCollections = "1"