Make FastLDF the default (#1139)

* Make FastLDF the default

* Add miscellaneous LogDensityProblems tests

* Use `init!!` instead of `fast_evaluate!!`

* Rename files, rebalance tests

Author: penelopeysm

HISTORY.md

@@ -4,6 +4,17 @@
 
 ### Breaking changes
 
+#### Fast Log Density Functions
+
+This version provides a reimplementation of `LogDensityFunction` that provides performance improvements on the order of 2–10× for both model evaluation as well as automatic differentiation.
+Exact speedups depend on the model size: larger models have less significant speedups because the bulk of the work is done in calls to `logpdf`.
+
+For more information about how this is accomplished, please see https://github.com/TuringLang/DynamicPPL.jl/pull/1113 as well as the `src/fasteval.jl` file, which contains extensive comments.
+
+As a result of this change, `LogDensityFunction` no longer stores a VarInfo inside it.
+In general, if `ldf` is a `LogDensityFunction`, it is now only valid to access `ldf.model` and `ldf.adtype`.
+If you were previously relying on this behaviour, you will need to store a VarInfo separately.
+
 #### Parent and leaf contexts
 
 The `DynamicPPL.NodeTrait` function has been removed.
@@ -24,18 +35,6 @@ Removed the method `returned(::Model, values, keys)`; please use `returned(::Mod
 The method `DynamicPPL.init` (for implementing `AbstractInitStrategy`) now has a different signature: it must return a tuple of the generated value, plus a transform function that maps it back to unlinked space.
 This is a generalisation of the previous behaviour, where `init` would always return an unlinked value (in effect forcing the transform to be the identity function).
 
-### Other changes
-
-#### FastLDF
-
-Added `DynamicPPL.Experimental.FastLDF`, a version of `LogDensityFunction` that provides performance improvements on the order of 2–10× for both model evaluation as well as automatic differentiation.
-Exact speedups depend on the model size: larger models have less significant speedups because the bulk of the work is done in calls to `logpdf`.
-
-Please note that `FastLDF` is currently considered internal and its API may change without warning.
-We intend to replace `LogDensityFunction` with `FastLDF` in a release in the near future, but until then we recommend not using it.
-
-For more information about `FastLDF`, please see https://github.com/TuringLang/DynamicPPL.jl/pull/1113 as well as the `src/fasteval.jl` file, which contains extensive comments.
-
 ## 0.38.9
 
 Remove warning when using Enzyme as the AD backend.

docs/src/api.md

@@ -66,6 +66,12 @@ The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) inte
 LogDensityFunction
 ```
 
+Internally, this is accomplished using [`init!!`](@ref) on:
+
+```@docs
+OnlyAccsVarInfo
+```
+
 ## Condition and decondition
 
 A [`Model`](@ref) can be conditioned on a set of observations with [`AbstractPPL.condition`](@ref) or its alias [`|`](@ref).
@@ -510,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.

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -6,8 +6,13 @@ using MarginalLogDensities: MarginalLogDensities
 # A thin wrapper to adapt a DynamicPPL.LogDensityFunction to the interface expected by
 # MarginalLogDensities. It's helpful to have a struct so that we can dispatch on its type
 # below.
-struct LogDensityFunctionWrapper{L<:DynamicPPL.LogDensityFunction}
+struct LogDensityFunctionWrapper{
+    L<:DynamicPPL.LogDensityFunction,V<:DynamicPPL.AbstractVarInfo
+}
     logdensity::L
+    # This field is used only to reconstruct the VarInfo later on; it's not needed for the
+    # actual log-density evaluation.
+    varinfo::V
 end
 function (lw::LogDensityFunctionWrapper)(x, _)
     return LogDensityProblems.logdensity(lw.logdensity, x)
@@ -101,7 +106,7 @@ function DynamicPPL.marginalize(
     # Construct the marginal log-density model.
     f = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
     mld = MarginalLogDensities.MarginalLogDensity(
-        LogDensityFunctionWrapper(f), varinfo[:], varindices, (), method; kwargs...
+        LogDensityFunctionWrapper(f, varinfo), varinfo[:], varindices, (), method; kwargs...
     )
     return mld
 end
@@ -190,7 +195,7 @@ function DynamicPPL.VarInfo(
     unmarginalized_params::Union{AbstractVector,Nothing}=nothing,
 )
     # Extract the original VarInfo. Its contents will in general be junk.
-    original_vi = mld.logdensity.logdensity.varinfo
+    original_vi = mld.logdensity.varinfo
     # Extract the stored parameters, which includes the modes for any marginalized
     # parameters
     full_params = MarginalLogDensities.cached_params(mld)

src/DynamicPPL.jl

@@ -92,8 +92,12 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
+    # evaluation
+    evaluate!!,
+    init!!,
     # LogDensityFunction
     LogDensityFunction,
+    OnlyAccsVarInfo,
     # Leaf contexts
     AbstractContext,
     contextualize,

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,
@@ -156,7 +156,7 @@ via `unflatten` plus re-evaluation. It is faster for two reasons:
 """
 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, OnlyAccsVarInfo(AccumulatorTuple(accs)), strategy)
     params = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
     if include_log_probs
         stats = merge(

src/experimental.jl

@@ -2,8 +2,6 @@ module Experimental
 
 using DynamicPPL: DynamicPPL
 
-include("fasteval.jl")
-
 # This file only defines the names of the functions, and their docstrings. The actual implementations are in `ext/DynamicPPLJETExt.jl`, since we don't want to depend on JET.jl other than as a weak dependency.
 """
     is_suitable_varinfo(model::Model, varinfo::AbstractVarInfo; kwargs...)