dummy implementation of flatten (#2205)

* reintroduce flatten * cleanup * fix docs * actually add docstring
FluxML · Mar 10, 2023 · 1258ddf · 1258ddf · CarloLucibello · Mar 10, 2023
1 parent a22a3bd
commit 1258ddf
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/docs/make.jl b/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],

diff --git a/src/deprecations.jl b/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: 

diff --git a/src/layers/stateless.jl b/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)