Merge pull request #322 from ztangent/remove-untraced-overhead

Remove overhead during untraced execution of dynamic @gen functions.

Author: ztangent

src/dynamic/dynamic.jl

@@ -45,13 +45,17 @@ end
 
 accepts_output_grad(gen_fn::DynamicDSLFunction) = gen_fn.accepts_output_grad
 
+mutable struct GFUntracedState
+    params::Dict{Symbol,Any}
+end
+
 function (gen_fn::DynamicDSLFunction)(args...)
-    (_, _, retval) = propose(gen_fn, args)
-    retval
+    state = GFUntracedState(gen_fn.params)
+    gen_fn.julia_function(state, args...)
 end
 
-function exec(gf::DynamicDSLFunction, state, args::Tuple)
-    gf.julia_function(state, args...)
+function exec(gen_fn::DynamicDSLFunction, state, args::Tuple)
+    gen_fn.julia_function(state, args...)
 end
 
 # whether there is a gradient of score with respect to each argument
@@ -76,6 +80,13 @@ function dynamic_trace_impl(expr::Expr)
     end
 end
 
+# Defaults for untraced execution
+@inline traceat(state::GFUntracedState, gen_fn::GenerativeFunction, args, key) =
+    gen_fn(args...)
+
+@inline traceat(state::GFUntracedState, dist::Distribution, args, key) =
+    random(dist, args...)
+
 ########################
 # trainable parameters #
 ########################