- if #self.outputModules == 1 then
- self.output = self.outputModules[1].output
- else
- self.output = { }
- for i, d in ipairs(self.outputModules) do
- self.output[i] = d.output
+function DAG:updateGradInput(input, gradOutput)
+ self:putInOrder()
+
+ self:nestedApply(
+ function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end,
+ self.outputModules, gradOutput
+ )
+
+ self:nestedApply(
+ function(nnm, i) self.node[nnm].input = i end,
+ self.inputModules, input
+ )
+
+ 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, gradInputSucc = node.pred, node.gradInputSucc
+
+ if #gradInputSucc > 0 then
+ nnm:updateGradInput(node.input, self:computeGradInput(gradInputSucc))
+ end
+
+ -- We fill the gradInputSucc of our predecessors
+ if #pred == 1 then
+ table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
+ elseif #pred > 1 then
+ if not torch.type(nnm.gradInput) == 'table' then
+ error('Should have a table gradInput since it has multiple predecessors')
+ end
+ for n = 1, #pred do
+ table.insert(self.node[node.pred[n]].gradInputSucc, nnm.gradInput[n])
+ end