delay -> offset

This commit is contained in:
Mike J Innes 2016-11-07 19:44:51 +00:00
parent c3d32c395c
commit 42a7a6ebf6
2 changed files with 42 additions and 49 deletions

View File

@ -6,7 +6,7 @@ export @net
function process_func(ex, params = []) function process_func(ex, params = [])
@capture(shortdef(ex), (args__,) -> body_) @capture(shortdef(ex), (args__,) -> body_)
body = @> body MacroTools.flatten liftloops!(params) graphm DataFlow.il body = @> body MacroTools.flatten liftloops(params) graphm DataFlow.il
body = mapconst(x -> x in params ? :(self.$x) : x, body) body = mapconst(x -> x in params ? :(self.$x) : x, body)
return args, body return args, body
end end
@ -15,7 +15,7 @@ function makegraph(graph, args)
@assert length(args) == 1 @assert length(args) == 1
mapconst(graph) do x mapconst(graph) do x
x == args[1] ? inputnode(1) : x == args[1] ? inputnode(1) :
isa(x, Delay) ? :(Delay($(Expr(:quote, x.name)), self.$(x.name))) : isa(x, Offset) ? :(Offset($(Expr(:quote, x.name)), $(x.n), self.$(x.name))) :
x x
end end
end end

View File

@ -1,26 +1,19 @@
export unroll export unroll
type Delay type Offset
name::Symbol name::Symbol
n::Int
default::Nullable{Param} default::Nullable{Param}
end end
Delay(name) = Delay(name, nothing) Offset(name, n) = Offset(name, n, nothing)
function liftloops!(ex, params) function liftloops(ex, params)
ex = DataFlow.normedges(ex) ex = DataFlow.normedges(ex)
hidden = intersect((b.args[1] for b in ex.args), params) MacroTools.postwalk(ex) do ex
edges = Dict(h => gensym("edge") for h in hidden) @capture(ex, x_{n_}) || return ex
declared = Dict(h => false for h in hidden) :($(Offset(x,n))($x))
liftvar(s) = get(declared, s, false) ? s : get(edges, s, s)
for b in ex.args
b.args[2] = MacroTools.postwalk(liftvar, b.args[2])
declared[b.args[1]] = true
end end
for (h, e) in edges
unshift!(ex.args, :($e = $(Delay(h))($h)))
end
return ex
end end
function hasloops(model) function hasloops(model)
@ -39,39 +32,39 @@ function atomise(model)
end end
end end
hiddeninput(n) = vertex(Split(n), inputnode(1)) # hiddeninput(n) = vertex(Split(n), inputnode(1))
#
function unroll!(delay::IVertex, n) # function unroll!(delay::IVertex, n)
prewalk!(delay[1]) do v # prewalk!(delay[1]) do v
v === delay ? hiddeninput(n) : v # v === delay ? hiddeninput(n) : v
end # end
end # end
#
function break!(g::IVertex) # function break!(g::IVertex)
g = bumpinputs(g) # g = bumpinputs(g)
loops = [] # loops = []
defaults = [] # defaults = []
g = prewalk!(g) do v # g = prewalk!(g) do v
isa(value(v), Delay) || return v # isa(value(v), Offset) || return v
n = length(loops)+1 # n = length(loops)+1
push!(loops, unroll!(v, n)) # push!(loops, unroll!(v, n))
push!(defaults, get(value(v).default)) # push!(defaults, get(value(v).default))
hiddeninput(n) # hiddeninput(n)
end # end
cse(group(group(loops...), g)), defaults # cse(group(group(loops...), g)), defaults
end # end
#
function unrollgraph(model, n) # function unrollgraph(model, n)
graph, defaults = break!(atomise(model)) # graph, defaults = break!(atomise(model))
outputs = [spliceinputs(graph, group([constant(splitnode(inputnode(1),i)) for i = 1:length(defaults)]...), # outputs = [spliceinputs(graph, group([constant(splitnode(inputnode(1),i)) for i = 1:length(defaults)]...),
splitnode(inputnode(2), 1))] # splitnode(inputnode(2), 1))]
for i = 2:n # for i = 2:n
push!(outputs, spliceinputs(graph, outputs[end][1], splitnode(inputnode(2), i))) # push!(outputs, spliceinputs(graph, outputs[end][1], splitnode(inputnode(2), i)))
end # end
state = outputs[end][1] # state = outputs[end][1]
outputs = map(x -> x[2], outputs) # outputs = map(x -> x[2], outputs)
(@> group(state, group(outputs...)) detuple), map(x->x.x, defaults) # (@> group(state, group(outputs...)) detuple), map(x->x.x, defaults)
end # end
type Unrolled <: Model type Unrolled <: Model
model model