OnlyAccsVarInfo -> AccumulatorTuple; fast_evaluate!! -> init!!

Author: penelopeysm

docs/src/api.md

@@ -66,13 +66,6 @@ The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) inte
 LogDensityFunction
 ```
 
-Internally, this is accomplished using:
-
-```@docs
-OnlyAccsVarInfo
-fast_evaluate!!
-```
-
 ## Condition and decondition
 
 A [`Model`](@ref) can be conditioned on a set of observations with [`AbstractPPL.condition`](@ref) or its alias [`|`](@ref).
@@ -371,6 +364,7 @@ The subtypes of [`AbstractVarInfo`](@ref) store the cumulative log prior and log
 
 ```@docs
 AbstractAccumulator
+AccumulatorTuple
 ```
 
 DynamicPPL provides the following default accumulators.

src/DynamicPPL.jl

@@ -51,6 +51,7 @@ export AbstractVarInfo,
     LogLikelihoodAccumulator,
     LogPriorAccumulator,
     LogJacobianAccumulator,
+    AccumulatorTuple,
     push!!,
     empty!!,
     subset,
@@ -92,10 +93,8 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
-    # LogDensityFunction and fasteval
+    # LogDensityFunction
     LogDensityFunction,
-    fast_evaluate!!,
-    OnlyAccsVarInfo,
     # Leaf contexts
     AbstractContext,
     contextualize,
@@ -197,7 +196,7 @@ include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")
 include("simple_varinfo.jl")
-include("onlyaccs.jl")
+include("accs_as_varinfo.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("fasteval.jl")

src/accs_as_varinfo.jl

@@ -0,0 +1,17 @@
+DynamicPPL.maybe_invlink_before_eval!!(at::AccumulatorTuple, ::Model) = at
+DynamicPPL.getaccs(at::AccumulatorTuple) = at
+DynamicPPL.setaccs!!(::AccumulatorTuple, new_at::AccumulatorTuple) = new_at
+
+function tilde_assume!!(
+    ctx::InitContext,
+    dist::Distribution,
+    vn::VarName,
+    vi::Union{AccumulatorTuple,ThreadSafeVarInfo{<:AccumulatorTuple}},
+)
+    # For AccumulatorTuple, since we don't need to write into the metadata part of the
+    # VarInfo, we can cut out a lot of the code.
+    val, transform = DynamicPPL.init(ctx.rng, vn, dist, ctx.strategy)
+    x, inv_logjac = Bijectors.with_logabsdet_jacobian(transform, val)
+    vi = DynamicPPL.accumulate_assume!!(vi, x, -inv_logjac, vn, dist)
+    return x, vi
+end

src/accumulators.jl

@@ -132,7 +132,7 @@ used by various AD backends, should implement a method for this function.
 convert_eltype(::Type, acc::AbstractAccumulator) = acc
 
 """
-    AccumulatorTuple{N,T<:NamedTuple}
+    AccumulatorTuple{N,T<:NamedTuple} <: AbstractVarInfo
 
 A collection of accumulators, stored as a `NamedTuple` of length `N`
 
@@ -144,8 +144,18 @@ constraint that the name in the tuple for each accumulator `acc` must be
 The constructor can be called with a tuple or a `VarArgs` of `AbstractAccumulators`. The
 names will be generated automatically. One can also call the constructor with a `NamedTuple`
 but the names in the argument will be discarded in favour of the generated ones.
+
+`AccumulatorTuple` is a subtype of `AbstractVarInfo`, but in reality it only implements the
+part of the interface that deals with accumulators. That is, it implements `getaccs`,
+`setaccs!!`, and `maybe_invlink_before_eval!!`. These definitions allow several other
+'derived' methods to work automatically, such as `getlogjoint`, `accumulate_assume!!`, and
+`accumulate_observe!!`. Note that `maybe_invlink_before_eval!!` should not be needed if/when
+SimpleVarInfo is removed.
+
+Unfortunately, because the necessary forward definitions are not present at this stage,
+these methods have to be defined in a separate file (`src/accs_as_varinfo.jl`).
 """
-struct AccumulatorTuple{N,T<:NamedTuple}
+struct AccumulatorTuple{N,T<:NamedTuple} <: AbstractVarInfo
     nt::T
 
     function AccumulatorTuple(t::T) where {N,T<:NTuple{N,AbstractAccumulator}}

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 `AccumulatorTuple` 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

@@ -142,7 +142,7 @@ This tells you how to obtain values for the random variable `vn` from `p.params`
 the last argument is `InitFromParams(params)`, not just `params` itself. Please see the
 docstring of [`DynamicPPL.init`](@ref) for more information on the expected behaviour.
 
-If you only use `InitFromParams` with `DynamicPPL.OnlyAccsVarInfo`, as is usually the case,
+If you only use `InitFromParams` with `DynamicPPL.AccumulatorTuple`, as is usually the case,
 then you will not need to implement anything else. So far, this is the same as you would do
 for creating any new `AbstractInitStrategy` subtype.
 
@@ -311,12 +311,13 @@ function tilde_assume!!(
         # calculation wastes a lot of time going from linked vectorised -> unlinked ->
         # 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.
+        # However, `VectorWithRanges` is only really used with `AccumulatorTuple`, in which
+        # case this method is never hit (since there's a special method for it, in
+        # `src/accs_as_varinfo.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

@@ -10,7 +10,6 @@ using DynamicPPL:
     Model,
     ThreadSafeVarInfo,
     VarInfo,
-    OnlyAccsVarInfo,
     RangeAndLinked,
     VectorWithRanges,
     Metadata,
@@ -29,60 +28,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,11 +99,11 @@ 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
+`AccumulatorTuple`. 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
+Because `AccumulatorTuple` 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
 `InitContext`.
 
@@ -274,7 +219,7 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real})
         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, accs, strategy)
     return f.getlogdensity(vi)
 end