fix: revert catch of nothing pullbacks with Zygote (#714)

JuliaDiff · Feb 1, 2025 · dccd07e · dccd07e · gdalle · Feb 1, 2025
1 parent 19803d1
commit dccd07e
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 34 deletions.
diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml
@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.38"
+version = "0.6.39"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/...ationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl b/...ationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl
@@ -13,24 +13,6 @@ using Zygote:
     withgradient,
     withjacobian
 
-struct ZygoteNothingError <: Exception
-    f
-    x
-    contexts
-end
-
-function Base.showerror(io::IO, e::ZygoteNothingError)
-    (; f, x, contexts) = e
-    sig = (typeof(x), map(typeof ∘ DI.unwrap, contexts)...)
-    return print(
-        io,
-        "Zygote failed to differentiate function `$f` with argument types `$sig` (the pullback returned `nothing`).",
-    )
-end
-
-check_nothing(::Nothing, f, x, contexts) = throw(ZygoteNothingError(f, x, contexts))
-check_nothing(::Any, f, x, contexts) = nothing
-
 DI.check_available(::AutoZygote) = true
 DI.inplace_support(::AutoZygote) = DI.InPlaceNotSupported()
 
@@ -64,7 +46,6 @@ function DI.value_and_pullback(
     tx = map(ty) do dy
         first(pb(dy))
     end
-    check_nothing(first(tx), f, x, contexts)
     return y, tx
 end
 
@@ -80,7 +61,6 @@ function DI.value_and_pullback(
     tx = map(ty) do dy
         first(pb(dy))
     end
-    check_nothing(first(tx), f, x, contexts)
     return copy(y), tx
 end
 
@@ -96,7 +76,6 @@ function DI.pullback(
     tx = map(ty) do dy
         first(pb(dy))
     end
-    check_nothing(first(tx), f, x, contexts)
     return tx
 end
 
@@ -110,15 +89,13 @@ function DI.value_and_gradient(
     f, ::DI.NoGradientPrep, ::AutoZygote, x, contexts::Vararg{DI.Context,C}
 ) where {C}
     (; val, grad) = withgradient(f, x, map(translate, contexts)...)
-    check_nothing(first(grad), f, x, contexts)
     return val, first(grad)
 end
 
 function DI.gradient(
     f, ::DI.NoGradientPrep, ::AutoZygote, x, contexts::Vararg{DI.Context,C}
 ) where {C}
     grad = gradient(f, x, map(translate, contexts)...)
-    check_nothing(first(grad), f, x, contexts)
     return first(grad)
 end
 
@@ -147,15 +124,13 @@ function DI.value_and_jacobian(
     y = f(x, map(translate, contexts)...)
     # https://github.com/FluxML/Zygote.jl/issues/1506
     jac = jacobian(f, x, map(translate, contexts)...)
-    check_nothing(first(jac), f, x, contexts)
     return y, first(jac)
 end
 
 function DI.jacobian(
     f, ::DI.NoJacobianPrep, ::AutoZygote, x, contexts::Vararg{DI.Context,C}
 ) where {C}
     jac = jacobian(f, x, map(translate, contexts)...)
-    check_nothing(first(jac), f, x, contexts)
     return first(jac)
 end
 
@@ -242,7 +217,6 @@ function DI.hessian(
 ) where {C}
     fc = DI.with_contexts(f, contexts...)
     hess = hessian(fc, x)
-    check_nothing(hess, f, x, contexts)
     return hess
 end
 

diff --git a/DifferentiationInterface/test/Back/Zygote/test.jl b/DifferentiationInterface/test/Back/Zygote/test.jl
@@ -56,10 +56,3 @@ test_differentiation(
         logging=LOGGING,
     )
 end
-
-## Errors
-
-@testset "Errors" begin
-    safe_log(x) = x > zero(x) ? log(x) : convert(typeof(x), NaN)
-    @test_throws "Zygote failed to differentiate" derivative(safe_log, AutoZygote(), 0.0)
-end