Merge pull request #109 from TuringLang/function_signature

Define model generator exactly as specified by the user

Author: yebai

src/compiler.jl

@@ -69,7 +69,7 @@ function model(expr, warn)
     modelinfo = build_model_info(expr)
 
     # Generate main body
-    modelinfo[:main_body] = generate_mainbody(modelinfo[:main_body], modelinfo[:args], warn)
+    modelinfo[:modelbody] = generate_mainbody(modelinfo[:body], modelinfo[:modelargs], warn)
 
     return build_output(modelinfo)
 end
@@ -87,7 +87,8 @@ function build_model_info(input_expr)
     # Construct model_info dictionary
 
     # Extracting the argument symbols from the model definition
-    arg_syms = map(modeldef[:args]) do arg
+    combinedargs = vcat(modeldef[:args], modeldef[:kwargs])
+    arg_syms = map(combinedargs) do arg
         # @model demo(x)
         if (arg isa Symbol)
             arg
@@ -113,7 +114,7 @@ function build_model_info(input_expr)
         )
         args_nt = Expr(:call, :($namedtuple), nt_type, Expr(:tuple, arg_syms...))
     end
-    args = map(modeldef[:args]) do arg
+    args = map(combinedargs) do arg
         if (arg isa Symbol)
             arg
         elseif MacroTools.@capture(arg, ::Type{T_} = Tval_)
@@ -134,7 +135,7 @@ function build_model_info(input_expr)
 
     default_syms = []
     default_vals = [] 
-    foreach(modeldef[:args]) do arg
+    foreach(combinedargs) do arg
         # @model demo(::Type{T}) where {T}
         if MacroTools.@capture(arg, ::Type{T_} = Tval_)
             push!(default_syms, T)
@@ -151,15 +152,13 @@ function build_model_info(input_expr)
     end
     defaults_nt = to_namedtuple_expr(default_syms, default_vals)
 
-    model_info = Dict(
-        :name => modeldef[:name],
-        :main_body => modeldef[:body],
-        :arg_syms => arg_syms,
-        :args_nt => args_nt,
-        :defaults_nt => defaults_nt,
-        :args => args,
-        :whereparams => modeldef[:whereparams]
+    modelderiv = Dict(
+        :modelargs => args,
+        :modelargsyms => arg_syms,
+        :modelargsnt => args_nt,
+        :modeldefaultsnt => defaults_nt,
     )
+    model_info = merge(modeldef, modelderiv)
 
     return model_info
 end
@@ -319,20 +318,18 @@ Builds the output expression.
 """
 function build_output(model_info)
     # Arguments with default values
-    args = model_info[:args]
+    args = model_info[:modelargs]
     # Argument symbols without default values
-    arg_syms = model_info[:arg_syms]
+    arg_syms = model_info[:modelargsyms]
     # Arguments namedtuple
-    args_nt = model_info[:args_nt]
+    args_nt = model_info[:modelargsnt]
     # Default values of the arguments
     # Arguments namedtuple
-    defaults_nt = model_info[:defaults_nt]
-    # Where parameters
-    whereparams = model_info[:whereparams]
+    defaults_nt = model_info[:modeldefaultsnt]
     # Model generator name
     model_gen = model_info[:name]
     # Main body of the model
-    main_body = model_info[:main_body]
+    main_body = model_info[:modelbody]
 
     unwrap_data_expr = Expr(:block)
     for var in arg_syms
@@ -341,9 +338,13 @@ function build_output(model_info)
     end
 
     @gensym(evaluator, generator)
-    generator_kw_form = isempty(args) ? () : (:($generator(;$(args...)) = $generator($(arg_syms...))),)
     model_gen_constructor = :($(DynamicPPL.ModelGen){$(Tuple(arg_syms))}($generator, $defaults_nt))
 
+    # construct the user-facing model generator
+    model_info[:name] = generator
+    model_info[:body] = :(return $(DynamicPPL.Model)($evaluator, $args_nt, $model_gen_constructor))
+    generator_expr = MacroTools.combinedef(model_info)
+
     return quote
         function $evaluator(
             _rng::$(Random.AbstractRNG),
@@ -356,8 +357,7 @@ function build_output(model_info)
             $main_body
         end
 
-        $generator($(args...)) = $(DynamicPPL.Model)($evaluator, $args_nt, $model_gen_constructor)
-        $(generator_kw_form...)
+        $(generator_expr)
 
         $(Base).@__doc__ $model_gen = $model_gen_constructor
     end

test/compiler.jl

@@ -135,7 +135,22 @@ end
         end
         f1_mm = testmodel1(1., 10.)
         @test f1_mm() == (1, 10)
-        f1_mm = testmodel1(x1=1., x2=10.)
+
+        # alternatives with keyword arguments
+        testmodel1kw(; x1, x2) = testmodel1(x1, x2)
+        f1_mm = testmodel1kw(x1 = 1., x2 = 10.)
+        @test f1_mm() == (1, 10)
+
+        @model function testmodel2(; x1, x2)
+            s ~ InverseGamma(2,3)
+            m ~ Normal(0,sqrt(s))
+
+            x1 ~ Normal(m, sqrt(s))
+            x2 ~ Normal(m, sqrt(s))
+
+            return x1, x2
+        end
+        f1_mm = testmodel2(x1=1., x2=10.)
         @test f1_mm() == (1, 10)
 
         @info "Testing the compiler's ability to catch bad models..."
@@ -199,7 +214,7 @@ end
             x ~ Bernoulli(0.5)
             return x
         end
-        @test_throws UndefKeywordError testmodel()
+        @test_throws MethodError testmodel()
 
         # Test missing initialization for vector observation turned parameter
         @model testmodel(x) = begin
@@ -285,7 +300,7 @@ end
         chain = sample(gauss(x), PG(10), 10)
         chain = sample(gauss(x), SMC(), 10)
 
-        @model gauss2(x, ::Type{TV}=Vector{Float64}) where {TV} = begin
+        @model function gauss2(::Type{TV} = Vector{Float64}; x) where {TV}
             priors = TV(undef, 2)
             priors[1] ~ InverseGamma(2,3)         # s
             priors[2] ~ Normal(0, sqrt(priors[1])) # m
@@ -295,10 +310,11 @@ end
             priors
         end
 
-        chain = sample(gauss2(x), PG(10), 10)
-        chain = sample(gauss2(x=x, TV=Vector{Float64}), PG(10), 10)
-        chain = sample(gauss2(x), SMC(), 10)
-        chain = sample(gauss2(x=x, TV=Vector{Float64}), SMC(), 10)
+        chain = sample(gauss2(x = x), PG(10), 10)
+        chain = sample(gauss2(x = x), SMC(), 10)
+
+        chain = sample(gauss2(Vector{Float64}; x = x), PG(10), 10)
+        chain = sample(gauss2(Vector{Float64}; x = x), SMC(), 10)
     end
     @testset "new interface" begin
         obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]
@@ -513,7 +529,7 @@ end
         setchunksize(N)
         alg = HMC(0.01, 5)
         x = randn(1000)
-        @model vdemo1(::Type{T}=Float64) where {T} = begin
+        @model function vdemo1(::Type{T}=Float64) where {T}
             x = Vector{T}(undef, N)
             for i = 1:N
                 x[i] ~ Normal(0, sqrt(4))
@@ -522,7 +538,9 @@ end
 
         t_loop = @elapsed res = sample(vdemo1(), alg, 250)
         t_loop = @elapsed res = sample(vdemo1(Float64), alg, 250)
-        t_loop = @elapsed res = sample(vdemo1(T=Float64), alg, 250)
+
+        vdemo1kw(; T) = vdemo1(T)
+        t_loop = @elapsed res = sample(vdemo1kw(T = Float64), alg, 250)
 
         @model vdemo2(::Type{T}=Float64) where {T <: Real} = begin
             x = Vector{T}(undef, N)
@@ -531,7 +549,9 @@ end
 
         t_vec = @elapsed res = sample(vdemo2(), alg, 250)
         t_vec = @elapsed res = sample(vdemo2(Float64), alg, 250)
-        t_vec = @elapsed res = sample(vdemo2(T=Float64), alg, 250)
+
+        vdemo2kw(; T) = vdemo2(T)
+        t_vec = @elapsed res = sample(vdemo2kw(T = Float64), alg, 250)
 
         @model vdemo3(::Type{TV}=Vector{Float64}) where {TV <: AbstractVector} = begin
             x = TV(undef, N)
@@ -540,7 +560,9 @@ end
 
         sample(vdemo3(), alg, 250)
         sample(vdemo3(Vector{Float64}), alg, 250)
-        sample(vdemo3(TV=Vector{Float64}), alg, 250)
+
+        vdemo3kw(; T) = vdemo3(T)
+        sample(vdemo3kw(T = Vector{Float64}), alg, 250)
     end
     @testset "var name splitting" begin
         var_expr = :(x)