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)

test/turing/compiler.jl

@@ -95,8 +95,12 @@
         @test_throws ErrorException chain = sample(gauss2(; x=x), PG(10), 10)
         @test_throws ErrorException chain = sample(gauss2(; x=x), SMC(), 10)
 
-        @test_throws ErrorException chain = sample(gauss2(Vector{Float64}; x=x), PG(10), 10)
-        @test_throws ErrorException chain = sample(gauss2(Vector{Float64}; x=x), SMC(), 10)
+        @test_throws ErrorException chain = sample(
+            gauss2(DynamicPPL.TypeWrap{Vector{Float64}}; x=x), PG(10), 10
+        )
+        @test_throws ErrorException chain = sample(
+            gauss2(DynamicPPL.TypeWrap{Vector{Float64}}; x=x), SMC(), 10
+        )
     end
     @testset "new interface" begin
         obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]
@@ -310,31 +314,35 @@
         end
 
         t_loop = @elapsed res = sample(vdemo1(), alg, 250)
-        t_loop = @elapsed res = sample(vdemo1(Float64), alg, 250)
+        t_loop = @elapsed res = sample(vdemo1(DynamicPPL.TypeWrap{Float64}()), alg, 250)
 
         vdemo1kw(; T) = vdemo1(T)
-        t_loop = @elapsed res = sample(vdemo1kw(; T=Float64), alg, 250)
+        t_loop = @elapsed res = sample(
+            vdemo1kw(; T=DynamicPPL.TypeWrap{Float64}()), alg, 250
+        )
 
         @model function vdemo2(::Type{T}=Float64) where {T<:Real}
             x = Vector{T}(undef, N)
             @. x ~ Normal(0, 2)
         end
 
         t_vec = @elapsed res = sample(vdemo2(), alg, 250)
-        t_vec = @elapsed res = sample(vdemo2(Float64), alg, 250)
+        t_vec = @elapsed res = sample(vdemo2(DynamicPPL.TypeWrap{Float64}()), alg, 250)
 
         vdemo2kw(; T) = vdemo2(T)
-        t_vec = @elapsed res = sample(vdemo2kw(; T=Float64), alg, 250)
+        t_vec = @elapsed res = sample(
+            vdemo2kw(; T=DynamicPPL.TypeWrap{Float64}()), alg, 250
+        )
 
         @model function vdemo3(::Type{TV}=Vector{Float64}) where {TV<:AbstractVector}
             x = TV(undef, N)
             @. x ~ InverseGamma(2, 3)
         end
 
         sample(vdemo3(), alg, 250)
-        sample(vdemo3(Vector{Float64}), alg, 250)
+        sample(vdemo3(DynamicPPL.TypeWrap{Vector{Float64}}()), alg, 250)
 
         vdemo3kw(; T) = vdemo3(T)
-        sample(vdemo3kw(; T=Vector{Float64}), alg, 250)
+        sample(vdemo3kw(; T=DynamicPPL.TypeWrap{Vector{Float64}}()), alg, 250)
     end
 end