-- Reference: https://arxiv.org/pdf/1607.06450.pdf (Section 3)

local LayerNormalization, parent = torch.class('nn.LayerNormalization', 'nn.Sequential')
function LayerNormalization:__init(nOutput, bias, eps, affine)
   parent.__init(self)
   eps = eps or 1e-10
   affine = (affine == nil) and true or affine
   bias = bias or 0

   self:add(nn.ConcatTable()
               :add(nn.Identity())
               :add(nn.Sequential()
                       :add(nn.Mean(1, 1))
                       :add(nn.Replicate(nOutput,1,1))))
      :add(nn.CSubTable())
      :add(nn.Normalize(2, eps))
      :add(nn.MulConstant(torch.sqrt(nOutput)))

   if affine then
      local biasTransform = nn.Add(nOutput, false)
      biasTransform.bias:fill(bias)
      local gainTransform = nn.CMul(nOutput)
      gainTransform.weight:fill(1.)
      self:add(gainTransform)
      self:add(biasTransform)
   end
end