Calling rand with float type arguments

Replacing the call to rand! in the _dropout_mask function to account
for complex (float) data types.

Author: ShoofLLC

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)))
+  fptype = float(real(eltype(x)))
+  y = rand!(rng, similar(x, fptype, _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)