manual gradient checks for RNN - implicit and explicit gradients

[jeremiedb]

This adds gradient tests for RNN in relation to #2185

Note that for implicit gradient mode, gradients successfully pass all tests on all Julia versions.
Implicit mode gradients only fail on Julia >= 1.7 when in REPL (and gradient isnt' call from within a function).

For explicit mode gradient (new Optimisers.jl), all gradients fail on Julia >= 1.7.
On Julia v1.6, all gradients other than state0 are correct. The correct state0 gradient is actually assigned to Recur's state rather than to the cell's state0.

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable

[ToucheSir]

Did you discover anything in your test creation which might give us a lead on where the incorrect grads are coming from?

[jeremiedb]

Did you discover anything in your test creation which might give us a lead on where the incorrect grads are coming from?

Sorry no new insights at the moment. For explicit mode state0 issue, I suspect some "alias" from

Flux.jl/src/layers/recurrent.jl

Line 179 in 1ca15f3

reset!(m::Recur) = (m.state = m.cell.state0)

.

[ToucheSir]

Thanks!

[mcabbott]

Would it be easy to add explicit errors to the wrong ones? E.g. by overloading Zygote.pullback(::Context{true}, ...) where the struct has a flag to indicate implicit mode.

[ToucheSir]

Unfortunately it's the explicit mode path which is the broken one.

[mcabbott]

Sure, but we can dispatch on that too? Haven't tried & not sure whether there's a point at which this could be attached.

[ToucheSir]

We can, but we have to find that point first. And if we do, it's likely we'll be able fo fix the bug then without having to put up a "this is broken" sign.