Merge pull request #101 from dfdx/more-dynamic-rrule-via-ad

More dynamic rrule via ad
dfdx · Jan 11, 2022 · 5583700 · 5583700
2 parents 2ad0d21 + 532308a
commit 5583700
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Yota"
 uuid = "cd998857-8626-517d-b929-70ad188a48f0"
 authors = ["Andrei Zhabinski <andrei.zhabinski@gmail.com>"]
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"

diff --git a/src/chainrules.jl b/src/chainrules.jl
@@ -135,11 +135,15 @@ function ChainRulesCore.rrule_via_ad(::YotaRuleConfig, f, args...)
     sig = call_signature(f, args...)
     if haskey(GENERATED_RRULE_CACHE, sig)
         rr = GENERATED_RRULE_CACHE[sig]
-        return Base.invokelatest(rr, f, args...)
+        # return Base.invokelatest(rr, f, args...)
+        val, pb = Base.invokelatest(rr, f, args...)
+        return val, dy -> Base.invokelatest(pb, dy)
     else
         rr = make_rrule(f, args...)
         GENERATED_RRULE_CACHE[sig] = rr
-        return Base.invokelatest(rr, f, args...)
+        # return Base.invokelatest(rr, f, args...)
+        val, pb = Base.invokelatest(rr, f, args...)
+        return val, dy -> Base.invokelatest(pb, dy)
     end
 end
 
@@ -194,7 +198,7 @@ end
 
 function ChainRulesCore.rrule(::typeof(tuple), args...)
     y = tuple(args...)
-    return y, dy -> (NoTangent(), collect(dy...)...)
+    return y, dy -> (NoTangent(), collect(dy)...)
 end
 
 # test_rrule(tuple, 1, 2, 3; output_tangent=Tangent{Tuple}((1, 2, 3)), check_inferred=false)