From 0e325e0425606161ded20064ba0c5d929f497fad Mon Sep 17 00:00:00 2001
From: Mike J Innes <mike.j.innes@gmail.com>
Date: Thu, 27 Jul 2017 20:44:45 +0100
Subject: [PATCH] Revert "Merge pull request #51 from ylxdzsw/tf-train"

This reverts commit b4a0841e9df6da38b19abcbdc8b5c1105a329dd2.
---
 src/backend/tensorflow/model.jl      | 52 +++++++++++++---------------
 src/backend/tensorflow/tensorflow.jl |  2 +-
 2 files changed, 26 insertions(+), 28 deletions(-)

diff --git a/src/backend/tensorflow/model.jl b/src/backend/tensorflow/model.jl
index cdb7f5e5..836861a0 100644
--- a/src/backend/tensorflow/model.jl
+++ b/src/backend/tensorflow/model.jl
@@ -4,22 +4,18 @@ struct Exec
   session ::Session
   input   ::Any
   output  ::Any
-  params  ::Dict{Param,Param{Tensor}}
+  grads   ::Any
+  params  ::Dict{Flux.Param,Tensor}
   stacks  ::Dict{Any,Any}
 end
 
 function makesession(model, inputs; session = Session(Graph()))
   inputs = mapt(_ -> placeholder(Float32), inputs)
   params, stacks, output = tograph(model, inputs...)
-  output = mapt(x->Param{Tensor}(x, placeholder(Float32)), output)
-  params = Dict(x=>Param{Tensor}(y, gradients(mapt(x->x.x, output),
-                                              y, mapt(x->x.Δx, output)))
-                for (x, y) in params)
-  inputs = mapt(x->Param{Tensor}(x, gradients(mapt(x->x.x, output),
-                                              x, mapt(x->x.Δx, output))),
-                inputs)
+  # grads = gradients(output, [collectt(inputs)..., values(params)...])
+  grads = placeholder(Float32)
   run(session, global_variables_initializer())
-  Exec(session, inputs, output, params, stacks)
+  Exec(session, inputs, output, grads, params, stacks)
 end
 
 retuple(xs) = xs
@@ -27,33 +23,35 @@ retuple(xs::AbstractArray{<:AbstractArray}) = (retuple.(xs)...,)
 
 dictt(xs, ys) = Dict(zip(collectt(xs), collectt(ys)))
 
-function (m::Exec)(args...)
-  dict = merge(
-    Dict(y.x=>x.x for (x, y) in m.params),
-    Dict(x.x=>y for (x, y) in zip(m.input, args))
-  )
-  retuple(run(m.session, mapt(x->x.x, m.output), dict))
+function params(m::Exec, args...)
+  shapecheckt(m.input, args)
+  idict = dictt(m.input, args)
+  pdict = Dict(t => p.x for (p, t) in m.params)
+  merge(idict, pdict)
 end
 
+function (m::Exec)(args...)
+  retuple(run(m.session, m.output, params(m, args...)))
+end
+
+pullt!(_, xs) = shift!(xs)
+pullt!(x::Tuple, xs) = map(x -> pullt!(x, xs), x)
+
+# TODO: gradients don't work yet
+# `gradients` lacks support for `grad_y`s and multiple `y`s
+
 function Flux.back!(m::Exec, Δ, args...)
-  dict = merge(
-    Dict(y.x=>x.x for (x, y) in m.params),
-    Dict(x.x=>y for (x, y) in zip(m.input, args)),
-    Dict(x.Δx=>y for (x, y) in zip(collectt(m.output), collectt(Δ)))
-  )
-
-  Δin, Δps = run(m.session, (mapt(x->x.Δx, m.input), map(x->x.Δx, values(m.params))), dict)
-
+  Δps = run(m.session, m.grads, params(m, args...))
+  Δin = pullt!(m.input, Δps)
   for (p, Δ) in zip(keys(m.params), Δps)
     p.Δx .+= Δ
   end
-
   Δin
 end
 
 function Flux.update!(m::Exec, η)
   for p in keys(m.params)
-    Flux.update!(p, η)
+    update!(p, η)
   end
   return m
 end
@@ -72,8 +70,8 @@ function (m::Model)(args...)
   @tferr m.exec.stacks m.exec(args...)
 end
 
-Flux.back!(m::Model, Δ, args...) = Flux.back!(m.exec, Δ, args...)
-Flux.update!(m::Model, η) = (Flux.update!(m.exec, η); m)
+Flux.back!(m::Model, Δ, args...) = back!(m.exec, Δ, args...)
+Flux.update!(m::Model, η) = (update!(m.exec, η); m)
 
 # Recurrent Models
 
diff --git a/src/backend/tensorflow/tensorflow.jl b/src/backend/tensorflow/tensorflow.jl
index 7536bdbc..74c94012 100644
--- a/src/backend/tensorflow/tensorflow.jl
+++ b/src/backend/tensorflow/tensorflow.jl
@@ -1,7 +1,7 @@
 module TF
 
 using ..Flux, DataFlow, TensorFlow, Juno
-import Flux: accuracy, convertel, Param
+import Flux: accuracy, convertel
 
 export tf