Fix for #596

src/compiler.jl

@@ -578,10 +578,42 @@ hasmissing(::Type{>:Missing}) = true
 hasmissing(::Type{<:AbstractArray{TA}}) where {TA} = hasmissing(TA)
 hasmissing(::Type{Union{}}) = false # issue #368
 
+"""
+    TypeWrap{T}
+
+A wrapper type used internally to make expressions such as `::Type{TV}` in the model arguments
+not ending up as a `DataType`.
+"""
+struct TypeWrap{T} end
+
+function arg_type_is_type(e)
+    return Meta.isexpr(e, :curly) && length(e.args) > 1 && e.args[1] === :Type
+end
+
 function splitarg_to_expr((arg_name, arg_type, is_splat, default))
     return is_splat ? :($arg_name...) : arg_name
 end
 
+"""
+    transform_args(args)
+
+Return transformed `args` used in both the model constructor and evaluator.
+
+Specifically, this replaces expressions of the form `::Type{TV}=Vector{Float64}`
+with `::TypeWrap{TV}=TypeWrap{Vector{Float64}}()` to avoid introducing `DataType`.
+"""
+function transform_args(args)
+    splitargs = map(args) do arg
+        arg_name, arg_type, is_splat, default = MacroTools.splitarg(arg)
+        return if arg_type_is_type(arg_type)
+            arg_name, :($TypeWrap{$(arg_type.args[2])}), is_splat, :($TypeWrap{$default}())
+        else
+            arg_name, arg_type, is_splat, default
+        end
+    end
+    return map(Base.splat(MacroTools.combinearg), splitargs)
+end
+
 function namedtuple_from_splitargs(splitargs)
     names = map(splitargs) do (arg_name, arg_type, is_splat, default)
         is_splat ? Symbol("#splat#$(arg_name)") : arg_name
@@ -599,7 +631,7 @@ is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#")
 Builds the output expression.
 """
 function build_output(modeldef, linenumbernode)
-    args = modeldef[:args]
+    args = transform_args(modeldef[:args])
     kwargs = modeldef[:kwargs]
 
     ## Build the anonymous evaluator from the user-provided model definition.
@@ -645,6 +677,8 @@ function build_output(modeldef, linenumbernode)
     args_nt = namedtuple_from_splitargs(args_split)
     kwargs_inclusion = map(splitarg_to_expr, kwargs_split)
 
+    # Need to update `args` since we might have added `TypeWrap` to the types.
+    modeldef[:args] = args
     # Update the function body of the user-specified model.
     # We use `MacroTools.@q begin ... end` instead of regular `quote ... end` to ensure
     # that no new `LineNumberNode`s are added apart from the reference `linenumbernode`
@@ -689,9 +723,13 @@ function matchingvalue(sampler, vi, value)
         return value
     end
 end
+# If we hit `Type` or `TypeWrap`, we immediately jump to `get_matching_type`.
 function matchingvalue(sampler::AbstractSampler, vi, value::FloatOrArrayType)
     return get_matching_type(sampler, vi, value)
 end
+function matchingvalue(sampler::AbstractSampler, vi, value::TypeWrap{T}) where {T}
+    return TypeWrap{get_matching_type(sampler, vi, T)}()
+end
 
 function matchingvalue(context::AbstractContext, vi, value)
     return matchingvalue(NodeTrait(matchingvalue, context), context, vi, value)
@@ -707,7 +745,7 @@ function matchingvalue(context::SamplingContext, vi, value)
 end
 
 """
-    get_matching_type(spl::AbstractSampler, vi, ::Type{T}) where {T}
+    get_matching_type(spl::AbstractSampler, vi, ::TypeWrap{T}) where {T}
 
 Get the specialized version of type `T` for sampler `spl`.
 

test/model.jl

@@ -350,9 +350,7 @@ is_typed_varinfo(varinfo::DynamicPPL.SimpleVarInfo{<:NamedTuple}) = true
 
     @testset "Type stability of models" begin
         models_to_test = [
-            # FIXME: Fix issues with type-stability in `DEMO_MODELS`.
-            # DynamicPPL.TestUtils.DEMO_MODELS...,
-            DynamicPPL.TestUtils.demo_lkjchol(2),
+            DynamicPPL.TestUtils.DEMO_MODELS..., DynamicPPL.TestUtils.demo_lkjchol(2)
         ]
         @testset "$(model.f)" for model in models_to_test
             vns = DynamicPPL.TestUtils.varnames(model)