Improvements for LayerNorm (#1911)

Author: theabhirath

src/layers/normalise.jl

@@ -164,22 +164,21 @@ struct LayerNorm{F,D,T,N}
   affine::Bool
 end
 
-function LayerNorm(sz, λ=identity; affine=true, ϵ=1f-5)
-  sz = sz isa Integer ? (sz,) : sz
-  diag = affine ? Diagonal(sz...) : nothing
-  return LayerNorm(λ, diag, ϵ, sz, affine)
+function LayerNorm(sz, λ=identity; affine::Bool=true, ϵ::Real=1f-5)
+  diag = affine ? Diagonal(sz...) : identity
+  return LayerNorm(λ, diag, ϵ, Tuple(sz), affine)
 end
 
 @functor LayerNorm
 
 function (a::LayerNorm)(x)
-  x = normalise(x, dims=1:length(a.size), ϵ=a.ϵ)
-  a.diag === nothing ? a.λ.(x) : a.λ.(a.diag(x))
+  x = a.diag(normalise(x, dims=1:length(a.size), ϵ=a.ϵ))
+  return a.λ === identity ? x : a.λ.(x)
 end
 
 function Base.show(io::IO, l::LayerNorm)
   print(io, "LayerNorm($(l.size)")
-  l.λ == identity || print(io, ", $(l.λ)")
+  l.λ === identity || print(io, ", ", l.λ)
   hasaffine(l) || print(io, ", affine=false")
   print(io, ")")
 end

src/layers/stateless.jl

@@ -33,9 +33,8 @@ Normalise `x` to mean 0 and standard deviation 1 across the dimension(s) given b
 Per default, `dims` is the last dimension. 
 `ϵ` is a small additive factor added to the denominator for numerical stability.
 """
-function normalise(x::AbstractArray; dims=ndims(x), ϵ=ofeltype(x, 1e-5))
+@inline function normalise(x::AbstractArray; dims=ndims(x), ϵ=ofeltype(x, 1e-5))
   μ = mean(x, dims=dims)
-    #   σ = std(x, dims=dims, mean=μ, corrected=false) # use this when Zygote#478 gets merged
-  σ = std(x, dims=dims, corrected=false)
-  return (x .- μ) ./ (σ .+ ϵ)
+  σ = std(x, dims=dims, mean=μ, corrected=false)
+  return @. (x - μ) / (σ + ϵ)
 end