split up Flatten layer to use the flatten function

src/layers/conv.jl

@@ -392,12 +392,20 @@ Flattening layer.
 Transforms (w,h,c,b)-shaped input into (w*h*c,b)-shaped output,
 by linearizing all values for each element in the batch.
 """
-struct Flatten end
+struct Flatten{F}
+  σ::F
+  function Flatten(σ::F = identity) where {F}
+    return new{F}(σ)
+  end
+end
 
-function (f::Flatten)(x)
-  return reshape(x, :, size(x)[end])
+function (f::Flatten)(x::AbstractArray)
+  σ = f.σ
+  σ(flatten(x))
 end
 
 function Base.show(io::IO, f::Flatten)
-  print(io, "Flatten()")
+  print(io, "Flatten(")
+  f.σ == identity || print(io, f.σ)
+  print(io, ")")
 end

src/layers/stateless.jl

@@ -84,3 +84,13 @@ function normalise(x::AbstractArray; dims=1)
   σ′ = std(x, dims = dims, mean = μ′, corrected=false)
   return (x .- μ′) ./ σ′
 end
+
+"""
+    flatten(x::AbstractArray)
+
+Transforms (w,h,c,b)-shaped input into (w*h*c,b)-shaped output,
+by linearizing all values for each element in the batch.
+"""
+function flatten(x::AbstractArray)
+  return reshape(x, :, size(x)[end])
+end

test/layers/stateless.jl

@@ -1,6 +1,6 @@
 using Test
 using Flux: onehotbatch, mse, crossentropy, logitcrossentropy,
-            σ, binarycrossentropy, logitbinarycrossentropy
+            σ, binarycrossentropy, logitbinarycrossentropy, flatten
 
 const ϵ = 1e-7
 
@@ -62,3 +62,10 @@ const ϵ = 1e-7
     end
   end
 end
+
+@testset "helpers" begin
+  @testset "flatten" begin
+    x = randn(Float32, 10, 10, 3, 2)
+    @test size(flatten(x)) == (300, 2)
+  end
+end