Refactor

Author: Anton Smirnov

src/layers/normalise.jl

@@ -206,16 +206,22 @@ function _promote_to_output(
     _maybe_promote_type(T, Vel), N), Wel)
 end
 
+function _basetype(::Type{T}) where T
+  if T <: Array
+    return Array
+  elseif T <: CuArray
+    return CuArray
+  end
+  throw("Unsupported type $T")
+end
+
 # For InstanceNorm, GroupNorm, and BatchNorm.
 # Compute the statistics on the slices specified by reduce_dims.
 # reduce_dims=[1,...,N-2,N] for BatchNorm
 # reduce_dims=[1,...,N-2] for InstanceNorm and GroupNorm
-_norm_layer_forward(l, x; reduce_dims, affine_shape) =
-    _norm_layer_forward(l, x, _promote_to_output(l, x); reduce_dims, affine_shape)
-
 function _norm_layer_forward(
-  l, x::Array{T, N}, ::Type{O}; reduce_dims, affine_shape,
-) where {T, N, O}
+  l, x::AbstractArray{T, N}; reduce_dims, affine_shape,
+) where {T, N}
   if !_isactive(l) && l.track_stats # testmode with tracked stats
     stats_shape = ntuple(i -> i == N-1 ? size(x, N-1) : 1, N)
     μ = reshape(l.μ, stats_shape)
@@ -228,7 +234,8 @@ function _norm_layer_forward(
     end
   end
 
-  o::Array{O, N} = ((x .- μ) ./ sqrt.(σ² .+ l.ϵ))
+  O = _promote_to_output(l, x)
+  o::_basetype(typeof(x)){O, N} = ((x .- μ) ./ sqrt.(σ² .+ l.ϵ))
   hasaffine(l) || return l.λ.(o)
 
   γ = reshape(l.γ, affine_shape)

test/runtests.jl

@@ -8,55 +8,55 @@ using CUDA
 
 Random.seed!(0)
 
-# @testset "Utils" begin
-#   include("utils.jl")
-# end
+@testset "Utils" begin
+  include("utils.jl")
+end
 
-# @testset "Onehot" begin
-#   include("onehot.jl")
-# end
+@testset "Onehot" begin
+  include("onehot.jl")
+end
 
-# @testset "Optimise" begin
-#   include("optimise.jl")
-# end
+@testset "Optimise" begin
+  include("optimise.jl")
+end
 
-# @testset "Data" begin
-#   include("data.jl")
-# end
+@testset "Data" begin
+  include("data.jl")
+end
 
-# @testset "Losses" begin
-#   include("losses.jl")
-#   include("ctc.jl")
-#   CUDA.functional() && include("ctc-gpu.jl")
-# end
+@testset "Losses" begin
+  include("losses.jl")
+  include("ctc.jl")
+  CUDA.functional() && include("ctc-gpu.jl")
+end
 
 @testset "Layers" begin
-  # include("layers/basic.jl")
+  include("layers/basic.jl")
   include("layers/normalisation.jl")
-  # include("layers/stateless.jl")
-  # include("layers/recurrent.jl")
-  # include("layers/conv.jl")
-  # include("layers/upsample.jl")
-  # include("layers/show.jl")
-end
-
-# @testset "outputsize" begin
-#   using Flux: outputsize
-#   include("outputsize.jl")
-# end
-
-# @testset "CUDA" begin
-#   if CUDA.functional()
-#     include("cuda/runtests.jl")
-#   else
-#     @warn "CUDA unavailable, not testing GPU support"
-#   end
-# end
-
-# @static if VERSION == v"1.6"
-#   using Documenter
-#   @testset "Docs" begin
-#     DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive=true)
-#     doctest(Flux)
-#   end
-# end
+  include("layers/stateless.jl")
+  include("layers/recurrent.jl")
+  include("layers/conv.jl")
+  include("layers/upsample.jl")
+  include("layers/show.jl")
+end
+
+@testset "outputsize" begin
+  using Flux: outputsize
+  include("outputsize.jl")
+end
+
+@testset "CUDA" begin
+  if CUDA.functional()
+    include("cuda/runtests.jl")
+  else
+    @warn "CUDA unavailable, not testing GPU support"
+  end
+end
+
+@static if VERSION == v"1.6"
+  using Documenter
+  @testset "Docs" begin
+    DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive=true)
+    doctest(Flux)
+  end
+end