ZygoteRules.jl
ZygoteRules.jl copied to clipboard
Allow for `unthunk` to return `nothing`
Aims to fix https://github.com/FluxML/Zygote.jl/issues/1567
In rules defined by @adjoint, there is always a second method back(::Nothing) = nothing, so that the method you write need not allow for nothing. However, the way https://github.com/FluxML/ZygoteRules.jl/pull/17 added unthunk means that if this returns nothing, it does not cause this shortcut.
Making instead a separate method back(Δ::AbstractThunk) = back(unthunk_tangent(Δ)) should avoid that. It assumes that (eventually) unthunk_tangent must give us a non-thunk.
Cc @oschulz and @pxl-th, for work on https://github.com/FluxML/Zygote.jl/pull/966