Merge #1432

1432: Generalize train/testmode! to all Functors r=CarloLucibello a=ToucheSir

Addresses #1044 (comment). See also https://discourse.julialang.org/t/do-i-have-to-implement-flux-testmode-for-my-own-models/52038.

### PR Checklist

- [x] Tests are added
- [ ] Entry in NEWS.md
- [ ] Final review from `@dhairyagandhi96` (for API changes).


Co-authored-by: Brian Chen <ToucheSir@users.noreply.github.com>

Author: bors[bot]

src/functor.jl

@@ -19,7 +19,7 @@ Possible values include:
 - `true` for testing
 - `:auto` or `nothing` for Flux to detect the mode automatically
 """
-testmode!(m, mode = true) = m
+testmode!(m, mode = true) = (foreach(x -> testmode!(x, mode), trainable(m)); m)
 
 """
     trainmode!(m, mode = true)

src/layers/basic.jl

@@ -39,8 +39,6 @@ applychain(fs::Tuple, x) = applychain(tail(fs), first(fs)(x))
 
 Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i]...)
 
-testmode!(m::Chain, mode = true) = (map(x -> testmode!(x, mode), m.layers); m)
-
 function Base.show(io::IO, c::Chain)
   print(io, "Chain(")
   join(io, c.layers, ", ")

test/utils.jl

@@ -272,4 +272,17 @@ end
       testdense(re(p), bt)
     end
   end
-end 
+end
+
+@testset "Train and test mode" begin
+  mutable struct DummyLayer
+    testing::Bool
+  end
+  Flux.testmode!(m::DummyLayer, testing=true) = (m.testing = testing; m)
+
+  c = Chain(DummyLayer(true))
+  testmode!(c)
+  @test c[1].testing
+  trainmode!(c)
+  @test !c[1].testing
+end