local CAddTensorTable, parent = torch.class('nn.CAddTensorTable', 'nn.Module')

function CAddTensorTable:__init()

parent.__init(self)

self.gradInput = {}

end

-- input is a table with 2 entries. input[1] is the vector to be added.

-- input[2] is the table to which we add the vector

function CAddTensorTable:updateOutput(input)

local currentOutput = {}

for i=1,#input[2] do

currentOutput[i] = currentOutput[i] or input[1].new()

currentOutput[i]:resizeAs(input[1])

currentOutput[i]:copy(input[2][i])

currentOutput[i]:add(input[1])

end

for i = #input[2]+1, #currentOutput do

currentOutput[i] = nil

end

self.output = currentOutput

return self.output

end

function CAddTensorTable:updateGradInput(input, gradOutput)

self.gradInput[1] = self.gradInput[1] or input[1].new()

self.gradInput[1]:resizeAs(input[1])

self.gradInput[1]:copy(gradOutput[1])

for i=2, #input[2] do

self.gradInput[1]:add(gradOutput[i])

end

self.gradInput[2] = self.gradInput[2] or {}

for i=1,#input[2] do

self.gradInput[2][i] = self.gradInput[2][i] or input[1].new()

self.gradInput[2][i]:resizeAs(input[1])

self.gradInput[2][i]:copy(gradOutput[i])

end

for i=#input[2]+1, #self.gradInput[2] do

self.gradInput[2][i] = nil

end

return self.gradInput

end