fix and add tests

Author: ToucheSir

src/layers/normalise.jl

@@ -2,9 +2,14 @@ istraining() = false
 
 ChainRulesCore.rrule(::typeof(istraining)) = true, _ -> (NoTangent(),)
 
-_isactive(m) = isnothing(m.active) ? istraining() : Bool(m.active)
+_isactive(m) = Bool(something(m.active, istraining()))
 
-ChainRulesCore.@non_differentiable _isactive(::Any)
+# Avoids instabilities from differentiating through getproperty(m, :active)
+function ChainRulesCore.rrule(::typeof(_isactive), m)
+  training, _ = rrule(istraining)
+  _isactive_pullback(_) = (NoTangent(), NoTangent())
+  return Bool(something(m.active, training)), _isactive_pullback
+end
 
 _dropout_shape(s, ::Colon) = size(s)
 _dropout_shape(s, dims) = tuple((i ∉ dims ? 1 : si for (i, si) ∈ enumerate(size(s)))...)
@@ -57,7 +62,7 @@ end
 
 function (pb::DropoutPullback)(dy)
   dx = pb.project(_apply_mask(dy, pb.mask))
-  return (NoTangent(), NoTangent(), dx, NoTangent())
+  return (NoTangent(), NoTangent(), dx, NoTangent(), NoTangent(), NoTangent())
 end
 
 _apply_mask(x, ::Nothing) = x

test/layers/normalisation.jl

@@ -73,10 +73,10 @@ evalwgrad(f, x...) = pullback(f, x...)[1]
       @test cpu(m).rng === only(values(rng_kwargs))
     end
   end
-  
+
   for active in (true, false)
     m = Dropout(0.5, :, active)
-    @inferred _, back = pullback(m, rand(10)) # _, DropoutPullback{Array{Float64}}
+    _, back = @inferred pullback(m, rand(10)) # _, DropoutPullback{Array{Float64}}
     @inferred back(ones(10)) # Array{Float64}
   end
 end
@@ -353,9 +353,9 @@ end
   x = rand(2)
   m = LayerNorm(2, tanh)
   @test m(x) ≈ tanh.(Flux.normalise(x, dims=1))
-  @inferred _, back = pullback(sum∘m, x)
-  @inferred back(1.0)
-
+  _, back = @inferred pullback(|>, x, m)
+  # TODO needs https://github.com/FluxML/Zygote.jl/pull/1248
+  # @inferred back(1.0)
 
   x = rand(2,3,4,5)
   @test LayerNorm((2,3))(x) ≈ Flux.normalise(x, dims=(1,2))