Make unflatten more type stable, and add a test for it

Author: mhauru

src/simple_varinfo.jl

@@ -264,8 +264,12 @@ function unflatten(svi::SimpleVarInfo, x::AbstractVector)
     vals = unflatten(svi.values, x)
     # TODO(mhauru) See comment in unflatten in src/varinfo.jl for why this conversion is
     # required but undesireable.
-    T = float_type_with_fallback(eltype(x))
-    accs = map(acc -> convert_eltype(T, acc), getaccs(svi))
+    # The below line is finicky for type stability. For instance, assigning the eltype to
+    # convert to into an intermediate variable makes this unstable (constant propagation)
+    # fails. Take care when editing.
+    accs = map(
+        acc -> convert_eltype(float_type_with_fallback(eltype(x)), acc), getaccs(svi)
+    )
     return SimpleVarInfo(vals, accs, svi.transformation)
 end
 

src/varinfo.jl

@@ -444,8 +444,13 @@ function unflatten(vi::VarInfo, x::AbstractVector)
     # element type of ThreadSafeVarInfo.accs_by_thread. However, doing this conversion here
     # messes with cases like using Float32 of logprobs and Float64 for x. Also, this is just
     # plain ugly and hacky.
-    T = float_type_with_fallback(eltype(x))
-    accs = map(acc -> convert_eltype(T, acc), deepcopy(getaccs(vi)))
+    # The below line is finicky for type stability. For instance, assigning the eltype to
+    # convert to into an intermediate variable makes this unstable (constant propagation)
+    # fails. Take care when editing.
+    accs = map(
+        acc -> convert_eltype(float_type_with_fallback(eltype(x)), acc),
+        deepcopy(getaccs(vi)),
+    )
     return VarInfo(md, accs)
 end
 

test/varinfo.jl

@@ -732,6 +732,29 @@ end
         end
     end
 
+    @testset "unflatten type stability" begin
+        @model function demo(y)
+            x ~ Normal()
+            y ~ Normal(x, 1)
+            return nothing
+        end
+
+        model = demo(0.0)
+        varinfos = DynamicPPL.TestUtils.setup_varinfos(
+            model, (; x=1.0), (@varname(x),); include_threadsafe=true
+        )
+        @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
+            # Skip the severely inconcrete `SimpleVarInfo` types, since checking for type
+            # stability for them doesn't make much sense anyway.
+            if varinfo isa SimpleVarInfo{OrderedDict{Any,Any}} ||
+                varinfo isa
+               DynamicPPL.ThreadSafeVarInfo{<:SimpleVarInfo{OrderedDict{Any,Any}}}
+                continue
+            end
+            @inferred DynamicPPL.unflatten(varinfo, varinfo[:])
+        end
+    end
+
     @testset "subset" begin
         @model function demo_subsetting_varinfo(::Type{TV}=Vector{Float64}) where {TV}
             s ~ InverseGamma(2, 3)