Merge pull request #25 from FluxML/bc/cr-zeroes

Pass through ChainRules zeros in @adjoint too
FluxML · Mar 21, 2023 · 19414a4 · 19414a4
2 parents da51da7 + 369160f
commit 19414a4
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 23 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+Manifest.toml
diff --git a/Manifest.toml b/Manifest.toml
diff --git a/Project.toml b/Project.toml
@@ -3,6 +3,7 @@ uuid = "700de1a5-db45-46bc-99cf-38207098b444"
 version = "0.2.2"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]

diff --git a/src/adjoint.jl b/src/adjoint.jl
@@ -1,5 +1,6 @@
 using MacroTools
 using MacroTools: @q, combinedef
+using ChainRulesCore: AbstractZero
 
 function named(arg)
   if isexpr(arg, :(::)) && length(arg.args) == 1
@@ -63,13 +64,19 @@ function gradm(ex, mut = false, keepthunks = false)
     $adj
     @inline function ZygoteRules._pullback($cx, $f::$T, $(args...)) where $(Ts...)
       y, _back = adjoint(__context__, $f, $(argnames...))
-      $(mut ? nothing : :(back(::Nothing) = nothing))
+      $(mut ? nothing : quote
+        back(::Nothing) = nothing
+        back(Δ::AbstractZero) = $gradtuple(ntuple(_ -> Δ, $(length(args))))
+      end)
       back(Δ) = $gradtuple(_back($maybe_unthunked_Δ))
       return y, back
     end
     @inline function ZygoteRules._pullback($cx, ::$kT, kw, $f::$T, $(args...)) where $(Ts...)
       y, _back = adjoint(__context__, $f, $(argnames...); kw...)
-      $(mut ? nothing : :(back(::Nothing) = nothing))
+      $(mut ? nothing : quote
+        back(::Nothing) = nothing
+        back(Δ::AbstractZero) = $gradtuplekw(ntuple(_ -> Δ, $(length(args))))
+      end)
       back(Δ) = $gradtuplekw(_back($maybe_unthunked_Δ))
       return y, back
     end