Cleanup BatchNorm implementation

This provides greater datatype persistence

src/layers/normalise.jl

@@ -134,21 +134,21 @@ function (BN::BatchNorm)(x)
     μ = reshape(BN.μ, affine_shape...)
     σ = reshape(BN.σ, affine_shape...)
   else
-    T = eltype(x)
+    T = eltype(data(x))
 
-    ϵ = data(convert(T, BN.ϵ))
     axes = [1:dims-2; dims] # axes to reduce along (all but channels axis)
     μ = mean(x, dims = axes)
-    σ = sqrt.(mean((x .- μ).^2, dims = axes) .+ ϵ)
+    meansub = (x .- μ)
+    σ = mean(meansub .* meansub, dims = axes)
 
     # update moving mean/std
-    mtm = data(convert(T, BN.momentum))
-    BN.μ = (1 - mtm) .* BN.μ .+ mtm .* dropdims(data(μ), dims = (axes...,))
-    BN.σ = (1 - mtm) .* BN.σ .+ mtm .* dropdims(data(σ), dims = (axes...,)) .* m ./ (m - 1)
+    mtm = convert(T, data(BN.momentum))
+    BN.μ = (1 - mtm) .* BN.μ .+ mtm .* data(reshape(μ, :))
+    BN.σ = ((1 - mtm) .* BN.σ .+ mtm .* data(reshape(σ, :)) .* m ./ (m - 1))
   end
 
-  let λ = BN.λ
-    λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ σ) .+ reshape(β, affine_shape...))
+  let λ = BN.λ, ϵ = eltype(data(σ²))(BN.ϵ)
+    λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ .+ ϵ)) .+ reshape(β, affine_shape...))
   end
 end