applied suggested changes

Author: shravanngoswamii

benchmarks/benchmarks.jl

@@ -1,15 +1,14 @@
-using DynamicPPL
-using DynamicPPLBenchmarks
-using BenchmarkTools
-using TuringBenchmarking
-using Distributions
-using PrettyTables
+using DynamicPPL: @model
+using DynamicPPLBenchmarks: make_suite
+using BenchmarkTools: median, run
+using Distributions: Normal, Beta, Bernoulli
+using PrettyTables: pretty_table, PrettyTables
 
 # Define models
 @model function demo1(x)
     m ~ Normal()
     x ~ Normal(m, 1)
-    return (m = m, x = x)
+    return (m=m, x=x)
 end
 
 @model function demo2(y)
@@ -28,60 +27,39 @@ demo2_data = rand(Bool, 10)
 demo1_instance = demo1(demo1_data)
 demo2_instance = demo2(demo2_data)
 
-# Define available AD backends
-available_ad_backends = Dict(
-    :forwarddiff => :forwarddiff,
-    :reversediff => :reversediff,
-    :zygote      => :zygote
-)
-
-# Define available VarInfo types.
-# Each entry is (Name, function to produce the VarInfo)
-available_varinfo_types = Dict(
-    :untyped           => ("UntypedVarInfo", VarInfo),
-    :typed             => ("TypedVarInfo", m -> VarInfo(m)),
-    :simple_namedtuple => ("SimpleVarInfo (NamedTuple)", m -> SimpleVarInfo{Float64}(m())),
-    :simple_dict       => ("SimpleVarInfo (Dict)", m -> begin
-        retvals = m()
-        varnames = map(keys(retvals)) do k
-            VarName{k}()
-        end
-        SimpleVarInfo{Float64}(Dict(zip(varnames, values(retvals))))
-    end)
-)
-
 # Specify the combinations to test:
 # (Model Name, model instance, VarInfo choice, AD backend)
 chosen_combinations = [
-    ("Demo1", demo1_instance, :typed,           :forwarddiff),
+    ("Demo1", demo1_instance, :typed, :forwarddiff),
     ("Demo1", demo1_instance, :simple_namedtuple, :zygote),
-    ("Demo2", demo2_instance, :untyped,           :reversediff),
-    ("Demo2", demo2_instance, :simple_dict,       :forwarddiff)
+    ("Demo2", demo2_instance, :untyped, :reversediff),
+    ("Demo2", demo2_instance, :simple_dict, :forwarddiff),
 ]
 
-# Store results as tuples: (Model, AD Backend, VarInfo Type, Eval Time, AD Eval Time)
-results_table = Tuple{String, String, String, Float64, Float64}[]
+results_table = Tuple{String,String,String,Float64,Float64}[]
 
 for (model_name, model, varinfo_choice, adbackend) in chosen_combinations
     suite = make_suite(model, varinfo_choice, adbackend)
     results = run(suite)
-    eval_time    = median(results["evaluation"]).time
-    ad_eval_time = median(results["AD_Benchmarking"]["evaluation"]["standard"]).time
-    push!(results_table, (model_name, string(adbackend), string(varinfo_choice), eval_time, ad_eval_time))
-end
 
-# Convert results to a 2D array for PrettyTables
-function to_matrix(tuples::Vector{<:NTuple{5,Any}})
-    n = length(tuples)
-    data = Array{Any}(undef, n, 5)
-    for i in 1:n
-        for j in 1:5
-            data[i, j] = tuples[i][j]
-        end
+    eval_time = median(results["AD_Benchmarking"]["evaluation"]["standard"]).time
+
+    grad_group = results["AD_Benchmarking"]["gradient"]
+    if isempty(grad_group)
+        ad_eval_time = NaN
+    else
+        grad_backend_key = first(keys(grad_group))
+        ad_eval_time = median(grad_group[grad_backend_key]["standard"]).time
     end
-    return data
+
+    push!(
+        results_table,
+        (model_name, string(adbackend), string(varinfo_choice), eval_time, ad_eval_time),
+    )
 end
 
-table_matrix = to_matrix(results_table)
-header = ["Model", "AD Backend", "VarInfo Type", "Evaluation Time (ns)", "AD Eval Time (ns)"]
+table_matrix = hcat(Iterators.map(collect, zip(results_table...))...)
+header = [
+    "Model", "AD Backend", "VarInfo Type", "Evaluation Time (ns)", "AD Eval Time (ns)"
+]
 pretty_table(table_matrix; header=header, tf=PrettyTables.tf_markdown)

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -1,7 +1,7 @@
 module DynamicPPLBenchmarks
 
-using DynamicPPL
-using BenchmarkTools
+using DynamicPPL: VarInfo, SimpleVarInfo, VarName
+using BenchmarkTools: BenchmarkGroup
 using TuringBenchmarking: make_turing_suite
 
 export make_suite
@@ -13,40 +13,40 @@ Create a benchmark suite for `model` using the selected varinfo type and AD back
 Available varinfo choices:
   • `:untyped`           → uses `VarInfo()`
   • `:typed`             → uses `VarInfo(model)`
-  • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())`
+  • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(free_nt)`
   • `:simple_dict`       → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs)
 
 The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversediff`, `:zygote`).
 """
 function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol)
     suite = BenchmarkGroup()
-    context = DefaultContext()
-    
-    # Create the chosen varinfo.
-    vi = nothing
-    if varinfo_choice == :untyped
-        vi = VarInfo()
-        model(vi)
+
+    vi = if varinfo_choice == :untyped
+        v = VarInfo()
+        model(v)
+        v
     elseif varinfo_choice == :typed
-        vi = VarInfo(model)
+        VarInfo(model)
     elseif varinfo_choice == :simple_namedtuple
-        vi = SimpleVarInfo{Float64}(model())
+        free_nt = NamedTuple{(:m,)}(model())  # Extract only the free parameter(s)
+        SimpleVarInfo{Float64}(free_nt)
     elseif varinfo_choice == :simple_dict
         retvals = model()
-        vns = map(keys(retvals)) do k
-            VarName{k}()
-        end
-        vi = SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals))))
+        vns = [VarName{k}() for k in keys(retvals)]
+        SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals))))
     else
         error("Unknown varinfo choice: $varinfo_choice")
     end
 
-    # Add the evaluation benchmark.
-    suite["evaluation"] = @benchmarkable $model($vi, $context)
-    
     # Add the AD benchmarking suite.
-    suite["AD_Benchmarking"] = make_turing_suite(model; adbackends=[adbackend])
-    
+    suite["AD_Benchmarking"] = make_turing_suite(
+        model;
+        adbackends=[adbackend],
+        varinfo=vi,
+        check_grads=true,
+        error_on_failed_backend=true,
+    )
+
     return suite
 end