ZygoteRules.jl icon indicating copy to clipboard operation
ZygoteRules.jl copied to clipboard

Allow for `unthunk` to return `nothing`

Open mcabbott opened this issue 6 months ago • 2 comments

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

mcabbott avatar Apr 09 '25 03:04 mcabbott