loop lifting

src/Flux.jl

@@ -9,6 +9,7 @@ include("utils.jl")
 
 include("compiler/diff.jl")
 include("compiler/code.jl")
+include("compiler/loops.jl")
 
 include("layers/dense.jl")
 include("layers/shape.jl")

src/compiler/code.jl

@@ -6,7 +6,7 @@ export @model
 
 function process_func(ex, params = [])
   @capture(shortdef(ex), (args__,) -> body_)
-  body = Flow.il(graphm(unblock(body)))
+  body = @> body MacroTools.flatten liftloops!(params) graphm Flow.il
   body = mapconst(x -> x in params ? :(self.$x) : x, body)
   return args, body
 end

src/compiler/loops.jl

@@ -0,0 +1,49 @@
+type Delay
+  name::Symbol
+end
+
+function liftloops!(ex, params)
+  e = Flow.normedges(ex)
+  hidden = intersect((b.args[1] for b in ex.args), params)
+  edges = Dict(h => gensym("edge") for h in hidden)
+  for b in ex.args
+    b.args[2] = MacroTools.postwalk(x -> get(edges, x, x), b.args[2])
+  end
+  for (h, e) in edges
+    unshift!(ex.args, :($e = $(Delay(h))($h)))
+  end
+  return ex
+end
+
+bumpinput(i::ModelInput) = isa(i.name, Integer) ? ModelInput(i.name + 1) : i
+bumpinput(x) = x
+
+bumpinputs(v::IVertex) = mapconst(bumpinput, v)
+
+function break!(model)
+  iscyclic(graph(model)) || return model
+  bumpinputs(graph(model))
+end
+
+# r = Recurrent(784, 10, 50)
+
+# break!(r)
+
+@model type Recurrent
+  Wxh; Whh; Bh
+  Wxy; Why; By
+  hidden
+
+  function (x)
+    hidden = σ( Wxh*x + Whh*hidden + Bh )
+    y = σ( Wxy*x + Why*hidden + By )
+  end
+end
+
+Recurrent(in::Integer, out::Integer, hidden::Integer; init = initn) =
+  Recurrent(init(hidden, in), init(hidden, hidden), init(hidden),
+            init(out, in), init(out, hidden), init(hidden),
+            zeros(hidden))
+
+Base.show(io::IO, r::Recurrent) =
+  print(io, "Flux.Recurrent(...)")

src/layers/dense.jl

@@ -21,13 +21,3 @@ end
 
 Sigmoid(in::Integer, out::Integer; init = randn) =
   Sigmoid(Dense(in, out, init = init))
-
-# @model type Recurrent
-#   Wxh; Whh; Bh
-#   Wxy; Why; By
-#
-#   function (x)
-#     hidden = σ( Wxh*x + Whh*hidden + Bh )
-#     y = σ( Wxy*x + Why*hidden + By )
-#   end
-# end