Merge remote-tracking branch 'origin/main' into breaking

Author: mhauru

.github/workflows/Benchmarking.yml

@@ -33,14 +33,19 @@ jobs:
           echo "$version_info" >> $GITHUB_ENV
           echo "EOF" >> $GITHUB_ENV
 
-          # Capture benchmark output into a variable
+          # Capture benchmark output into a variable. The sed and tail calls cut out anything but the
+          # final block of results.
           echo "Running Benchmarks..."
-          benchmark_output=$(julia --project=benchmarks benchmarks/benchmarks.jl)
-          
+          benchmark_output=$(\
+            julia --project=benchmarks benchmarks/benchmarks.jl \
+            | sed -n '/Final results:/,$p' \
+            | tail -n +2\
+          )
+
           # Print benchmark results directly to the workflow log
           echo "Benchmark Results:"
           echo "$benchmark_output"
-          
+
           # Set the benchmark output as an env var for later steps
           echo "BENCHMARK_OUTPUT<<EOF" >> $GITHUB_ENV
           echo "$benchmark_output" >> $GITHUB_ENV

HISTORY.md

@@ -4,6 +4,10 @@
 
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
 
+## 0.38.4
+
+Improve performance of VarNamedVector. It should now be very nearly on par with Metadata for all models we've benchmarked on.
+
 ## 0.38.3
 
 Add an implementation of `returned(::Model, ::AbstractDict{<:VarName})`.

benchmarks/Project.toml

@@ -30,4 +30,4 @@ LogDensityProblems = "2.1.2"
 Mooncake = "0.4"
 PrettyTables = "3"
 ReverseDiff = "1.15.3"
-StableRNGs = "1"
+StableRNGs = "1"

benchmarks/benchmarks.jl

@@ -7,6 +7,27 @@ using StableRNGs: StableRNG
 
 rng = StableRNG(23)
 
+function print_results(results_table)
+    table_matrix = hcat(Iterators.map(collect, zip(results_table...))...)
+    header = [
+        "Model",
+        "Dim",
+        "AD Backend",
+        "VarInfo",
+        "Linked",
+        "t(eval)/t(ref)",
+        "t(grad)/t(eval)",
+    ]
+    return pretty_table(
+        table_matrix;
+        column_labels=header,
+        backend=:text,
+        formatters=[fmt__printf("%.1f", [6, 7])],
+        fit_table_in_display_horizontally=false,
+        fit_table_in_display_vertically=false,
+    )
+end
+
 # Create DynamicPPL.Model instances to run benchmarks on.
 smorgasbord_instance = Models.smorgasbord(randn(rng, 100), randn(rng, 100))
 loop_univariate1k, multivariate1k = begin
@@ -41,6 +62,8 @@ chosen_combinations = [
     ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :typed_vector, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :untyped_vector, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true),
@@ -82,17 +105,9 @@ for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinati
             relative_ad_eval_time,
         ),
     )
+    println("Results so far:")
+    print_results(results_table)
 end
 
-table_matrix = hcat(Iterators.map(collect, zip(results_table...))...)
-header = [
-    "Model", "Dim", "AD Backend", "VarInfo", "Linked", "t(eval)/t(ref)", "t(grad)/t(eval)"
-]
-pretty_table(
-    table_matrix;
-    column_labels=header,
-    backend=:text,
-    formatters=[fmt__printf("%.1f", [6, 7])],
-    fit_table_in_display_horizontally=false,
-    fit_table_in_display_vertically=false,
-)
+println("Final results:")
+print_results(results_table)

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -80,6 +80,10 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
         retvals = model(rng)
         vns = [VarName{k}() for k in keys(retvals)]
         SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals))))
+    elseif varinfo_choice == :typed_vector
+        DynamicPPL.typed_vector_varinfo(rng, model)
+    elseif varinfo_choice == :untyped_vector
+        DynamicPPL.untyped_vector_varinfo(rng, model)
     else
         error("Unknown varinfo choice: $varinfo_choice")
     end

docs/src/api.md

@@ -414,7 +414,7 @@ DynamicPPL.reset!
 DynamicPPL.update!
 DynamicPPL.insert!
 DynamicPPL.loosen_types!!
-DynamicPPL.tighten_types
+DynamicPPL.tighten_types!!
 ```
 
 ```@docs

src/contexts/init.jl

@@ -180,7 +180,9 @@ function tilde_assume!!(
     end
     # Neither of these set the `trans` flag so we have to do it manually if
     # necessary.
-    insert_transformed_value && set_transformed!!(vi, true, vn)
+    if insert_transformed_value
+        vi = set_transformed!!(vi, true, vn)
+    end
     # `accumulate_assume!!` wants untransformed values as the second argument.
     vi = accumulate_assume!!(vi, x, logjac, vn, dist)
     # We always return the untransformed value here, as that will determine

src/debug_utils.jl

@@ -27,7 +27,7 @@ add_io_context(io::IO) = IOContext(io, :compact => true, :limit => true)
 show_varname(io::IO, varname::VarName) = print(io, varname)
 function show_varname(io::IO, varname::Array{<:VarName,N}) where {N}
     # Attempt to make the type concrete in case the symbol is shared.
-    return _show_varname(io, map(identity, varname))
+    return _show_varname(io, [vn for vn in varname])
 end
 function _show_varname(io::IO, varname::Array{<:VarName,N}) where {N}
     # Print the first and last element of the array.
@@ -407,7 +407,7 @@ julia> @model function demo_incorrect()
        end
 demo_incorrect (generic function with 2 methods)
 
-julia> # Notice that VarInfo(model_incorrect) evaluates the model, but doesn't actually 
+julia> # Notice that VarInfo(model_incorrect) evaluates the model, but doesn't actually
        # alert us to the issue of `x` being sampled twice.
        model = demo_incorrect(); varinfo = VarInfo(model);
 

src/logdensityfunction.jl

@@ -49,7 +49,7 @@ box:
 - [`getlogprior`](@ref): calculate the log prior in the model space, ignoring
   any effects of linking
 - [`getloglikelihood`](@ref): calculate the log likelihood (this is unaffected
-  by linking, since transforms are only applied to random variables) 
+  by linking, since transforms are only applied to random variables)
 
 !!! note
     By default, `LogDensityFunction` uses `getlogjoint_internal`, i.e., the
@@ -146,7 +146,7 @@ struct LogDensityFunction{
             is_supported(adtype) ||
                 @warn "The AD backend $adtype is not officially supported by DynamicPPL. Gradient calculations may still work, but compatibility is not guaranteed."
             # Get a set of dummy params to use for prep
-            x = map(identity, varinfo[:])
+            x = [val for val in varinfo[:]]
             if use_closure(adtype)
                 prep = DI.prepare_gradient(
                     LogDensityAt(model, getlogdensity, varinfo), adtype, x
@@ -282,7 +282,7 @@ function LogDensityProblems.logdensity_and_gradient(
 ) where {M,F,V,AD<:ADTypes.AbstractADType}
     f.prep === nothing &&
         error("Gradient preparation not available; this should not happen")
-    x = map(identity, x)  # Concretise type
+    x = [val for val in x]  # Concretise type
     # Make branching statically inferrable, i.e. type-stable (even if the two
     # branches happen to return different types)
     return if use_closure(f.adtype)

src/simple_varinfo.jl

@@ -484,6 +484,7 @@ function set_transformed!!(vi::SimpleOrThreadSafeSimple, trans::Bool, ::VarName)
             "Individual variables in SimpleVarInfo cannot have different `set_transformed` statuses.",
         )
     end
+    return vi
 end
 
 is_transformed(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)