Fix issues with RangeAndLinked and VNT

Author: mhauru

ext/DynamicPPLMarginalLogDensitiesExt.jl

@@ -1,6 +1,6 @@
 module DynamicPPLMarginalLogDensitiesExt
 
-using DynamicPPL: DynamicPPL, LogDensityProblems, VarName
+using DynamicPPL: DynamicPPL, LogDensityProblems, VarName, RangeAndLinked
 using MarginalLogDensities: MarginalLogDensities
 
 # A thin wrapper to adapt a DynamicPPL.LogDensityFunction to the interface expected by
@@ -105,11 +105,9 @@ function DynamicPPL.marginalize(
     ldf = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
     # Determine the indices for the variables to marginalise out.
     varindices = mapreduce(vcat, marginalized_varnames) do vn
-        if DynamicPPL.getoptic(vn) === identity
-            ldf._iden_varname_ranges[DynamicPPL.getsym(vn)].range
-        else
-            ldf._varname_ranges[vn].range
-        end
+        # The type assertion helps in cases where the model is type unstable and thus
+        # `varname_ranges` may have an abstract element type.
+        (ldf._varname_ranges[vn]::RangeAndLinked).range
     end
     mld = MarginalLogDensities.MarginalLogDensity(
         LogDensityFunctionWrapper(ldf, varinfo),

src/contexts/init.jl

@@ -206,14 +206,16 @@ an unlinked value.
 
 $(TYPEDFIELDS)
 """
-struct RangeAndLinked
+struct RangeAndLinked{T<:Tuple}
     # indices that the variable corresponds to in the vectorised parameter
     range::UnitRange{Int}
     # whether it's linked
     is_linked::Bool
+    # original size of the variable before vectorisation
+    original_size::T
 end
 
-Base.size(ral::RangeAndLinked) = size(ral.range)
+Base.size(ral::RangeAndLinked) = ral.original_size
 
 """
     VectorWithRanges{Tlink}(
@@ -249,7 +251,12 @@ struct VectorWithRanges{Tlink,VNT<:VarNamedTuple,T<:AbstractVector{<:Real}}
 end
 
 function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
-    return vr.varname_ranges[vn]
+    # The type assertion does nothing if VectorWithRanges has concrete element types, as is
+    # the case for all type stable models. However, if the model is not type stable,
+    # vr.varname_ranges[vn] may infer to have type `Any`. In this case it is helpful to
+    # assert that it is a RangeAndLinked, because even though it remains non-concrete,
+    # it'll allow the compiler to infer the types of `range` and `is_linked`.
+    return vr.varname_ranges[vn]::RangeAndLinked
 end
 function init(
     ::Random.AbstractRNG,

src/logdensityfunction.jl

@@ -330,7 +330,10 @@ function get_ranges_and_linked_metadata(md::Metadata, start_offset::Int)
     for (vn, idx) in md.idcs
         is_linked = md.is_transformed[idx]
         range = md.ranges[idx] .+ (start_offset - 1)
-        all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
+        orig_size = varnamesize(vn)
+        all_ranges = BangBang.setindex!!(
+            all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
+        )
         offset += length(range)
     end
     return all_ranges, offset
@@ -341,7 +344,10 @@ function get_ranges_and_linked_metadata(vnv::VarNamedVector, start_offset::Int)
     for (vn, idx) in vnv.varname_to_index
         is_linked = vnv.is_unconstrained[idx]
         range = vnv.ranges[idx] .+ (start_offset - 1)
-        all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
+        orig_size = varnamesize(vn)
+        all_ranges = BangBang.setindex!!(
+            all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
+        )
         offset += length(range)
     end
     return all_ranges, offset

src/varname.jl

@@ -41,3 +41,28 @@ Possibly existing indices of `varname` are neglected.
 ) where {s,missings,_F,_a,_T}
     return s in missings
 end
+
+# TODO(mhauru) This should probably be Base.size(::VarName) in AbstractPPL.
+"""
+    varnamesize(vn::VarName)
+
+Return the size of the object referenced by this VarName.
+
+```jldoctest
+julia> varnamesize(@varname(a))
+()
+
+julia> varnamesize(@varname(b[1:3, 2]))
+(3,)
+
+julia> varnamesize(@varname(c.d[4].e[3, 2:5, 2, 1:4, 1]))
+(4, 4)
+"""
+function varnamesize(vn::VarName)
+    l = AbstractPPL._last(vn.optic)
+    if l isa Accessors.IndexLens
+        return reduce((x, y) -> tuple(x..., y...), map(size, l.indices))
+    else
+        return ()
+    end
+end

src/varnamedtuple.jl

@@ -352,7 +352,13 @@ function _concretise_eltype!!(pa::PartialArray)
     if isconcretetype(eltype(pa))
         return pa
     end
-    new_et = promote_type((typeof(pa.data[i]) for i in eachindex(pa.mask) if pa.mask[i])...)
+    # We could use promote_type here, instead of typejoin. However, that would e.g.
+    # cause Ints to be converted to Float64s, since
+    # promote_type(Int, Float64) == Float64, which can cause problems. See
+    # https://github.com/TuringLang/DynamicPPL.jl/pull/1098#discussion_r2472636188.
+    # Base.promote_typejoin would be like typejoin, but creates Unions out of Nothing
+    # and Missing, rather than falling back on Any. However, it's not exported.
+    new_et = typejoin((typeof(pa.data[i]) for i in eachindex(pa.mask) if pa.mask[i])...)
     # TODO(mhauru) Should we check as below, or rather isconcretetype(new_et)?
     # In other words, does it help to be more concrete, even if we aren't fully concrete?
     if new_et === eltype(pa)
@@ -588,8 +594,8 @@ function _setindex!!(pa::PartialArray, value, inds::Vararg{INDEX_TYPES})
         if size(value) != inds_size
             throw(
                 DimensionMismatch(
-                    "Assigned value has size $(size(value)), which does not match the size " *
-                    "implied by the indices $(map(x -> _length_needed(x), inds)).",
+                    "Assigned value has size $(size(value)), which does not match the " *
+                    "size implied by the indices $(map(x -> _length_needed(x), inds)).",
                 ),
             )
         end
@@ -659,7 +665,14 @@ function _merge_recursive(pa1::PartialArray, pa2::PartialArray)
             result
         else
             # Neither is strictly bigger than the other.
-            et = promote_type(eltype(pa1), eltype(pa2))
+            # We could use promote_type here, instead of typejoin. However, that would e.g.
+            # cause Ints to be converted to Float64s, since
+            # promote_type(Int, Float64) == Float64, which can cause problems. See
+            # https://github.com/TuringLang/DynamicPPL.jl/pull/1098#discussion_r2472636188.
+            # Base.promote_typejoin would be like typejoin, but creates Unions out of
+            # Nothing and Missing, rather than falling back on Any. However, it's not
+            # exported.
+            et = typejoin(eltype(pa1), eltype(pa2))
             new_data = Array{et,num_dims}(undef, merge_size)
             new_mask = fill(false, merge_size)
             result = PartialArray(new_data, new_mask)