Proper support for kwargs + splatting of args and kwargs (#477)

* initial work on simplification of compiler + supporting splatting of
args and kwargs

* fixed prob_macro

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fixed doctests

* dont use splat due to backwards compat

* Update src/compiler.jl

Co-authored-by: David Widmann <[email protected]>

* added a make_evaluate_args_and_kwargs method

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* renamed build_model_info

* bumped minor mode because Turing.jl breaks

* improved docstring

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: David Widmann <[email protected]>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.22.4"
+version = "0.23.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/compiler.jl

@@ -209,20 +209,21 @@ macro model(expr, warn=false)
 end
 
 function model(mod, linenumbernode, expr, warn)
-    modelinfo = build_model_info(expr)
+    modeldef = build_model_definition(expr)
 
     # Generate main body
-    modelinfo[:body] = generate_mainbody(mod, modelinfo[:modeldef][:body], warn)
+    modeldef[:body] = generate_mainbody(mod, modeldef[:body], warn)
 
-    return build_output(modelinfo, linenumbernode)
+    return build_output(modeldef, linenumbernode)
 end
 
 """
-    build_model_info(input_expr)
+    build_model_definition(input_expr)
 
-Builds the `model_info` dictionary from the model's expression.
+Builds the `modeldef` dictionary from the model's expression, where
+`modeldef` is a dictionary compatible with `MacroTools.combinedef`.
 """
-function build_model_info(input_expr)
+function build_model_definition(input_expr)
     # Break up the model definition and extract its name, arguments, and function body
     modeldef = MacroTools.splitdef(input_expr)
 
@@ -238,66 +239,13 @@ function build_model_info(input_expr)
 
     # Shortcut if the model does not have any arguments
     if !haskey(modeldef, :args) && !haskey(modeldef, :kwargs)
-        modelinfo = Dict(
-            :allargs_exprs => [],
-            :allargs_syms => [],
-            :allargs_namedtuple => NamedTuple(),
-            :defaults_namedtuple => NamedTuple(),
-            :modeldef => modeldef,
-        )
-        return modelinfo
+        return modeldef
     end
 
     # Ensure that all arguments have a name, i.e., are of the form `name` or `name::T`
     addargnames!(modeldef[:args])
 
-    # Extract the positional and keyword arguments from the model definition.
-    allargs = vcat(modeldef[:args], modeldef[:kwargs])
-
-    # Split the argument expressions and the default values.
-    allargs_exprs_defaults = map(allargs) do arg
-        MacroTools.@match arg begin
-            (x_ = val_) => (x, val)
-            x_ => (x, NO_DEFAULT)
-        end
-    end
-
-    # Extract the expressions of the arguments, without default values.
-    allargs_exprs = first.(allargs_exprs_defaults)
-
-    # Extract the names of the arguments.
-    allargs_syms = map(allargs_exprs) do arg
-        MacroTools.@match arg begin
-            (name_::_) => name
-            x_ => x
-        end
-    end
-
-    # Build named tuple expression of the argument symbols and variables of the same name.
-    allargs_namedtuple = to_namedtuple_expr(allargs_syms)
-
-    # Extract default values of the positional and keyword arguments.
-    default_syms = []
-    default_vals = []
-    for (sym, (expr, val)) in zip(allargs_syms, allargs_exprs_defaults)
-        if val !== NO_DEFAULT
-            push!(default_syms, sym)
-            push!(default_vals, val)
-        end
-    end
-
-    # Build named tuple expression of the argument symbols with default values.
-    defaults_namedtuple = to_namedtuple_expr(default_syms)
-
-    modelinfo = Dict(
-        :allargs_exprs => allargs_exprs,
-        :allargs_syms => allargs_syms,
-        :allargs_namedtuple => allargs_namedtuple,
-        :defaults_namedtuple => defaults_namedtuple,
-        :modeldef => modeldef,
-    )
-
-    return modelinfo
+    return modeldef
 end
 
 """
@@ -561,14 +509,32 @@ hasmissing(::Type{>:Missing}) = true
 hasmissing(::Type{<:AbstractArray{TA}}) where {TA} = hasmissing(TA)
 hasmissing(::Type{Union{}}) = false # issue #368
 
+function splitarg_to_expr((arg_name, arg_type, is_splat, default))
+    return is_splat ? :($arg_name...) : arg_name
+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
+    end
+    names_expr = Expr(:tuple, map(QuoteNode, names)...)
+    vals = Expr(:tuple, map(first, splitargs)...)
+    return :(NamedTuple{$names_expr}($vals))
+end
+
+is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#")
+
 """
-    build_output(modelinfo, linenumbernode)
+    build_output(modeldef, linenumbernode)
 
 Builds the output expression.
 """
-function build_output(modelinfo, linenumbernode)
+function build_output(modeldef, linenumbernode)
+    args = modeldef[:args]
+    kwargs = modeldef[:kwargs]
+
     ## Build the anonymous evaluator from the user-provided model definition.
-    evaluatordef = deepcopy(modelinfo[:modeldef])
+    evaluatordef = deepcopy(modeldef)
 
     # Add the internal arguments to the user-specified arguments (positional + keywords).
     evaluatordef[:args] = vcat(
@@ -577,12 +543,9 @@ function build_output(modelinfo, linenumbernode)
             :(__varinfo__::$(DynamicPPL.AbstractVarInfo)),
             :(__context__::$(DynamicPPL.AbstractContext)),
         ],
-        modelinfo[:allargs_exprs],
+        args,
     )
 
-    # Delete the keyword arguments.
-    evaluatordef[:kwargs] = []
-
     # Replace the user-provided function body with the version created by DynamicPPL.
     # 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`
@@ -593,18 +556,13 @@ function build_output(modelinfo, linenumbernode)
     # See the docstrings of `replace_returns` for more info.
     evaluatordef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        $(replace_returns(make_returns_explicit!(modelinfo[:body])))
+        $(replace_returns(make_returns_explicit!(modeldef[:body])))
     end
 
     ## Build the model function.
 
-    # Extract the named tuple expression of all arguments and the default values.
-    allargs_namedtuple = modelinfo[:allargs_namedtuple]
-    defaults_namedtuple = modelinfo[:defaults_namedtuple]
-
     # Obtain or generate the name of the model to support functors:
     # https://github.com/TuringLang/DynamicPPL.jl/issues/367
-    modeldef = modelinfo[:modeldef]
     if MacroTools.@capture(modeldef[:name], ::T_)
         name = gensym(:f)
         modeldef[:name] = Expr(:(::), name, T)
@@ -613,13 +571,18 @@ function build_output(modelinfo, linenumbernode)
         throw(ArgumentError("unsupported format of model function"))
     end
 
+    args_split = map(MacroTools.splitarg, args)
+    kwargs_split = map(MacroTools.splitarg, kwargs)
+    args_nt = namedtuple_from_splitargs(args_split)
+    kwargs_inclusion = map(splitarg_to_expr, kwargs_split)
+
     # 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`
     # to the call site
     modeldef[:body] = MacroTools.@q begin
         $(linenumbernode)
-        return $(DynamicPPL.Model)($name, $allargs_namedtuple, $defaults_namedtuple)
+        return $(DynamicPPL.Model)($name, $args_nt; $(kwargs_inclusion...))
     end
 
     return MacroTools.@q begin

src/model.jl

@@ -67,13 +67,17 @@ model with different arguments.
 @generated function Model(
     f::F,
     args::NamedTuple{argnames,Targs},
-    defaults::NamedTuple=NamedTuple(),
+    defaults::NamedTuple,
     context::AbstractContext=DefaultContext(),
 ) where {F,argnames,Targs}
     missings = Tuple(name for (name, typ) in zip(argnames, Targs.types) if typ <: Missing)
     return :(Model{$missings}(f, args, defaults, context))
 end
 
+function Model(f, args::NamedTuple, context::AbstractContext=DefaultContext(); kwargs...)
+    return Model(f, args, NamedTuple(kwargs), context)
+end
+
 function contextualize(model::Model, context::AbstractContext)
     return Model(model.f, model.args, model.defaults, context)
 end
@@ -177,7 +181,7 @@ julia> @model function demo_mv(::Type{TV}=Float64) where {TV}
            m[2] ~ Normal()
            return m
        end
-demo_mv (generic function with 3 methods)
+demo_mv (generic function with 4 methods)
 
 julia> model = demo_mv();
 
@@ -376,7 +380,7 @@ julia> @model function demo_mv(::Type{TV}=Float64) where {TV}
            m[2] ~ Normal()
            return m
        end
-demo_mv (generic function with 3 methods)
+demo_mv (generic function with 4 methods)
 
 julia> model = demo_mv();
 
@@ -573,12 +577,27 @@ end
 
 Evaluate the `model` with the arguments matching the given `context` and `varinfo` object.
 """
-@generated function _evaluate!!(
-    model::Model{_F,argnames}, varinfo, context
+function _evaluate!!(model::Model, varinfo::AbstractVarInfo, context::AbstractContext)
+    args, kwargs = make_evaluate_args_and_kwargs(model, varinfo, context)
+    return model.f(args...; kwargs...)
+end
+
+"""
+    make_evaluate_args_and_kwargs(model, varinfo, context)
+
+Return the arguments and keyword arguments to be passed to the evaluator of the model, i.e. `model.f`e.
+"""
+@generated function make_evaluate_args_and_kwargs(
+    model::Model{_F,argnames}, varinfo::AbstractVarInfo, context::AbstractContext
 ) where {_F,argnames}
     unwrap_args = [
-        :($matchingvalue(context_new, varinfo, model.args.$var)) for var in argnames
+        if is_splat_symbol(var)
+            :($matchingvalue(context_new, varinfo, model.args.$var)...)
+        else
+            :($matchingvalue(context_new, varinfo, model.args.$var))
+        end for var in argnames
     ]
+
     # We want to give `context` precedence over `model.context` while also
     # preserving the leaf context of `context`. We can do this by
     # 1. Set the leaf context of `model.context` to `leafcontext(context)`.
@@ -590,7 +609,7 @@ Evaluate the `model` with the arguments matching the given `context` and `varinf
         context_new = setleafcontext(
             context, setleafcontext(model.context, leafcontext(context))
         )
-        model.f(
+        args = (
             model,
             # Maybe perform `invlink!!` once prior to evaluation to avoid
             # lazy `invlink`-ing of the parameters. This can be useful for
@@ -600,6 +619,8 @@ Evaluate the `model` with the arguments matching the given `context` and `varinf
             context_new,
             $(unwrap_args...),
         )
+        kwargs = model.defaults
+        return args, kwargs
     end
 end
 

src/prob_macro.jl

@@ -92,6 +92,8 @@ function probtype(
             return getfield(right, arg)
         elseif arg in defaultnames
             return getfield(defaults, arg)
+        elseif arg in argnames
+            return getfield(model.args, arg)
         else
             return nothing
         end
@@ -170,7 +172,7 @@ end
             push!(argvals, :(model.defaults.$argname))
         else
             push!(warnings, :(@warn($(warn_msg(argname)))))
-            push!(argvals, :(nothing))
+            push!(argvals, :(model.args.$argname))
         end
     end
 
@@ -184,7 +186,7 @@ end
     end
 end
 
-warn_msg(arg) = "Argument $arg is not defined. A value of `nothing` is used."
+warn_msg(arg) = "Argument $arg is not defined. Using the value from the model."
 
 function Distributions.loglikelihood(
     left::NamedTuple, right::NamedTuple, _model::Model, _vi::Union{Nothing,VarInfo}
@@ -227,6 +229,8 @@ end
             push!(missings, argname)
         elseif argname in defaultnames
             push!(argvals, :(model.defaults.$argname))
+        elseif argname in argnames
+            push!(argvals, :(model.args.$argname))
         else
             throw(
                 "This point should not be reached. Please open an issue in the DynamicPPL.jl repository.",

src/varname.jl

@@ -22,8 +22,10 @@ Statically check whether the variable of name `varname` is an argument of the `m
 
 Possibly existing indices of `varname` are neglected.
 """
-@generated function inargnames(::VarName{s}, ::Model{_F,argnames}) where {s,argnames,_F}
-    return s in argnames
+@generated function inargnames(
+    ::VarName{s}, ::Model{_F,argnames,defaultnames}
+) where {s,argnames,defaultnames,_F}
+    return s in argnames || s in defaultnames
 end
 
 """

test/compiler.jl

@@ -52,7 +52,7 @@ end
             x ~ Normal()
             return x
         end
-        @test length(methods(testmodel01)) == 3
+        @test length(methods(testmodel01)) == 4
         f0_mm = testmodel01()
         @test mean(f0_mm() for _ in 1:1000) ≈ 0.0 atol = 0.1
 
@@ -65,7 +65,7 @@ end
             x[2] ~ Normal()
             return x
         end
-        @test length(methods(testmodel02)) == 3
+        @test length(methods(testmodel02)) == 4
         f0_mm = testmodel02()
         @test all(x -> isapprox(x, 0; atol=0.1), mean(f0_mm() for _ in 1:1000))
 
@@ -74,7 +74,7 @@ end
             return x
         end
         f01_mm = testmodel03()
-        @test length(methods(testmodel03)) == 3
+        @test length(methods(testmodel03)) == 4
         @test mean(f01_mm() for _ in 1:1000) ≈ 0.5 atol = 0.1
 
         # test if we get the correct return values
@@ -620,4 +620,32 @@ end
         @test f_393()() == 1
         @test f_393(Val(true))() == 0
     end
+
+    @testset "splatting of args and kwargs" begin
+        @model function f_splat_test_1(x; y::T=1, kwargs...) where {T}
+            x ~ Normal(y, 1)
+            return x, y, T, NamedTuple(kwargs)
+        end
+
+        # Non-empty `kwargs...`.
+        res = f_splat_test_1(1; z=2, w=3)()
+        @test res == (1, 1, Int, (z=2, w=3))
+
+        # Empty `kwargs...`.
+        res = f_splat_test_1(1)()
+        @test res == (1, 1, Int, NamedTuple())
+
+        @model function f_splat_test_2(x, args...; y::T=1, kwargs...) where {T}
+            x ~ Normal(y, 1)
+            return x, args, y, T, NamedTuple(kwargs)
+        end
+
+        # Non-empty `args...` and non-empty `kwargs...`.
+        res = f_splat_test_2(1, 2, 3; z=2, w=3)()
+        @test res == (1, (2, 3), 1, Int, (z=2, w=3))
+
+        # Empty `args...` and empty `kwargs...`.
+        res = f_splat_test_2(1)()
+        @test res == (1, (), 1, Int, NamedTuple())
+    end
 end