Choose whether to show linked or unlinked benchmark times

Author: mhauru

benchmarks/Project.toml

@@ -6,5 +6,6 @@ version = "0.1.0"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f"

benchmarks/benchmarks.jl

@@ -26,24 +26,30 @@ lda_instance = begin
 end
 
 # Specify the combinations to test:
-# (Model Name, model instance, VarInfo choice, AD backend)
+# (Model Name, model instance, VarInfo choice, AD backend, linked)
 chosen_combinations = [
-    ("Simple assume observe", Models.simple_assume_observe(randn()), :typed, :forwarddiff),
-    ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff),
-    ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff),
-    ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff),
-    ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff),
-    ("Smorgasbord", smorgasbord_instance, :typed, :reversediff),
+    (
+        "Simple assume observe",
+        Models.simple_assume_observe(randn()),
+        :typed,
+        :forwarddiff,
+        false,
+    ),
+    ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff, false),
+    ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),
     # TODO(mhauru) Add Mooncake once TuringBenchmarking.jl supports it. Consider changing
     # all the below :reversediffs to :mooncakes too.
-    #("Smorgasbord", smorgasbord_instance, :typed, :mooncake),
-    ("Loop univariate 1k", loop_univariate1k, :typed, :reversediff),
-    ("Multivariate 1k", multivariate1k, :typed, :reversediff),
-    ("Loop univariate 10k", loop_univariate10k, :typed, :reversediff),
-    ("Multivariate 10k", multivariate10k, :typed, :reversediff),
-    ("Dynamic", Models.dynamic(), :typed, :reversediff),
-    ("Submodel", Models.parent(randn()), :typed, :reversediff),
-    ("LDA", lda_instance, :typed, :reversediff),
+    #("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true),
+    ("Loop univariate 1k", loop_univariate1k, :typed, :reversediff, true),
+    ("Multivariate 1k", multivariate1k, :typed, :reversediff, true),
+    ("Loop univariate 10k", loop_univariate10k, :typed, :reversediff, true),
+    ("Multivariate 10k", multivariate10k, :typed, :reversediff, true),
+    ("Dynamic", Models.dynamic(), :typed, :reversediff, true),
+    ("Submodel", Models.parent(randn()), :typed, :reversediff, true),
+    ("LDA", lda_instance, :typed, :reversediff, true),
 ]
 
 # Time running a model-like function that does not use DynamicPPL, as a reference point.
@@ -53,21 +59,22 @@ reference_time = begin
     median(@benchmark Models.simple_assume_observe_non_model(obs)).time
 end
 
-results_table = Tuple{String,String,String,Float64,Float64}[]
+results_table = Tuple{String,String,String,Bool,Float64,Float64}[]
 
-for (model_name, model, varinfo_choice, adbackend) in chosen_combinations
+for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations
     suite = make_suite(model, varinfo_choice, adbackend)
     results = run(suite)
+    result_key = islinked ? "linked" : "standard"
 
-    eval_time = median(results["evaluation"]["standard"]).time
+    eval_time = median(results["evaluation"][result_key]).time
     relative_eval_time = eval_time / reference_time
 
     grad_group = results["gradient"]
     if isempty(grad_group)
         relative_ad_eval_time = NaN
     else
         grad_backend_key = first(keys(grad_group))
-        ad_eval_time = median(grad_group[grad_backend_key]["standard"]).time
+        ad_eval_time = median(grad_group[grad_backend_key][result_key]).time
         relative_ad_eval_time = ad_eval_time / eval_time
     end
 
@@ -77,6 +84,7 @@ for (model_name, model, varinfo_choice, adbackend) in chosen_combinations
             model_name,
             string(adbackend),
             string(varinfo_choice),
+            islinked,
             relative_eval_time,
             relative_ad_eval_time,
         ),
@@ -88,12 +96,13 @@ header = [
     "Model",
     "AD Backend",
     "VarInfo Type",
+    "Linked",
     "Evaluation Time / Reference Time",
     "AD Time / Eval Time",
 ]
 PrettyTables.pretty_table(
     table_matrix;
     header=header,
     tf=PrettyTables.tf_markdown,
-    formatters=ft_printf("%.1f", [4, 5]),
+    formatters=ft_printf("%.1f", [5, 6]),
 )

benchmarks/src/Models.jl

@@ -6,18 +6,19 @@ observed (this is used for constructing SimpleVarInfos).
 """
 module Models
 
-using DynamicPPL: @model, to_submodel
 using Distributions:
     Categorical,
     Dirichlet,
     Exponential,
     Gamma,
     LKJCholesky,
-    MatrixBeta,
+    InverseWishart,
     Normal,
     logpdf,
     product_distribution,
     truncated
+using DynamicPPL: @model, to_submodel
+using LinearAlgebra: cholesky
 
 export simple_assume_observe_non_model,
     simple_assume_observe, smorgasbord, loop_univariate, multivariate, parent, dynamic, lda
@@ -114,19 +115,14 @@ Like simple_assume_observe, but with a submodel for the assumed random variable.
 end
 
 """
-A model with random variables that have changing support.
-
-Includes both variables the dimension of which depends on other variables, and variables
-the support of which changes under linking.
+A model with random variables that have changing support under linking, or otherwise
+complicated bijectors.
 """
 @model function dynamic(::Type{T}=Vector{Float64}) where {T}
-    eta ~ truncated(Normal(); lower=0.0)
+    eta ~ truncated(Normal(); lower=0.0, upper=0.1)
     mat1 ~ LKJCholesky(4, eta)
-    mat2 ~ MatrixBeta(5, 6.0, 8.0)
-    dim = eta > 0.2 ? 2 : 3
-    vec = T(undef, dim)
-    vec .~ truncated(Exponential(0.5); lower=0.0, upper=1.0)
-    return (; eta=eta, mat1=mat1, mat2=mat2, vec=vec)
+    mat2 ~ InverseWishart(3.2, cholesky([1.0 0.5; 0.5 1.0]))
+    return (; eta=eta, mat1=mat1, mat2=mat2)
 end
 
 """