Error with Zygote and return statement

[cossio]

I am just copying from https://github.com/FluxML/Zygote.jl/issues/668. Not sure if this should be fixed here or in Zygote.

julia> using Zygote, ProgressMeter
julia> function train()
       @showprogress for d in 1:10
       gs = Zygote.gradient(1.0) do x
       return sin(x)
       end
       end
       end
train (generic function with 1 method)
julia> train()
ERROR: Compiling Tuple{typeof(lock),ProgressMeter.var"#13#14"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Progress},ReentrantLock}: try/catch is not supported.

Stacktrace

Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] instrument(::IRTools.Inner.IR) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/reverse.jl:89
 [3] #Primal#16 at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/reverse.jl:170 [inlined]
 [4] Zygote.Adjoint(::IRTools.Inner.IR; varargs::Nothing, normalise::Bool) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/reverse.jl:283
 [5] _lookup_grad(::Type{T} where T) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/emit.jl:101
 [6] #s3320#1912 at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:21 [inlined]
 [7] #s3320#1912(::Any, ::Any, ::Any) at ./none:0
 [8] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:526
 [9] #next!#12 at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:330 [inlined]
 [10] (::typeof(∂(#next!#12)))(::Nothing) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:0
 [11] next! at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:330 [inlined]
 [12] (::typeof(∂(next!)))(::Nothing) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:0
 [13] #finish!#33 at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:424 [inlined]
 [14] (::typeof(∂(#finish!#33)))(::Nothing) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:0
 [15] finish! at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:423 [inlined]
 [16] (::typeof(∂(finish!)))(::Nothing) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:0
 [17] macro expansion at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:525 [inlined]
 [18] #9 at ./REPL[8]:4 [inlined]
 [19] (::typeof(∂(λ)))(::Float64) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface2.jl:0
 [20] (::Zygote.var"#36#37"{typeof(∂(λ))})(::Float64) at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface.jl:46
 [21] gradient at /home/cossio/.julia/packages/Zygote/YeCEW/src/compiler/interface.jl:55 [inlined]
 [22] macro expansion at ./REPL[8]:3 [inlined]
 [23] macro expansion at /home/cossio/.julia/packages/ProgressMeter/lMMY9/src/ProgressMeter.jl:745 [inlined]
 [24] train() at ./REPL[8]:2
 [25] top-level scope at REPL[9]:1

If I remove the return keyword,

function train()
       @showprogress for d in 1:10
       gs = Zygote.gradient(1.0) do x
         sin(x)
       end
   end
end

there is no error.

@darsnack pointed out this line: https://github.com/timholy/ProgressMeter.jl/blob/1aba22127a903f60a034fbb2fecb8ca6f0b499b0/src/ProgressMeter.jl#L522

[darsnack]

Could we just ignore the do keyword? Anytime a do-block is used, the user could have provided a defined function instead of an anonymous one. In the case of a defined function, we wouldn't descend into that function's body to detect early returns, so we shouldn't for do-blocks either.

[KristofferC]

Zygote says:

"Source-to-source" means that Zygote hooks into Julia's compiler, and generates the backwards pass for you – as if you had written it by hand.

Without compromising on performance, Zygote supports the full flexibility and dynamism of the Julia language, including control flow, recursion, closures, structs, dictionaries, and more.

So it seems this should be opened as a Zygote issue?

[darsnack]

True, Zygote should support the try/catch handling, but I suggested opening an issue here because the source for ProgressMeter.jl specifically mentions not descending inner function returns.