Merge pull request #1002 from FluxML/ox/clean

Make interface2.jl code around generating pullbacks via decomposition
FluxML · Jun 19, 2021 · 531da8b · 531da8b
2 parents 13647cd + 8d9fac7
commit 531da8b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/src/compiler/emit.jl b/src/compiler/emit.jl
@@ -95,7 +95,7 @@ end
 
 varargs(m::Method, n) = m.isva ? n - m.nargs + 1 : nothing
 
-function _lookup_grad(T)
+function _generate_pullback_via_decomposition(T)
   (m = meta(T)) === nothing && return
   va = varargs(m.method, length(T.parameters))
   forw, back = stacks!(Adjoint(IR(m), varargs = va, normalise = false), T)

diff --git a/src/compiler/interface2.jl b/src/compiler/interface2.jl
@@ -23,8 +23,8 @@ end
 
   hascr && return :($chain_rrule_f(ZygoteRuleConfig(ctx), f, args...))
 
-  g = try _lookup_grad(T) catch e e end
-  !(g isa Tuple) && return :(f(args...), Pullback{$T}((f,)))
+  g = try _generate_pullback_via_decomposition(T) catch e e end
+  g === nothing && return :(f(args...), Pullback{$T}((f,)))
   meta, forw, _ = g
   argnames!(meta, Symbol("#self#"), :ctx, :f, :args)
   forw = varargs!(meta, forw, 3)
@@ -37,7 +37,8 @@ end
 
 @generated function (j::Pullback{T})(Δ) where T
   ignore_sig(T) && return :nothing
-  g = try _lookup_grad(T)
+  g = try 
+    _generate_pullback_via_decomposition(T)
   catch e
     rethrow(CompileError(T,e))
   end