add support for n-dimensional input to normalise layer

src/layers/stateless.jl

@@ -40,12 +40,12 @@ but it is more numerically stable.
 logitbinarycrossentropy(logŷ, y) = (1 - y)*logŷ - logσ(logŷ)
 
 """
-    normalise(x::AbstractVecOrMat)
+    normalise(x::AbstractArray, dims::Int=1)
 
-Normalise each column of `x` to mean 0 and standard deviation 1.
+    Normalises x to mean 0 and standard deviation 1, across the dimensions given by dims. Defaults to normalising over columns.
 """
-function normalise(x::AbstractVecOrMat)
-  μ′ = mean(x, dims = 1)
-  σ′ = std(x, dims = 1, mean = μ′, corrected=false)
+function normalise(x::AbstractArray, dims::Int=1)
+  μ′ = mean(x, dims = dims)
+  σ′ = std(x, dims = dims, mean = μ′)
   return (x .- μ′) ./ σ′
 end