FastLDF / InitContext unified

HISTORY.md

@@ -21,6 +21,21 @@ You should not need to use these directly, please use `AbstractPPL.condition` an
 
 Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
 
+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

@@ -170,6 +170,12 @@ DynamicPPL.prefix
 
 ## Utilities
 
+`typed_identity` is the same as `identity`, but with an overload for `with_logabsdet_jacobian` that ensures that it never errors.
+
+```@docs
+typed_identity
+```
+
 It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function.
 
 ```@docs
@@ -517,10 +523,12 @@ InitFromParams
 ```
 
 If you wish to write your own, you have to subtype [`DynamicPPL.AbstractInitStrategy`](@ref) and implement the `init` method.
+In very rare situations, you may also need to implement `get_param_eltype`, which defines the element type of the parameters generated by the strategy.
 
 ```@docs
-DynamicPPL.AbstractInitStrategy
-DynamicPPL.init
+AbstractInitStrategy
+init
+get_param_eltype
 ```
 
 ### Choosing a suitable VarInfo

ext/DynamicPPLEnzymeCoreExt.jl

@@ -1,16 +1,15 @@
 module DynamicPPLEnzymeCoreExt
 
-if isdefined(Base, :get_extension)
-    using DynamicPPL: DynamicPPL
-    using EnzymeCore
-else
-    using ..DynamicPPL: DynamicPPL
-    using ..EnzymeCore
-end
+using DynamicPPL: DynamicPPL
+using EnzymeCore
 
 # Mark is_transformed as having 0 derivative. The `nothing` return value is not significant, Enzyme
 # only checks whether such a method exists, and never runs it.
 @inline EnzymeCore.EnzymeRules.inactive(::typeof(DynamicPPL.is_transformed), args...) =
     nothing
+# Likewise for get_range_and_linked.
+@inline EnzymeCore.EnzymeRules.inactive(
+    ::typeof(DynamicPPL._get_range_and_linked), args...
+) = nothing
 
 end

ext/DynamicPPLMooncakeExt.jl

@@ -5,5 +5,8 @@ using Mooncake: Mooncake
 
 # This is purely an optimisation.
 Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(is_transformed),Vararg}
+Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{
+    typeof(DynamicPPL._get_range_and_linked),Vararg
+}
 
 end # module

src/DynamicPPL.jl

@@ -84,8 +84,8 @@ export AbstractVarInfo,
     # Compiler
     @model,
     # Utilities
-    init,
     OrderedDict,
+    typed_identity,
     # Model
     Model,
     getmissings,
@@ -113,6 +113,8 @@ export AbstractVarInfo,
     InitFromPrior,
     InitFromUniform,
     InitFromParams,
+    init,
+    get_param_eltype,
     # Pseudo distributions
     NamedDist,
     NoDist,
@@ -193,6 +195,7 @@ include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")
 include("simple_varinfo.jl")
+include("onlyaccs.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("logdensityfunction.jl")

src/compiler.jl

@@ -718,14 +718,15 @@ end
 # TODO(mhauru) matchingvalue has methods that can accept both types and values. Why?
 # TODO(mhauru) This function needs a more comprehensive docstring.
 """
-    matchingvalue(vi, value)
+    matchingvalue(param_eltype, value)
 
-Convert the `value` to the correct type for the `vi` object.
+Convert the `value` to the correct type, given the element type of the parameters
+being used to evaluate the model.
 """
-function matchingvalue(vi, value)
+function matchingvalue(param_eltype, value)
     T = typeof(value)
     if hasmissing(T)
-        _value = convert(get_matching_type(vi, T), value)
+        _value = convert(get_matching_type(param_eltype, T), value)
         # TODO(mhauru) Why do we make a deepcopy, even though in the !hasmissing branch we
         # are happy to return `value` as-is?
         if _value === value
@@ -738,29 +739,30 @@ function matchingvalue(vi, value)
     end
 end
 
-function matchingvalue(vi, value::FloatOrArrayType)
-    return get_matching_type(vi, value)
+function matchingvalue(param_eltype, value::FloatOrArrayType)
+    return get_matching_type(param_eltype, value)
 end
-function matchingvalue(vi, ::TypeWrap{T}) where {T}
-    return TypeWrap{get_matching_type(vi, T)}()
+function matchingvalue(param_eltype, ::TypeWrap{T}) where {T}
+    return TypeWrap{get_matching_type(param_eltype, T)}()
 end
 
 # TODO(mhauru) This function needs a more comprehensive docstring. What is it for?
 """
-    get_matching_type(vi, ::TypeWrap{T}) where {T}
+    get_matching_type(param_eltype, ::TypeWrap{T}) where {T}
 
-Get the specialized version of type `T` for `vi`.
+Get the specialized version of type `T`, given an element type of the parameters
+being used to evaluate the model.
 """
 get_matching_type(_, ::Type{T}) where {T} = T
-function get_matching_type(vi, ::Type{<:Union{Missing,AbstractFloat}})
-    return Union{Missing,float_type_with_fallback(eltype(vi))}
+function get_matching_type(param_eltype, ::Type{<:Union{Missing,AbstractFloat}})
+    return Union{Missing,float_type_with_fallback(param_eltype)}
 end
-function get_matching_type(vi, ::Type{<:AbstractFloat})
-    return float_type_with_fallback(eltype(vi))
+function get_matching_type(param_eltype, ::Type{<:AbstractFloat})
+    return float_type_with_fallback(param_eltype)
 end
-function get_matching_type(vi, ::Type{<:Array{T,N}}) where {T,N}
-    return Array{get_matching_type(vi, T),N}
+function get_matching_type(param_eltype, ::Type{<:Array{T,N}}) where {T,N}
+    return Array{get_matching_type(param_eltype, T),N}
 end
-function get_matching_type(vi, ::Type{<:Array{T}}) where {T}
-    return Array{get_matching_type(vi, T)}
+function get_matching_type(param_eltype, ::Type{<:Array{T}}) where {T}
+    return Array{get_matching_type(param_eltype, T)}
 end