dummy implementation of flatten (#2205)

* reintroduce flatten

* cleanup

* fix docs

* actually add docstring

Author: CarloLucibello

docs/make.jl

@@ -1,7 +1,9 @@
 using Documenter, Flux, NNlib, Functors, MLUtils, BSON, Optimisers, OneHotArrays, Zygote, ChainRulesCore, Plots, MLDatasets, Statistics, DataFrames
 
 
-DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive = true)
+DocMeta.setdocmeta!(Flux, :DocTestSetup, 
+    :(using Flux; using Flux: flatten); 
+    recursive = true)
 
 makedocs(
     modules = [Flux, NNlib, Functors, MLUtils, BSON, Optimisers, OneHotArrays, Zygote, ChainRulesCore, Base, Plots, MLDatasets, Statistics, DataFrames],

src/deprecations.jl

@@ -187,9 +187,6 @@ function update!(opt::Optimise.AbstractOptimiser, ::Params, grads::Union{Tuple,
     """)
 end
 
-# https://github.com/FluxML/Flux.jl/issues/2195
-const flatten = MLUtils.flatten
-
 # v0.14 deprecations
 
 # Enable these when 0.14 is released, and delete const ClipGrad = Optimise.ClipValue etc: 

src/layers/stateless.jl

@@ -81,3 +81,14 @@ function ChainRulesCore.rrule(::typeof(_match_eltype), layer, x::AbstractArray)
   _match_eltype(layer, x), dx -> (ZeroTangent(), NoTangent(), dx)  # does not un-thunk dx
 end
 
+# We have to define our own flatten in order 
+# to load previously saved models. 
+# See #2195 #2204
+"""
+  flatten(x)
+
+Same as [`MLUtils.flatten`](@ref), which 
+should be prefered to this method existing 
+only for backward compatibility.
+"""
+flatten(x) = MLUtils.flatten(x)