fixup

Author: mcabbott

src/deprecations.jl

@@ -36,3 +36,9 @@ zeros32(::Type, dims...) = throw(ArgumentError("Flux.zeros32 is always Float32,
 
 
 # v0.13 deprecations
+function Broadcast.broadcasted(f::Recur, args...)
+  # This had an explicit @adjoint rule, calling Zygote.∇map(__context__, f, args...), until v0.12
+  Base.depwarn("""Broadcasting is not safe to use with RNNs, as it does not guarantee an iteration order.
+    Re-writing this as a comprehension would be better.""", :broadcasted)
+  map(f, args...)  # map isn't really safe either, but 
+end

src/layers/recurrent.jl

@@ -435,8 +435,3 @@ julia> g(rand(Float32, 3, 10)) |> size # batch size of 10
 """
 GRUv3(a...; ka...) = Recur(GRUv3Cell(a...; ka...))
 Recur(m::GRUv3Cell) = Recur(m, m.state0)
-
-# TODO move to ChainRulesCore? 
-@adjoint function Broadcast.broadcasted(f::Recur, args...)
-  Zygote.∇map(__context__, f, args...)
-end

src/losses/ctc.jl

@@ -134,8 +134,9 @@ for mathematical details.
 ctc_loss(ŷ::AbstractArray, y) = ctc_alpha(ŷ, y).loss
 
 function ChainRulesCore.rrule(::typeof(ctc_loss), ŷ, y)
-  ctc_loss_pullback(Δ) = (NoTangent(), Δ .* ∇ctc_loss(ŷ, y, out), NoTangent())
-  return ctc_loss(ŷ, y), ctc_loss_pullback
+  tmp = ctc_alpha(ŷ, y)
+  ctc_loss_pullback(Δ) = (NoTangent(), Δ .* ∇ctc_loss(ŷ, y, tmp), NoTangent())
+  return tmp.loss, ctc_loss_pullback
 end
 
 

src/losses/utils.jl

@@ -23,7 +23,7 @@ end
   res, Δ -> (nothing, Zygote.unbroadcast(x, xlogy.(Δ, y)), Zygote.unbroadcast(y, Δ .* x ./ y))
 end
 
-ChainRulesCore.@scalar_rule xlogy(x, y) (log(y), x/y)  # is this good enough?
+ChainRulesCore.@scalar_rule xlogy(x, y) (log(y), x/y)  # should help Diffractor's broadcasting
 ChainRulesCore.@scalar_rule xlogx(x) (log(y) + true)
 
 # This can be made an error in Flux v0.13, for now just a warning

src/utils.jl

@@ -793,8 +793,10 @@ true
 """
 modules(m) = [x for x in Functors.fcollect(m) if !isleaflike(x)]
 
-@nograd modules
-ChainRulesCore.@non_differentiable modules(::Any)  # is this correct?
+@nograd modules # TODO: is this correct? might fail with explicit parameters.
+function ChainRulesCore.rrule(::typeof(modules), m)
+  modules(m), dm -> error("Flux.modules is not at present differentiable, sorry")
+end
 
 isleaflike(x) = Functors.isleaf(x)
 isleaflike(::Tuple{Vararg{<:Number}}) = true