Add faster ParamsWithStats constructor

Author: penelopeysm

src/chains.jl

@@ -120,13 +120,27 @@ via `unflatten!!` plus re-evaluation. It is faster for two reasons:
    otherwise re-evaluation would mutate the VarInfo, rendering it unusable for subsequent
    MCMC iterations).
 2. The re-evaluation is faster as it uses `OnlyAccsVarInfo`.
+
+Furthermore, if the `LogDensityFunction` has all fixed transforms (i.e., was constructed
+with `fix_transforms=true`), and neither `include_log_probs` nor `include_colon_eq` is
+set, then model re-evaluation is skipped entirely and the raw parameter values are
+extracted directly from the parameter vector using the cached transforms.
 """
 function ParamsWithStats(
     param_vector::AbstractVector,
     ldf::DynamicPPL.LogDensityFunction,
     stats::NamedTuple=NamedTuple();
     include_colon_eq::Bool=true,
     include_log_probs::Bool=true,
+)
+    return pws_with_eval(param_vector, ldf, stats; include_colon_eq, include_log_probs)
+end
+function pws_with_eval(
+    param_vector::AbstractVector,
+    ldf::DynamicPPL.LogDensityFunction,
+    stats::NamedTuple=NamedTuple();
+    include_colon_eq::Bool=true,
+    include_log_probs::Bool=true,
 )
     strategy = InitFromVector(param_vector, ldf)
     accs = if include_log_probs
@@ -158,6 +172,30 @@ function ParamsWithStats(
     return ParamsWithStats(params, stats)
 end
 
+# Specialisation for when the LDF is known to have all fixed transforms
+function ParamsWithStats(
+    param_vector::AbstractVector,
+    ldf::LogDensityFunction{M,A,L,F,V,D,X,C,true},
+    stats::NamedTuple=NamedTuple();
+    include_colon_eq::Bool=true,
+    include_log_probs::Bool=true,
+) where {M,A,L,F,V,D,X,C}
+    if include_log_probs || include_colon_eq
+        return pws_with_eval(param_vector, ldf, stats; include_colon_eq, include_log_probs)
+    end
+    # Fast path: extract raw values directly from the parameter vector using the fixed
+    # transforms, without re-evaluating the model.
+    params = VarNamedTuple()
+    for (vn, rat) in pairs(ldf._varname_ranges)
+        top_sym = AbstractPPL.getsym(vn)
+        template = get(ldf._varname_ranges.data, top_sym, DynamicPPL.NoTemplate())
+        raw_val = rat.transform.transform(param_vector[rat.range])
+        params = DynamicPPL.templated_setindex!!(params, raw_val, vn, template)
+    end
+    params = densify!!(params)
+    return ParamsWithStats(params, stats)
+end
+
 function Base.show(io::IO, ::MIME"text/plain", pws::ParamsWithStats)
     printstyled(io, "ParamsWithStats"; bold=true)
     print(io, "\n ├─ ")

src/logdensityfunction.jl

@@ -189,6 +189,8 @@ struct LogDensityFunction{
     # type of the vector passed to logdensity functions
     X<:AbstractVector,
     AC<:AccumulatorTuple,
+    # whether all transforms are FixedTransforms, enabling fast parameter extraction
+    AllFixed,
 }
     model::M
     adtype::AD
@@ -231,6 +233,12 @@ struct LogDensityFunction{
         end
         ranges_and_transforms = get_rangeandtransforms(vnt)
 
+        # Determine whether all transforms are fixed. This enables fast parameter
+        # extraction in ParamsWithStats without model re-evaluation.
+        all_fixed = all(
+            rat -> rat.transform isa FixedTransform, values(ranges_and_transforms)
+        )
+
         # Get vectorised parameters. Note that `internal_values_as_vector` just concatenates
         # all the vectors inside in iteration order of the VNT's keys. *In principle*, the
         # result of that should always be consistent with the ranges extracted above via
@@ -268,6 +276,7 @@ struct LogDensityFunction{
             typeof(prep),
             typeof(x),
             typeof(accs),
+            all_fixed,
         }(
             model,
             adtype,

test/chains.jl

@@ -73,7 +73,7 @@ end
     @testset "$(m.f)" for m in DynamicPPL.TestUtils.ALL_MODELS
         @testset "$transform_strategy" for transform_strategy in (UnlinkAll(), LinkAll())
             # Get the ParamsWithStats using LogDensityFunction
-            ldf = DynamicPPL.LogDensityFunction(m, getlogjoint, transform_strategy)
+            ldf = LogDensityFunction(m, getlogjoint, transform_strategy)
             param_vector = rand(ldf)
             # This will give us a VNT of values.params`.
             actual_vnt = ParamsWithStats(param_vector, ldf).params
@@ -93,6 +93,54 @@ end
     end
 end
 
+@testset "ParamsWithStats from LogDensityFunction with fixed transforms" begin
+    # Note: can't use ALL_MODELS here because that contains a model with dynamic transforms,
+    # which would yield incorrect results with fix_transforms.
+    @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
+        @testset "$transform_strategy" for transform_strategy in (UnlinkAll(), LinkAll())
+            ldf_fixed = LogDensityFunction(
+                m, getlogjoint_internal, transform_strategy; fix_transforms=true
+            )
+            ldf_dynamic = LogDensityFunction(m, getlogjoint_internal, transform_strategy)
+            param_vector = rand(ldf_fixed)
+
+            # Fast path (no log probs, no colon eq): should match the model-evaluation path
+            fast = ParamsWithStats(
+                param_vector, ldf_fixed; include_log_probs=false, include_colon_eq=false
+            )
+            slow = ParamsWithStats(
+                param_vector, ldf_dynamic; include_log_probs=false, include_colon_eq=false
+            )
+            @test fast == slow
+        end
+    end
+
+    @testset "check that model is actually not evaluated" begin
+        should_error = false
+        @model function prickly()
+            x ~ Normal()
+            return should_error && error("nope")
+        end
+        # need to construct LDF without erroring
+        ldf = LogDensityFunction(
+            prickly(), getlogjoint_internal, LinkAll(); fix_transforms=true
+        )
+        # now make the model error
+        should_error = true
+        @test_throws ErrorException prickly()()
+        # check that ParamsWithStats doesn't error
+        @test ParamsWithStats(
+            [0.5], ldf; include_log_probs=false, include_colon_eq=false
+        ) isa Any
+        # but it does if you set either of them to true
+        for (ilp, ice) in ((true, false), (false, true), (true, true))
+            @test_throws ErrorException ParamsWithStats(
+                [0.5], ldf; include_log_probs=ilp, include_colon_eq=ice
+            )
+        end
+    end
+end
+
 @info "Completed $(@__FILE__) in $(now() - __now__)."
 
 end # module

test/logdensityfunction.jl

@@ -549,12 +549,21 @@ end
 end
 
 @testset "LogDensityFunction: fix_transforms correctness" begin
+    # Helper function to check whether the AllFixed type parameter on LogDensityFunction is
+    # set correctly.
+    function has_all_fixed_transforms(
+        ::LogDensityFunction{M,A,L,F,V,D,X,C,AF}
+    ) where {M,A,L,F,V,D,X,C,AF}
+        return AF
+    end
+
     @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
         @testset "$strategy" for strategy in (UnlinkAll(), LinkAll())
             ldf_dynamic = LogDensityFunction(m, getlogjoint_internal, strategy)
             ldf_fixed = LogDensityFunction(
                 m, getlogjoint_internal, strategy; fix_transforms=true
             )
+            @test has_all_fixed_transforms(ldf_fixed)
             # Check that the transform strategy does contain fixed transforms
             tfm_strategy = ldf_fixed.transform_strategy
             @test tfm_strategy isa WithTransforms