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

Author: devmotion

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"

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
 

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