Threaded assume

ext/DynamicPPLForwardDiffExt.jl

@@ -8,12 +8,8 @@ use_dynamicppl_tag(::ADTypes.AutoForwardDiff{<:Any,Nothing}) = true
 use_dynamicppl_tag(::ADTypes.AutoForwardDiff) = false
 
 function DynamicPPL.tweak_adtype(
-    ad::ADTypes.AutoForwardDiff{chunk_size},
-    ::DynamicPPL.Model,
-    vi::DynamicPPL.AbstractVarInfo,
+    ad::ADTypes.AutoForwardDiff{chunk_size}, ::DynamicPPL.Model, params::AbstractVector
 ) where {chunk_size}
-    params = vi[:]
-
     # Use DynamicPPL tag to improve stack traces
     # https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/
     # NOTE: DifferentiationInterface disables tag checking if the

src/chains.jl

@@ -156,13 +156,15 @@ via `unflatten` plus re-evaluation. It is faster for two reasons:
 """
 function ParamsWithStats(
     param_vector::AbstractVector,
-    ldf::DynamicPPL.LogDensityFunction,
+    ldf::DynamicPPL.LogDensityFunction{Tlink},
     stats::NamedTuple=NamedTuple();
     include_colon_eq::Bool=true,
     include_log_probs::Bool=true,
-)
+) where {Tlink}
     strategy = InitFromParams(
-        VectorWithRanges(ldf._iden_varname_ranges, ldf._varname_ranges, param_vector),
+        VectorWithRanges{Tlink}(
+            ldf._iden_varname_ranges, ldf._varname_ranges, param_vector
+        ),
         nothing,
     )
     accs = if include_log_probs

src/contexts/init.jl

@@ -214,7 +214,7 @@ struct RangeAndLinked
 end
 
 """
-    VectorWithRanges(
+    VectorWithRanges{Tlink}(
         iden_varname_ranges::NamedTuple,
         varname_ranges::Dict{VarName,RangeAndLinked},
         vect::AbstractVector{<:Real},
@@ -231,13 +231,19 @@ 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}}
+struct VectorWithRanges{Tlink,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
+
+    function VectorWithRanges{Tlink}(
+        iden_varname_ranges::N, varname_ranges::Dict{VarName,RangeAndLinked}, vect::T
+    ) where {Tlink,N,T}
+        return new{Tlink,N,T}(iden_varname_ranges, varname_ranges, vect)
+    end
 end
 
 function _get_range_and_linked(
@@ -252,11 +258,15 @@ function init(
     ::Random.AbstractRNG,
     vn::VarName,
     dist::Distribution,
-    p::InitFromParams{<:VectorWithRanges},
-)
+    p::InitFromParams{<:VectorWithRanges{T}},
+) where {T}
     vr = p.params
     range_and_linked = _get_range_and_linked(vr, vn)
-    transform = if range_and_linked.is_linked
+    # T can either be `nothing` (i.e., link status is mixed, in which
+    # case we use the stored link status), or `true` / `false`, which
+    # indicates that all variables are linked / unlinked.
+    linked = isnothing(T) ? range_and_linked.is_linked : T
+    transform = if linked
         from_linked_vec_transform(dist)
     else
         from_vec_transform(dist)