Implement ParamsWithStats from FastLDF

Author: penelopeysm

src/chains.jl

@@ -90,6 +90,42 @@ function ParamsWithStats(
     return ParamsWithStats(params, stats)
 end
 
+function ParamsWithStats(
+    param_vector::AbstractVector,
+    ldf::DynamicPPL.Experimental.FastLDF,
+    stats::NamedTuple=NamedTuple();
+    include_colon_eq::Bool=true,
+    include_log_probs::Bool=true,
+)
+    ctx = DynamicPPL.Experimental.FastEvalVectorContext(
+        ldf._iden_varname_ranges, ldf._varname_ranges, param_vector
+    )
+    accs = if include_log_probs
+        (
+            DynamicPPL.LogPriorAccumulator(),
+            DynamicPPL.LogLikelihoodAccumulator(),
+            DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),
+        )
+    else
+        (DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),)
+    end
+    return _, varinfo = DynamicPPL.Experimental.fast_evaluate!!(
+        ldf.model, ctx, AccumulatorTuple(accs)
+    )
+    params = DynamicPPL.getacc(varinfo, Val(:ValuesAsInModel)).values
+    if include_log_probs
+        stats = merge(
+            stats,
+            (
+                logprior=DynamicPPL.getlogprior(varinfo),
+                loglikelihood=DynamicPPL.getloglikelihood(varinfo),
+                lp=DynamicPPL.getlogjoint(varinfo),
+            ),
+        )
+    end
+    return ParamsWithStats(params, stats)
+end
+
 # Re-evaluating the model is unconscionably slow for untyped VarInfo. It's much faster to
 # convert it to a typed varinfo first, hence this method.
 # https://github.com/TuringLang/Turing.jl/issues/2604

src/fasteval.jl

@@ -352,16 +352,9 @@ end
 fast_ldf_accs(::typeof(getlogprior)) = AccumulatorTuple((LogPriorAccumulator(),))
 fast_ldf_accs(::typeof(getloglikelihood)) = AccumulatorTuple((LogLikelihoodAccumulator(),))
 
-struct FastLogDensityAt{M<:Model,F<:Function,N<:NamedTuple}
-    _model::M
-    _getlogdensity::F
-    _iden_varname_ranges::N
-    _varname_ranges::Dict{VarName,RangeAndLinked}
-end
-function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
-    ctx = FastEvalVectorContext(f._iden_varname_ranges, f._varname_ranges, params)
-    model = DynamicPPL.setleafcontext(f._model, ctx)
-    accs = fast_ldf_accs(f._getlogdensity)
+function fast_evaluate!!(model::Model, ctx::FastEvalVectorContext, accs::AccumulatorTuple)
+    model = DynamicPPL.setleafcontext(model, ctx)
+    vi = OnlyAccsVarInfo(accs)
     # Calling `evaluate!!` would be fine, but would lead to an extra call to resetaccs!!,
     # which is unnecessary. So we shortcircuit this by simply calling `_evaluate!!`
     # directly. To preserve thread-safety we need to reproduce the ThreadSafeVarInfo logic
@@ -378,7 +371,18 @@ function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
     else
         OnlyAccsVarInfo(accs)
     end
-    _, vi = DynamicPPL._evaluate!!(model, vi)
+    return DynamicPPL._evaluate!!(model, vi)
+end
+
+struct FastLogDensityAt{M<:Model,F<:Function,N<:NamedTuple}
+    _model::M
+    _getlogdensity::F
+    _iden_varname_ranges::N
+    _varname_ranges::Dict{VarName,RangeAndLinked}
+end
+function (f::FastLogDensityAt)(params::AbstractVector{<:Real})
+    ctx = FastEvalVectorContext(f._iden_varname_ranges, f._varname_ranges, params)
+    _, vi = fast_evaluate!!(f._model, ctx, fast_ldf_accs(f._getlogdensity))
     return f._getlogdensity(vi)
 end