Use init!! instead of fast_evaluate!!

Author: penelopeysm

docs/src/api.md

@@ -66,11 +66,10 @@ The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) inte
 LogDensityFunction
 ```
 
-Internally, this is accomplished using:
+Internally, this is accomplished using [`init!!`](@ref) on:
 
 ```@docs
 OnlyAccsVarInfo
-fast_evaluate!!
 ```
 
 ## Condition and decondition
@@ -517,7 +516,7 @@ The function `init!!` is used to initialise, or overwrite, values in a VarInfo.
 It is really a thin wrapper around using `evaluate!!` with an `InitContext`.
 
 ```@docs
-DynamicPPL.init!!
+init!!
 ```
 
 To accomplish this, an initialisation _strategy_ is required, which defines how new values are to be obtained.

src/DynamicPPL.jl

@@ -92,10 +92,8 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
-    # LogDensityFunction and fasteval
+    # LogDensityFunction
     LogDensityFunction,
-    fast_evaluate!!,
-    OnlyAccsVarInfo,
     # Leaf contexts
     AbstractContext,
     contextualize,
@@ -110,6 +108,9 @@ export AbstractVarInfo,
     # Tilde pipeline
     tilde_assume!!,
     tilde_observe!!,
+    # Evaluation
+    evaluate!!,
+    init!!,
     # Initialisation
     AbstractInitStrategy,
     InitFromPrior,

src/chains.jl

@@ -137,7 +137,7 @@ end
 """
     ParamsWithStats(
         param_vector::AbstractVector,
-        ldf::DynamicPPL.Experimental.FastLDF,
+        ldf::DynamicPPL.LogDensityFunction,
         stats::NamedTuple=NamedTuple();
         include_colon_eq::Bool=true,
         include_log_probs::Bool=true,
@@ -152,11 +152,11 @@ via `unflatten` plus re-evaluation. It is faster for two reasons:
 1. It does not rely on `deepcopy`-ing the VarInfo object (this used to be mandatory as
    otherwise re-evaluation would mutate the VarInfo, rendering it unusable for subsequent
    MCMC iterations).
-2. The re-evaluation is faster as it uses `OnlyAccsVarInfo`.
+2. The re-evaluation is faster as it uses `OnlyAccsVarInfo` rather than a full VarInfo.
 """
 function ParamsWithStats(
     param_vector::AbstractVector,
-    ldf::DynamicPPL.Experimental.FastLDF,
+    ldf::DynamicPPL.LogDensityFunction,
     stats::NamedTuple=NamedTuple();
     include_colon_eq::Bool=true,
     include_log_probs::Bool=true,
@@ -174,9 +174,7 @@ function ParamsWithStats(
     else
         (DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),)
     end
-    _, vi = DynamicPPL.Experimental.fast_evaluate!!(
-        ldf.model, strategy, AccumulatorTuple(accs)
-    )
+    _, vi = DynamicPPL.init!!(ldf.model, AccumulatorTuple(accs), strategy)
     params = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
     if include_log_probs
         stats = merge(

src/contexts/init.jl

@@ -312,11 +312,12 @@ function tilde_assume!!(
         # linked, and `inv_logjac` will also just be the negative of `fwd_logjac`.
         #
         # However, `VectorWithRanges` is only really used with `OnlyAccsVarInfo`, in which
-        # case this branch is never hit (since `in_varinfo` will always be false). It does
-        # mean that the combination of InitFromParams{<:VectorWithRanges} with a full,
-        # linked, VarInfo will be very slow. That should never really be used, though. So
-        # (at least for now) we can leave this branch in for full generality with other
-        # combinations of init strategies / VarInfo.
+        # case this method is never hit (since there's a special method for it, in
+        # `src/onlyaccs.jl`). It does mean that the combination of
+        # InitFromParams{<:VectorWithRanges} with a full, linked, VarInfo will be very slow.
+        # That should never really be used, though. So (at least for now) we can leave this
+        # branch in for full generality with other combinations of init strategies /
+        # VarInfo.
         #
         # TODO(penelopeysm): Figure out one day how to refactor this. The crux of the issue
         # is that the transform used by `VectorWithRanges` is `from_linked_VEC_transform`,

src/fasteval.jl

@@ -29,60 +29,6 @@ using LogDensityProblems: LogDensityProblems
 import DifferentiationInterface as DI
 using Random: Random
 
-"""
-    DynamicPPL.fast_evaluate!!(
-        [rng::Random.AbstractRNG,]
-        model::Model,
-        strategy::AbstractInitStrategy,
-        accs::AccumulatorTuple,
-    )
-
-Evaluate a model using parameters obtained via `strategy`, and only computing the results in
-the provided accumulators.
-
-It is assumed that the accumulators passed in have been initialised to appropriate values,
-as this function will not reset them. The default constructors for each accumulator will do
-this for you correctly.
-
-Returns a tuple of the model's return value, plus an `OnlyAccsVarInfo`. Note that the `accs`
-argument may be mutated (depending on how the accumulators are implemented); hence the `!!`
-in the function name.
-"""
-@inline function fast_evaluate!!(
-    # Note that this `@inline` is mandatory for performance. If it's not inlined, it leads
-    # to extra allocations (even for trivial models) and much slower runtime.
-    rng::Random.AbstractRNG,
-    model::Model,
-    strategy::AbstractInitStrategy,
-    accs::AccumulatorTuple,
-)
-    ctx = InitContext(rng, strategy)
-    model = DynamicPPL.setleafcontext(model, ctx)
-    # 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
-    # here.
-    # TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
-    # it _should_ do, but this is wrong regardless.
-    # https://github.com/TuringLang/DynamicPPL.jl/issues/1086
-    vi = if Threads.nthreads() > 1
-        param_eltype = DynamicPPL.get_param_eltype(strategy)
-        accs = map(accs) do acc
-            DynamicPPL.convert_eltype(float_type_with_fallback(param_eltype), acc)
-        end
-        ThreadSafeVarInfo(OnlyAccsVarInfo(accs))
-    else
-        OnlyAccsVarInfo(accs)
-    end
-    return DynamicPPL._evaluate!!(model, vi)
-end
-@inline function fast_evaluate!!(
-    model::Model, strategy::AbstractInitStrategy, accs::AccumulatorTuple
-)
-    # This `@inline` is also mandatory for performance
-    return fast_evaluate!!(Random.default_rng(), model, strategy, accs)
-end
-
 """
     DynamicPPL.LogDensityFunction(
         model::Model,
@@ -154,9 +100,9 @@ metadata can often be quite wasteful. In particular, it is very common that the
 we care about from model evaluation are those which are stored in accumulators, such as log
 probability densities, or `ValuesAsInModel`.
 
-To avoid this issue, we use `OnlyAccsVarInfo`, which is a VarInfo that only contains
-accumulators. It implements enough of the `AbstractVarInfo` interface to not error during
-model evaluation.
+To avoid this issue, instead of evaluating a model with a full `VarInfo`, we use just an
+`OnlyAccsVarInfo`. It implements enough of the `AbstractVarInfo` interface to not error
+during model evaluation.
 
 Because `OnlyAccsVarInfo` does not store any parameter values, when evaluating a model with
 it, it is mandatory that parameters are provided from outside the VarInfo, namely via
@@ -199,6 +145,7 @@ struct LogDensityFunction{
     F<:Function,
     N<:NamedTuple,
     ADP<:Union{Nothing,DI.GradientPrep},
+    AT<:AccumulatorTuple,
 }
     model::M
     adtype::AD
@@ -207,13 +154,15 @@ struct LogDensityFunction{
     _varname_ranges::Dict{VarName,RangeAndLinked}
     _adprep::ADP
     _dim::Int
+    _accs::AccumulatorTuple
 
     function LogDensityFunction(
         model::Model,
         getlogdensity::Function=getlogjoint_internal,
         varinfo::AbstractVarInfo=VarInfo(model);
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing,
     )
+        accs = fast_ldf_accs(getlogdensity)
         # Figure out which variable corresponds to which index, and
         # which variables are linked.
         all_iden_ranges, all_ranges = get_ranges_and_linked(varinfo)
@@ -226,7 +175,7 @@ struct LogDensityFunction{
             # Make backend-specific tweaks to the adtype
             adtype = DynamicPPL.tweak_adtype(adtype, model, varinfo)
             DI.prepare_gradient(
-                LogDensityAt(model, getlogdensity, all_iden_ranges, all_ranges),
+                LogDensityAt(model, getlogdensity, all_iden_ranges, all_ranges, accs),
                 adtype,
                 x,
             )
@@ -237,8 +186,9 @@ struct LogDensityFunction{
             typeof(getlogdensity),
             typeof(all_iden_ranges),
             typeof(prep),
+            typeof(accs),
         }(
-            model, adtype, getlogdensity, all_iden_ranges, all_ranges, prep, dim
+            model, adtype, getlogdensity, all_iden_ranges, all_ranges, prep, dim, accs
         )
     end
 end
@@ -268,21 +218,25 @@ struct LogDensityAt{M<:Model,F<:Function,N<:NamedTuple}
     getlogdensity::F
     iden_varname_ranges::N
     varname_ranges::Dict{VarName,RangeAndLinked}
+    accs::AccumulatorTuple
 end
 function (f::LogDensityAt)(params::AbstractVector{<:Real})
     strategy = InitFromParams(
         VectorWithRanges(f.iden_varname_ranges, f.varname_ranges, params), nothing
     )
-    accs = fast_ldf_accs(f.getlogdensity)
-    _, vi = DynamicPPL.fast_evaluate!!(f.model, strategy, accs)
+    _, vi = DynamicPPL.init!!(f.model, OnlyAccsVarInfo(f.accs), strategy)
     return f.getlogdensity(vi)
 end
 
 function LogDensityProblems.logdensity(
     ldf::LogDensityFunction, params::AbstractVector{<:Real}
 )
     return LogDensityAt(
-        ldf.model, ldf._getlogdensity, ldf._iden_varname_ranges, ldf._varname_ranges
+        ldf.model,
+        ldf._getlogdensity,
+        ldf._iden_varname_ranges,
+        ldf._varname_ranges,
+        ldf._accs,
     )(
         params
     )
@@ -293,7 +247,11 @@ function LogDensityProblems.logdensity_and_gradient(
 )
     return DI.value_and_gradient(
         LogDensityAt(
-            ldf.model, ldf._getlogdensity, ldf._iden_varname_ranges, ldf._varname_ranges
+            ldf.model,
+            ldf._getlogdensity,
+            ldf._iden_varname_ranges,
+            ldf._varname_ranges,
+            ldf._accs,
         ),
         ldf._adprep,
         ldf.adtype,