Zygote.jl tries to back-propagate nothings when branching happens inside broadcasting

[tkf]

julia> relu(x) = x > 0 ? x : zero(x)
relu (generic function with 1 method)

julia> gradient(x -> relu(x * 10), -1)
(nothing,)

julia> gradient(x -> sum(relu.(x .* 10)), [-1])
ERROR: MethodError: no method matching *(::Nothing, ::Int64)
Closest candidates are:
  *(::Any, ::Any, ::Any, ::Any...) at operators.jl:529
  *(::Complex{Bool}, ::Real) at complex.jl:305
  *(::Missing, ::Number) at missing.jl:115
  ...
Stacktrace:
 [1] _broadcast_getindex_evalf at ./broadcast.jl:630 [inlined]
 [2] _broadcast_getindex at ./broadcast.jl:603 [inlined]
 [3] getindex at ./broadcast.jl:563 [inlined]
 [4] copy at ./broadcast.jl:853 [inlined]
 [5] materialize at ./broadcast.jl:819 [inlined]
 [6] (::Zygote.var"##1208#1209"{Array{Int64,1},Int64})(::Array{Nothing,1}) at /home/takafumi/.julia/packages/Zygote/bdE6T/src/lib/broadcast.jl:63
 [7] (::Zygote.var"##3102#back#1210"{Zygote.var"##1208#1209"{Array{Int64,1},Int64}})(::Array{Nothing,1}) at /home/takafumi/.julia/packages/ZygoteRules/Jn4LO/src/adjoint.jl:48
 [8] (::typeof(∂(#121)))(::Int64) at ./REPL[121]:1
 [9] (::Zygote.var"##32#33"{typeof(∂(#121))})(::Int64) at /home/takafumi/.julia/packages/Zygote/bdE6T/src/compiler/interface.jl:38
 [10] gradient(::Function, ::Array{Int64,1}, ::Vararg{Array{Int64,1},N} where N) at /home/takafumi/.julia/packages/Zygote/bdE6T/src/compiler/interface.jl:47
 [11] top-level scope at REPL[121]:1

I checked this with Zygote v0.3.4 and current master d74f3cf.

[tkf]

This still fails even with #319

using LinearAlgebra
relu(x) = x > 0 ? x : zero(x)
M = fill(1, (1, 1))
gradient(x -> sum(relu.(M * x)), [-1])

[tkf]

Maybe it makes sense to have something like this in utils.jl so that users can workaround bugs like this?

fixnothing(x) = x
@adjoint fixnothing(x) = fixnothing(x), function(y)
    if y === nothing || y isa AbstractArray{Nothing}
        return (zero(x),)
    else
        return (y,)
    end
end

Usage:

julia> gradient(x -> sum(relu.(fixnothing(M * x))), [-1])
([0],)

[MikeInnes]

I tried a slightly simpler fix for this, can you see how that goes?

[tkf]

Thanks a lot! It solves the problems I reported.