Retrieve pullback only once for ReverseRuleConfigBackend (#51)

* Retrieve pullback only once for ReverseRuleConfigBackend * Update Project.toml * `HasReverseMode` is not exported... * Fix CRC imports * Missing bracket
JuliaDiff · Feb 9, 2022 · 031f7f4 · 031f7f4 · devmotion · Feb 9, 2022
1 parent 8f0d6db
commit 031f7f4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 11 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "AbstractDifferentiation"
 uuid = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 authors = ["Mohamed Tarek <mohamed82008@gmail.com> and contributors"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

diff --git a/src/AbstractDifferentiation.jl b/src/AbstractDifferentiation.jl
@@ -1,7 +1,7 @@
 module AbstractDifferentiation
 
 using LinearAlgebra, ExprTools, Requires, Compat
-using ChainRulesCore: RuleConfig, rrule_via_ad
+using ChainRulesCore: ChainRulesCore
 
 export AD
 

diff --git a/src/ruleconfig.jl b/src/ruleconfig.jl
@@ -3,17 +3,13 @@
 
 AD backend that uses reverse mode with any ChainRules-compatible reverse-mode AD package.
 """
-struct ReverseRuleConfigBackend{RC <: RuleConfig} <: AbstractReverseMode
+struct ReverseRuleConfigBackend{RC<:ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}} <: AbstractReverseMode
     ruleconfig::RC
 end
 
 AD.@primitive function pullback_function(ab::ReverseRuleConfigBackend, f, xs...)
-    return (vs) -> begin
-        _, back = rrule_via_ad(ab.ruleconfig, f, xs...)
-        if vs isa Tuple && length(vs) === 1
-            return Base.tail(back(vs[1]))
-        else
-            return Base.tail(back(vs))
-        end
-    end
+    _, back = ChainRulesCore.rrule_via_ad(ab.ruleconfig, f, xs...)
+    pullback(vs) = Base.tail(back(vs))
+    pullback(vs::Tuple{Any}) = Base.tail(back(first(vs)))
+    return pullback
 end