Fix AlphaDropout implementation and add tests

Behaviour and outputs are adapted from the PyTorch and TF implementations

Author: ToucheSir

src/layers/normalise.jl

@@ -101,17 +101,20 @@ mutable struct AlphaDropout{F}
   end
 end
 
-function (a::AlphaDropout)(x)
+function (a::AlphaDropout)(x::AbstractArray{T}) where T
   _isactive(a) || return x
-  λ = eltype(x)(1.0507009873554804934193349852946)
-  α = eltype(x)(1.6732632423543772848170429916717)
-  α1 = eltype(x)(-λ*α)
-  noise = randn(eltype(x), size(x))
-  x = @. x*(noise > (1 - a.p)) + α1 * (noise < (1 - a.p))
-  A = sqrt(a.p + a.p * (1 - a.p) * α1^2)
-  B = -A * α1 * (1 - a.p)
-  x = @. A * x + B
-  return x
+  p = a.p
+  iszero(p) && return x
+  isone(p) && return sign.(x) .* T(0)
+
+  λ = T(1.0507009873554804934193349852946)
+  α = T(1.6732632423543772848170429916717)
+  α1 = T(-λ * α)
+  A = inv(sqrt((1 - p) * (1 + p * α1^2)))
+  B = -A * α1 * p
+
+  noise = rand!(similar(x))
+  return A .* ifelse.(noise .> p, x, α1) .+ B
 end
 
 testmode!(m::AlphaDropout, mode=true) =

test/layers/normalisation.jl

@@ -57,6 +57,34 @@ evalwgrad(f, x...) = pullback(f, x...)[1]
   @test count(a->a == 0, y) == 0
 end
 
+@testset "AlphaDropout" begin
+  x = [1., 2., 3.]
+  @test x == AlphaDropout(0.1)(x)
+  @test x == evalwgrad(AlphaDropout(0), x)
+  @test zero(x) == evalwgrad(AlphaDropout(1), x)
+
+  x = randn(1000) # large enough to prevent flaky test
+  m = AlphaDropout(0.5)
+
+  y = evalwgrad(m, x)
+  # Should preserve unit mean and variance
+  @test mean(y) ≈ 0 atol=0.1
+  @test var(y) ≈ 1 atol=0.1
+
+  testmode!(m, true) # should override istraining
+  @test evalwgrad(m, x) == x
+
+  testmode!(m, false)
+  y = evalwgrad(m, x)
+  @test mean(y) ≈ 0 atol=0.1
+  @test var(y) ≈ 1 atol=0.1
+  
+  # Known good value ranges
+  # Values taken from https://github.com/pytorch/pytorch/blob/v1.10.0/test/cpp/api/modules.cpp#L1337-L1338
+  x = ones(100)
+  @test 40 < sum(evalwgrad(m, x)) < 130
+end
+
 @testset "BatchNorm" begin
   let m = BatchNorm(2), x = [1.0 3.0 5.0;
                              2.0 4.0 6.0]