-----------------------------------------------------------------------
-
-a = nn.Linear(10, 10)
-b = nn.ReLU()
-c = nn.Linear(10, 3)
-d = nn.Linear(10, 3)
-e = nn.CMulTable()
-f = nn.Linear(3, 2)
-
---[[
-
- a -----> b ---> c ----> e ---
- \ /
- \--> d ---/
- \
- \---> f ---
-]]--
-
-g = Graph:new()
-
-g:setInput(a)
-g:setOutput({ e, f })
-g:addEdge(c, e)
-g:addEdge(a, b)
-g:addEdge(d, e)
-g:addEdge(b, c)
-g:addEdge(b, d)
-g:addEdge(d, f)
-
-g:order()
-
-g:print(graph)
+function DAG:updateGradInput(input, gradOutput)
+ self:sort()
+
+ self:applyOnModules(function(m, i, go) m:updateGradInput(i, go) end, self.outputModules, input, gradOutput)
+
+ for k = self.sorted, 1, -1 do
+ local m = sorted[k]
+ if self.succ[d] then
+ if #self.succ[d] == 1 then
+ d:updateGradInput(self.succ[d][1].gradInput)
+ elseif #self.succ[d] > 1 then
+ local sum
+ for k = 1, #self.succ[d] do
+ if sum then
+ sum:add(self.succ[d][k].gradInput)
+ else
+ sum = self.succ[d][k].gradInput:clone()
+ end
+ end
+ d:updateGradInput(sum)
+ end
+ end
+ end