Refactor FastLDF to use InitContext

Author: penelopeysm

ext/DynamicPPLEnzymeCoreExt.jl

@@ -9,7 +9,7 @@ using EnzymeCore
     nothing
 # Likewise for get_range_and_linked.
 @inline EnzymeCore.EnzymeRules.inactive(
-    ::typeof(DynamicPPL.Experimental.get_range_and_linked), args...
+    ::typeof(DynamicPPL._get_range_and_linked), args...
 ) = nothing
 
 end

ext/DynamicPPLMooncakeExt.jl

@@ -6,7 +6,7 @@ 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.Experimental.get_range_and_linked),Vararg
+    typeof(DynamicPPL._get_range_and_linked),Vararg
 }
 
 end # module

src/DynamicPPL.jl

@@ -188,6 +188,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/contexts/init.jl

@@ -9,16 +9,24 @@ Any subtype of `AbstractInitStrategy` must implement the
 """
 abstract type AbstractInitStrategy end
 
+"""
+    InitValue{T,F}
+
+A wrapper type representing a value of type `T`. The function `F` indicates what transform
+to apply to the value to convert it back to the unlinked space. If `value` is already in
+unlinked space, then `transform` can be `identity`.
+"""
+struct InitValue{T,F}
+    value::T
+    transform::F
+end
+
 """
     init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, strategy::AbstractInitStrategy)
 
 Generate a new value for a random variable with the given distribution.
 
-!!! warning "Return values must be unlinked"
-    The values returned by `init` must always be in the untransformed space, i.e.,
-    they must be within the support of the original distribution. That means that,
-    for example, `init(rng, dist, u::InitFromUniform)` will in general return values that
-    are outside the range [u.lower, u.upper].
+This function must return a `InitValue`.
 """
 function init end
 
@@ -29,7 +37,7 @@ Obtain new values by sampling from the prior distribution.
 """
 struct InitFromPrior <: AbstractInitStrategy end
 function init(rng::Random.AbstractRNG, ::VarName, dist::Distribution, ::InitFromPrior)
-    return rand(rng, dist)
+    return InitValue(rand(rng, dist), identity)
 end
 
 """
@@ -69,7 +77,7 @@ function init(rng::Random.AbstractRNG, ::VarName, dist::Distribution, u::InitFro
     if x isa Array{<:Any,0}
         x = x[]
     end
-    return x
+    return InitValue(x, identity)
 end
 
 """
@@ -93,19 +101,20 @@ will be thrown. The default for `fallback` is `InitFromPrior()`.
 struct InitFromParams{P,S<:Union{AbstractInitStrategy,Nothing}} <: AbstractInitStrategy
     params::P
     fallback::S
+
     function InitFromParams(
-        params::AbstractDict{<:VarName},
-        fallback::Union{AbstractInitStrategy,Nothing}=InitFromPrior(),
-    )
-        return new{typeof(params),typeof(fallback)}(params, fallback)
-    end
-    function InitFromParams(
-        params::NamedTuple, fallback::Union{AbstractInitStrategy,Nothing}=InitFromPrior()
-    )
-        return new{typeof(params),typeof(fallback)}(params, fallback)
+        params::P, fallback::Union{AbstractInitStrategy,Nothing}=InitFromPrior()
+    ) where {P}
+        return new{P,typeof(fallback)}(params, fallback)
     end
 end
-function init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, p::InitFromParams)
+
+function init(
+    rng::Random.AbstractRNG,
+    vn::VarName,
+    dist::Distribution,
+    p::InitFromParams{<:Union{AbstractDict{<:VarName},NamedTuple}},
+)
     # TODO(penelopeysm): It would be nice to do a check to make sure that all
     # of the parameters in `p.params` were actually used, and either warn or
     # error if they aren't. This is actually quite non-trivial though because
@@ -119,14 +128,82 @@ function init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, p::InitF
         else
             # TODO(penelopeysm): Since x is user-supplied, maybe we could also
             # check here that the type / size of x matches the dist?
-            x
+            InitValue(x, identity)
         end
     else
         p.fallback === nothing && error("No value was provided for the variable `$(vn)`.")
         init(rng, vn, dist, p.fallback)
     end
 end
 
+"""
+    RangeAndLinked
+
+Suppose we have vectorised parameters `params::AbstractVector{<:Real}`. Each random variable
+in the model will in general correspond to a sub-vector of `params`. This struct stores
+information about that range, as well as whether the sub-vector represents a linked value or
+an unlinked value.
+
+$(TYPEDFIELDS)
+"""
+struct RangeAndLinked
+    # indices that the variable corresponds to in the vectorised parameter
+    range::UnitRange{Int}
+    # whether it's linked
+    is_linked::Bool
+end
+
+"""
+    VectorWithRanges(
+        iden_varname_ranges::NamedTuple,
+        varname_ranges::Dict{VarName,RangeAndLinked},
+        vect::AbstractVector{<:Real},
+    )
+
+A struct that wraps a vector of parameter values, plus information about how random
+variables map to ranges in that vector.
+
+In the simplest case, this could be accomplished only with a single dictionary mapping
+VarNames to ranges and link status. However, for performance reasons, we separate out
+VarNames with identity optics into a NamedTuple (`iden_varname_ranges`). All
+non-identity-optic VarNames are stored in the `varname_ranges` Dict.
+
+It would be nice to improve the NamedTuple and Dict approach. See, e.g.
+https://github.com/TuringLang/DynamicPPL.jl/issues/1116.
+"""
+struct VectorWithRanges{N<:NamedTuple,T<:AbstractVector{<:Real}}
+    # This NamedTuple stores the ranges for identity VarNames
+    iden_varname_ranges::N
+    # This Dict stores the ranges for all other VarNames
+    varname_ranges::Dict{VarName,RangeAndLinked}
+    # The full parameter vector which we index into to get variable values
+    vect::T
+end
+
+function _get_range_and_linked(
+    vr::VectorWithRanges, ::VarName{sym,typeof(identity)}
+) where {sym}
+    return vr.iden_varname_ranges[sym]
+end
+function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
+    return vr.varname_ranges[vn]
+end
+function init(
+    ::Random.AbstractRNG,
+    vn::VarName,
+    dist::Distribution,
+    p::InitFromParams{<:VectorWithRanges},
+)
+    vr = p.params
+    range_and_linked = _get_range_and_linked(vr, vn)
+    transform = if range_and_linked.is_linked
+        from_linked_vec_transform(dist)
+    else
+        from_vec_transform(dist)
+    end
+    return InitValue((@view vr.vect[range_and_linked.range]), transform)
+end
+
 """
     InitContext(
             [rng::Random.AbstractRNG=Random.default_rng()],
@@ -156,27 +233,43 @@ function tilde_assume!!(
     ctx::InitContext, dist::Distribution, vn::VarName, vi::AbstractVarInfo
 )
     in_varinfo = haskey(vi, vn)
-    # `init()` always returns values in original space, i.e. possibly
-    # constrained
-    x = init(ctx.rng, vn, dist, ctx.strategy)
+    init_val = init(ctx.rng, vn, dist, ctx.strategy)
+    x, inv_logjac = with_logabsdet_jacobian(init_val.transform, init_val.value)
     # Determine whether to insert a transformed value into the VarInfo.
     # If the VarInfo alrady had a value for this variable, we will
     # keep the same linked status as in the original VarInfo. If not, we
     # check the rest of the VarInfo to see if other variables are linked.
     # is_transformed(vi) returns true if vi is nonempty and all variables in vi
     # are linked.
     insert_transformed_value = in_varinfo ? is_transformed(vi, vn) : is_transformed(vi)
-    y, logjac = if insert_transformed_value
-        with_logabsdet_jacobian(link_transform(dist), x)
+    val_to_insert, logjac = if insert_transformed_value
+        # Calculate the forward logjac and sum them up.
+        y, fwd_logjac = with_logabsdet_jacobian(link_transform(dist), x)
+        # Note that if we use VectorWithRanges with a full VarInfo, this double-Jacobian
+        # 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). So 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`,
+        # which is NOT the same as `inverse(link_transform)` (because there is an additional
+        # vectorisation step). We need `init` and `tilde_assume!!` to share this information
+        # but it's not clear right now how to do this. In my opinion, the most productive
+        # way forward would be to standardise the behaviour of bijectors so that we can have
+        # a clean separation between the linking and vectorisation parts of it.
+        y, inv_logjac + fwd_logjac
     else
-        x, zero(LogProbType)
+        x, inv_logjac
     end
     # Add the new value to the VarInfo. `push!!` errors if the value already
     # exists, hence the need for setindex!!.
     if in_varinfo
-        vi = setindex!!(vi, y, vn)
+        vi = setindex!!(vi, val_to_insert, vn)
     else
-        vi = push!!(vi, vn, y, dist)
+        vi = push!!(vi, vn, val_to_insert, dist)
     end
     # Neither of these set the `trans` flag so we have to do it manually if
     # necessary.