FluxML / Zygote.jl

21st century AD
https://fluxml.ai/Zygote.jl/
Other
1.48k stars 213 forks source link

Support try/catch on the happy (nothrow) path #1474

Closed Pangoraw closed 1 month ago

Pangoraw commented 11 months ago

This enables differentation support to try/catch when no exception is thrown on the try block.

The changes are twofolds:

Here is an example of the behavior:

julia> using Zygote

julia> function foo(f)
         y = 1
         try
           y = f()
         catch
           y
         end
         y
       end
foo (generic function with 1 method)

julia> gradient(x -> foo(() -> x), 1) # 1
(1.0,)

julia> _, pull = pullback(foo, () -> 0//0)
(1, Zygote.var"#76#77"{Zygote.Pullback{Tuple{typeof(foo), var"#9#10"}, Any}}(∂(foo)))

julia> pull(1.)
ERROR: Can't differentiate function execution in catch block at #= REPL[4]:4 =#.
Stacktrace:
...

https://github.com/FluxML/IRTools.jl/pull/117 is a companion change that fixes some of the error that can currently happen with try/catch blocks in IRTools.jl.

Supersedes https://github.com/FluxML/Zygote.jl/pull/466.

ToucheSir commented 11 months ago

I'm not really sure how to evaluate any risks from this change other than reading through the old PR. @oxinabox as the reviewer on that, would you be able to take a look here?

Pangoraw commented 10 months ago

Tests are failing because an update to IRTools is needed (https://github.com/FluxML/IRTools.jl/pull/120).

ToucheSir commented 1 month ago

@oxinabox is this still good to merge?

oxinabox commented 1 month ago

I see no reason why it wouldn't be. I thought it was merged back in January