Generating Julia function for log density evaluation

.github/workflows/Tests.yml

@@ -12,15 +12,20 @@ jobs:
   test:
     name: Julia ${{ matrix.version }} on ${{ matrix.os }} (${{ matrix.arch }})
     runs-on: ${{ matrix.os }}
-    continue-on-error: ${{ matrix.version == 'nightly' }}
+    continue-on-error: ${{ matrix.version == 'pre' }}
     strategy:
-      fail-fast: false
       matrix:
-        version: ['1', '1.10', 'nightly']
-        os: [ubuntu-latest, windows-latest]
-        arch: [x64]
+        version:
+          - '1'
+          - 'min'
+          - 'pre'
+        os:
+          - ubuntu-latest
+          - windows-latest
+        arch:
+          - x64
         include:
-          - version: '1'
+          - version: 'min'
             os: ubuntu-latest
             arch: x64
             coverage: true
@@ -60,6 +65,11 @@ jobs:
         env:
           TEST_GROUP: "log_density"
 
+      - name: Running `source_gen` tests
+        uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "source_gen"
+
       - name: Running `gibbs` tests
         uses: nick-fields/retry@v3
         with:

Project.toml

@@ -1,6 +1,6 @@
 name = "JuliaBUGS"
 uuid = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
-version = "0.9"
+version = "0.9.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -39,11 +39,10 @@ MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 JuliaBUGSAdvancedHMCExt = ["AdvancedHMC", "MCMCChains"]
 JuliaBUGSAdvancedMHExt = ["AdvancedMH", "MCMCChains"]
 JuliaBUGSGraphMakieExt = ["GraphMakie", "GLMakie"]
-JuliaBUGSGraphPlotExt = ["GraphPlot"]
-JuliaBUGSMCMCChainsExt = ["MCMCChains"]
+JuliaBUGSGraphPlotExt = "GraphPlot"
+JuliaBUGSMCMCChainsExt = "MCMCChains"
 
 [compat]
-ADTypes = "1.6"
 AbstractMCMC = "5"
 AbstractPPL = "0.8.4, 0.9, 0.10, 0.11"
 Accessors = "0.1"
@@ -52,6 +51,7 @@ AdvancedMH = "0.8"
 BangBang = "0.4.1"
 Bijectors = "0.13, 0.14, 0.15.5"
 ChainRules = "1"
+DifferentiationInterface = "0.6.42"
 Distributions = "0.23.8, 0.24, 0.25"
 Documenter = "0.27, 1"
 GLMakie = "0.10, 0.11"
@@ -67,6 +67,7 @@ LogExpFunctions = "0.3"
 MCMCChains = "6"
 MacroTools = "0.5"
 MetaGraphsNext = "0.6, 0.7"
+Mooncake = "0.4"
 OrderedCollections = "1"
 PDMats = "0.10, 0.11"
 Serialization = "1.10"
@@ -76,15 +77,27 @@ Statistics = "1.10"
 julia = "1.10.8"
 
 [extras]
-ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AbstractMCMC", "ADTypes", "AdvancedHMC", "AdvancedMH", "ChainRules", "MCMCChains", "LogDensityProblemsAD", "ReverseDiff", "Test"]
+test = [
+    "AbstractMCMC",
+    "AdvancedHMC",
+    "AdvancedMH",
+    "ChainRules",
+    "DifferentiationInterface",
+    "LogDensityProblemsAD",
+    "MCMCChains",
+    "Mooncake",
+    "ReverseDiff",
+    "Test"
+]

benchmark/Project.toml

@@ -1,12 +1,13 @@
 [deps]
-ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BridgeStan = "c88b6f0a-829e-4b0b-94b7-f06ab5908f5a"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

benchmark/benchmark.jl

@@ -1,11 +1,11 @@
-module Benchmark
-
 using Pkg
 Pkg.develop(; path=joinpath(@__DIR__, ".."))
 
 using JuliaBUGS
-using ADTypes
-using ReverseDiff
+
+using DifferentiationInterface
+using Mooncake: Mooncake
+
 using MetaGraphsNext
 using BridgeStan
 using StanLogDensityProblems
@@ -96,5 +96,3 @@ function _print_results_table(
         backend=backend,
     )
 end
-
-end

benchmark/juliabugs.jl

@@ -71,15 +71,28 @@ function _create_JuliaBUGS_model(model_name::Symbol)
     return compile(model_def, data, inits)
 end
 
-function benchmark_JuliaBUGS_model(model::JuliaBUGS.BUGSModel)
-    ad_model = ADgradient(AutoReverseDiff(true), model)
+# ! writing a _function_ to benchmark all models won't work because of world-age error
+
+function benchmark_JuliaBUGS_model_with_Mooncake(model::JuliaBUGS.BUGSModel)
+    p = Base.Fix1(LogDensityProblems.logdensity, model)
+    backend = AutoMooncake(; config=nothing)
     dim = LogDensityProblems.dimension(model)
     params_values = JuliaBUGS.getparams(model)
-    density_time = Chairmarks.@be LogDensityProblems.logdensity($ad_model, $params_values)
-    density_and_gradient_time = Chairmarks.@be LogDensityProblems.logdensity_and_gradient(
-        $ad_model, $params_values
-    )
+    prep = prepare_gradient(p, backend, params_values)
+    density_time = Chairmarks.@be LogDensityProblems.logdensity($model, $params_values)
+    density_and_gradient_time = Chairmarks.@be gradient($p, $prep, $backend, $params_values)
     return BenchmarkResult(:juliabugs, dim, density_time, density_and_gradient_time)
 end
 
-# writing a _function_ to benchmark all models won't work because of worldage error
+# function benchmark_JuliaBUGS_model_with_Enzyme(model::JuliaBUGS.BUGSModel)
+#     f(params, model) = LogDensityProblems.logdensity(model, params)
+#     backend = AutoEnzyme()
+#     dim = LogDensityProblems.dimension(model)
+#     params_values = JuliaBUGS.getparams(model)
+#     prep = prepare_gradient(f, backend, params_values, Constant(model))
+#     density_time = Chairmarks.@be LogDensityProblems.logdensity($model, $params_values)
+#     density_and_gradient_time = Chairmarks.@be gradient(
+#         $f, $prep, $backend, $params_values, $(Constant(model))
+#     )
+#     return BenchmarkResult(:juliabugs_enzyme, dim, density_time, density_and_gradient_time)
+# end

benchmark/run_benchmarks.jl

@@ -1,22 +1,35 @@
 include("benchmark.jl")
-using OrderedCollections
 
 examples_to_benchmark = [
     :rats, :pumps, :bones, :oxford, :epil, :lsat, :schools, :beetles, :air
 ]
 
-stan_results = Benchmark.benchmark_Stan_models(examples_to_benchmark)
+stan_results = benchmark_Stan_models(examples_to_benchmark)
 
 juliabugs_models = [
-    Benchmark._create_JuliaBUGS_model(model_name) for model_name in examples_to_benchmark
+    JuliaBUGS.set_evaluation_mode(
+        _create_JuliaBUGS_model(model_name), JuliaBUGS.UseGeneratedLogDensityFunction()
+    ) for model_name in examples_to_benchmark
 ]
-juliabugs_results = OrderedDict{Symbol,Benchmark.BenchmarkResult}()
+juliabugs_results = OrderedDict{Symbol,BenchmarkResult}()
 for (model_name, model) in zip(examples_to_benchmark, juliabugs_models)
-    @info "Benchmarking $model_name"
-    juliabugs_results[model_name] = Benchmark.benchmark_JuliaBUGS_model(model)
+    @info "Benchmarking $model_name with Mooncake"
+    juliabugs_results[model_name] = benchmark_JuliaBUGS_model_with_Mooncake(model)
 end
 
+# juliabugs_enzyme_results = OrderedDict{Symbol,BenchmarkResult}()
+# for (model_name, model) in zip(examples_to_benchmark, juliabugs_models)
+#     @info "Benchmarking $model_name with Enzyme"
+#     try
+#         juliabugs_enzyme_results[model_name] = benchmark_JuliaBUGS_model_with_Enzyme(model)
+#     catch e
+#         @warn "Error benchmarking $model_name with Enzyme: $e"
+#     end
+# end
+
 println("### Stan results:")
-Benchmark._print_results_table(stan_results; backend=Val(:markdown))
-println("### JuliaBUGS results:")
-Benchmark._print_results_table(juliabugs_results; backend=Val(:markdown))
+_print_results_table(stan_results; backend=Val(:markdown))
+println("### JuliaBUGS Mooncake results:")
+_print_results_table(juliabugs_results; backend=Val(:markdown))
+# println("### JuliaBUGS Enzyme results:")
+# _print_results_table(juliabugs_enzyme_results; backend=Val(:markdown))