Merge pull request #314 from ztangent/typed-gen-fns

Support AD-safe type checking in dynamic @gen functions

Author: ztangent

src/dsl/dsl.jl

@@ -61,6 +61,22 @@ function parse_arg(expr)
     arg
 end
 
+function resolve_grad_arg(arg, __module__)
+    # Ensure that differentiable arguments are supported by ReverseDiff
+    if !(DSL_ARG_GRAD_ANNOTATION in arg.annotations) return arg end
+    typ = Core.eval(__module__, arg.typ)
+    if typ <: Real
+        new_typ = :Real
+    elseif typ <: AbstractArray{<:Real} && IndexStyle(typ) == IndexLinear()
+        new_typ = :(AbstractArray{<:Real})
+    elseif Real <: typ || AbstractArray{<:Real} <: typ
+        new_typ = arg.typ
+    else
+        error("Type of $(arg.name)::$(arg.typ) does not support differentiation.")
+    end
+    return Argument(arg.name, new_typ, arg.annotations, arg.default)
+end
+
 include("dynamic.jl")
 include("static.jl")
 
@@ -124,27 +140,16 @@ end
 
 function parse_gen_function(ast, annotations, __module__)
     ast = MacroTools.longdef(ast)
-    if ast.head != :function
-        error("syntax error at $ast in $(ast.head)")
-    end
-    if length(ast.args) != 2
-        error("syntax error at $ast in $(ast.args)")
-    end
-    signature = ast.args[1]
-    if signature.head == :(::)
-        (call_signature, return_type) = signature.args
-    elseif signature.head == :call
-        (call_signature, return_type) = (signature, :Any)
-    else
-        error("syntax error at $(signature)")
-    end
-    body = preprocess_body(ast.args[2], __module__)
-    name = call_signature.args[1]
-    args = map(parse_arg, call_signature.args[2:end])
+    def = MacroTools.splitdef(ast)
+    name = def[:name]
+    args = map(parse_arg, def[:args])
+    body = preprocess_body(def[:body], __module__)
+    return_type = get(def, :rtype, :Any)
     static = DSL_STATIC_ANNOTATION in annotations
     if static
         make_static_gen_function(name, args, body, return_type, annotations)
     else
+        args = map(a -> resolve_grad_arg(a, __module__), args)
         make_dynamic_gen_function(name, args, body, return_type, annotations)
     end
 end

src/dsl/dynamic.jl

@@ -2,12 +2,12 @@ const DYNAMIC_DSL_TRACE = Symbol("@trace")
 
 "Convert Argument structs to ASTs."
 function arg_to_ast(arg::Argument)
-    ast = esc(arg.name)
+    ast = Expr(:(::), esc(arg.name), esc(arg.typ))
     if (arg.default != nothing)
         default = something(arg.default)
         ast = Expr(:kw, ast, esc(default))
     end
-    ast
+    return ast
 end
 
 "Escape argument defaults (if present)."
@@ -41,8 +41,8 @@ function make_dynamic_gen_function(name, args, body, return_type, annotations)
         esc(body))
     arg_types = map((arg) -> esc(arg.typ), args)
     arg_defaults = map(escape_default, args)
-    has_argument_grads = map(
-        (arg) -> (DSL_ARG_GRAD_ANNOTATION in arg.annotations), args)
+    has_argument_grads =
+        map((arg) -> (DSL_ARG_GRAD_ANNOTATION in arg.annotations), args)
     accepts_output_grad = DSL_RET_GRAD_ANNOTATION in annotations
 
     quote

test/inference/particle_filter.jl

@@ -31,7 +31,7 @@ end
     # test the hmm_forward_alg on a hand-calculated example
     prior = [0.4, 0.6]
     emission_dists = [0.1 0.9; 0.7 0.3]'
-    transition_dists = [0.5 0.5; 0.2 0.8']
+    transition_dists = [0.5 0.5; 0.2 0.8]'
     obs = [2, 1]
     expected_marg_lik = 0.
     # z = [1, 1]
@@ -114,7 +114,7 @@ end
 
     # do particle filter steps
 
-    @gen function step_proposal(prev_trace, T::Int, x::Float64)
+    @gen function step_proposal(prev_trace, T::Int, x::Int)
         @assert T > 1
         choices = get_choices(prev_trace)
         if T > 2