- 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)
+ self:putInOrder()
+
+ self:nestApply(
+ function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end,
+ self.outputModules, gradOutput
+ )
+
+ for _, node in pairs(self.node) do
+ node.gradInputSucc = {}
+ end
+
+ for k = #self.sorted, 1, -1 do
+ local nnm = self.sorted[k]
+ local node = self.node[nnm]
+ local pred, succ, gradInputSucc = node.pred, node.succ, node.gradInputSucc
+
+ if #gradInputSucc > 0 then
+ -- We update nnm:gradInput
+ local gi
+ if #gradInputSucc == 1 then
+ gi = gradInputSucc[1] -- we avoid a clone()
+ elseif #gradInputSucc > 1 then
+ for k = 1, #gradInputSucc do
+ if gi then
+ gi:add(gradInputSucc[k])