Merge pull request #1867 from ShoofLLC/master

Updated Dropout for more input types.

Author: ToucheSir

NEWS.md

@@ -6,6 +6,7 @@ been removed in favour of MLDatasets.jl.
 * `params` is not exported anymore since it is a common name and is also exported by Distributions.jl 
 * `flatten` is not exported anymore due to clash with Iterators.flatten.
 * Remove Juno.jl progress bar support as it is now obsolete.
+* Improved compatibility of Dropout with Int and Complex types. 
 
 ## v0.12.10
 * `Dropout`/`AlphaDropout` now supports [user-specified RNGs](https://github.com/FluxML/Flux.jl/pull/1838)

src/layers/normalise.jl

@@ -49,7 +49,8 @@ dropout_mask(rng, x::CuArray, p; kwargs...) =
   throw(ArgumentError("x isa CuArray, but rng isa $(typeof(rng)). dropout_mask only support CUDA.RNG for CuArrays."))
 dropout_mask(rng, x, p; kwargs...) = _dropout_mask(rng, x, p; kwargs...)
 function _dropout_mask(rng, x, p; dims=:)
-  y = rand!(rng, similar(x, _dropout_shape(x, dims)))
+  realfptype = float(real(eltype(x)))
+  y = rand!(rng, similar(x, realfptype, _dropout_shape(x, dims)))
   y .= _dropout_kernel.(y, p, 1 - p)
   return y
 end

test/layers/normalisation.jl

@@ -5,6 +5,11 @@ evalwgrad(f, x...) = pullback(f, x...)[1]
 
 @testset "Dropout" begin
   @testset for rng_kwargs in ((), (; rng = MersenneTwister()))
+    x = [1.0+0im,2.0+1im,3.0+3im]
+    @test x == Dropout(0.1; rng_kwargs...)(x)
+    @test x == evalwgrad(Dropout(0; rng_kwargs...), x)
+    @test zero(x) == evalwgrad(Dropout(1; rng_kwargs...), x)
+
     x = [1.,2.,3.]
     @test x == Dropout(0.1; rng_kwargs...)(x)
     @test x == evalwgrad(Dropout(0; rng_kwargs...), x)