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

function Add:__init(inputSize,scalar)
   parent.__init(self)
  
   local size = inputSize
   if scalar then size=1 end
   self.scalar = scalar
   self.bias = torch.Tensor(size)
   self.gradBias = torch.Tensor(size)
   
   self._ones = torch.Tensor{1}

   self:reset()
end

function Add:reset(stdv)
   if stdv then 
      stdv = stdv * math.sqrt(3)
   else
      stdv = 1./math.sqrt(self.bias:size(1))
   end

   self.bias:uniform(-stdv, stdv)
end

function Add:updateOutput(input)
   self.output:resizeAs(input):copy(input)
   if self.scalar then
      self.output:add(self.bias[1]);
   else
      if input:isSameSizeAs(self.bias) then
         self.output:add(self.bias)
      else
         local batchSize = input:size(1)
         if self._ones:size(1) ~= batchSize then
            self._ones:resize(batchSize):fill(1)
         end
         local bias = self.bias:view(-1)
         local output = self.output:view(batchSize, -1)
         output:addr(1, self._ones, bias)
      end
   end
   return self.output
end

function Add:updateGradInput(input, gradOutput)
   if self.gradInput then
      self.gradInput:resizeAs(gradOutput):copy(gradOutput) 
      return self.gradInput
   end
end

function Add:accGradParameters(input, gradOutput, scale)
   scale = scale or 1
   if self.gradBias:size(1) == 1 then
      self.gradBias[1] = self.gradBias[1] + scale*gradOutput:sum();
   else
      if input:isSameSizeAs(self.bias) then
         self.gradBias:add(scale, gradOutput)
      else
         local gradOutput = gradOutput:view(input:size(1), -1)
         self.gradBias:view(-1):addmv(scale, gradOutput:t(), self._ones)
      end
   end
end