Add support for rng_from_array and restrict to dropout(::CUDA.RNG, ::CuArray)

Author: darsnack

src/layers/normalise.jl

@@ -10,7 +10,7 @@ _dropout_shape(s, dims) = tuple((i ∉ dims ? 1 : si for (i, si) ∈ enumerate(s
 _dropout_kernel(y::T, p, q) where {T} = y > p ? T(1 / q) : T(0)
 
 """
-    dropout([rng = default_rng()], x, p; dims=:, active=true)
+    dropout([rng = rng_from_array(x)], x, p; dims=:, active=true)
 
 The dropout function. If `active` is `true`,
 for each input, either sets that input to `0` (with probability
@@ -36,16 +36,20 @@ function dropout(rng, x, p; dims=:, active::Bool=true)
   y = dropout_mask(rng, x, p, dims=dims)
   return x .* y
 end
-dropout(x, p; kwargs...) = dropout(Random.default_rng(), x, p; kwargs...)
-dropout(x::CuArray, p; kwargs...) = dropout(CUDA.default_rng(), x, p; kwargs...)
+dropout(x, p; kwargs...) = dropout(rng_from_array(x), x, p; kwargs...)
+dropout(x::CuArray, p; kwargs...) = dropout(rng_from_array(x), x, p; kwargs...)
 
 @adjoint function dropout(rng, x, p; dims=:, active::Bool=true)
   active || return x, Δ -> (Δ, nothing)
   y = dropout_mask(rng, x, p, dims=dims)
   return x .* y, Δ -> (nothing, Δ .* y, nothing)
 end
 
-function dropout_mask(rng::AbstractRNG, x, p; dims=:)
+dropout_mask(rng::CUDA.RNG, x::CuArray, p; kwargs...) = _dropout_mask(rng, x, p; kwargs...)
+dropout_mask(rng, x::CuArray, p; kwargs...) =
+  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)))
   y .= _dropout_kernel.(y, p, 1 - p)
   return y
@@ -71,9 +75,9 @@ mutable struct Dropout{F,D,R<:AbstractRNG}
   active::Union{Bool, Nothing}
   rng::R
 end
-Dropout(p, dims, active) = Dropout(p, dims, active, Random.default_rng())
+Dropout(p, dims, active) = Dropout(p, dims, active, rng_from_array())
 
-function Dropout(p; dims=:, rng = Random.default_rng())
+function Dropout(p; dims=:, rng = rng_from_array())
   @assert 0 ≤ p ≤ 1
   Dropout(p, dims, nothing, rng)
 end
@@ -115,8 +119,8 @@ mutable struct AlphaDropout{F,R<:AbstractRNG}
     new{typeof(p), typeof(rng)}(p, active, rng)
   end
 end
-AlphaDropout(p, active) = AlphaDropout(p, active, Random.default_rng())
-AlphaDropout(p; rng = Random.default_rng()) = AlphaDropout(p, nothing, rng)
+AlphaDropout(p, active) = AlphaDropout(p, active, rng_from_array())
+AlphaDropout(p; rng = rng_from_array()) = AlphaDropout(p, nothing, rng)
 
 @functor AlphaDropout
 

src/utils.jl

@@ -33,6 +33,26 @@ nfan(dims...) = prod(dims[1:end-2]) .* (dims[end-1], dims[end]) # In case of con
 ofeltype(x, y) = convert(float(eltype(x)), y)
 epseltype(x) = eps(float(eltype(x)))
 
+"""
+    rng_from_array()
+    rng_from_array(x)
+
+Create an instance of the RNG most appropriate for `x`.
+The current defaults are:
+- `x isa AbstractArray`
+  - Julia version is < 1.7: `Random.GLOBAL_RNG`
+  - Julia version is >= 1.7: `Random.default_rng()`
+- `x isa CuArray`: `CUDA.default_rng()`
+When `x` is unspecified, it is assumed to be a `AbstractArray`.
+"""
+if VERSION >= v"1.7"
+  rng_from_array() = Random.default_rng()
+else
+  rng_from_array() = Random.GLOBAL_RNG
+end
+rng_from_array(::AbstractArray) = rng_from_array()
+rng_from_array(::CuArray) = CUDA.default_rng()
+
 """
     glorot_uniform([rng=GLOBAL_RNG], dims...)