make tests pass

Author: oxinabox

src/ChainRulesOverloadGeneration.jl

@@ -1,13 +1,15 @@
+module ChainRulesOverloadGeneration
+
 using ChainRulesCore
 
 export on_new_rule, refresh_rules
 
-include("precompile.jl")
 include("ruleset_loading.jl")
+include("precompile.jl")
 
 function __init__()
     # Need to refresh rules when a package is loaded
     push!(Base.package_callbacks, _package_hook)
 end
 
-end
+end # module

src/ruleset_loading.jl

@@ -55,10 +55,10 @@ Returns a list of all the methods of the currently defined rules of the given ki
 Excluding the fallback rule that returns `nothing` for every input.
 """
 function _rule_list end
-# The fallback rules are the only rules defined in ChainRulesCore & that is how we skip them
-# TODO this needs to be changed to work now it is in it's own repo
-_rule_list(rule_kind) = (m for m in methods(rule_kind) if m.module != @__MODULE__)
+_rule_list(rule_kind) = (m for m in methods(rule_kind) if !_is_fallback(rule_kind, m))
 
+"check if this is the fallback-frule/rrule that always returns `nothing`"
+_is_fallback(rule_kind, m::Method) = m.sig === Tuple{typeof(rule_kind), Any, Vararg{Any}}
 
 const LAST_REFRESH_RRULE = Ref(0)
 const LAST_REFRESH_FRULE = Ref(0)

test/demos/forwarddiffzero.jl

@@ -2,6 +2,10 @@
 module ForwardDiffZero
 using ChainRulesCore
 using ChainRulesOverloadGeneration
+# resolve conflicts while this code exists in both.
+const on_new_rule = ChainRulesOverloadGeneration.on_new_rule
+const refresh_rules = ChainRulesOverloadGeneration.refresh_rules
+
 using Test
 
 #########################################

test/demos/reversediffzero.jl

@@ -2,6 +2,10 @@
 module ReverseDiffZero
 using ChainRulesCore
 using ChainRulesOverloadGeneration
+# resolve conflicts while this code exists in both.
+const on_new_rule = ChainRulesOverloadGeneration.on_new_rule
+const refresh_rules = ChainRulesOverloadGeneration.refresh_rules
+
 using Test
 
 #########################################

test/ruleset_loading.jl

@@ -22,8 +22,8 @@
         end
 
         @testset "# Make sure nothing happens anymore once we clear the hooks" begin
-            ChainRulesCore.clear_new_rule_hooks!(frule)
-            ChainRulesCore.clear_new_rule_hooks!(rrule)
+            ChainRulesOverloadGeneration.clear_new_rule_hooks!(frule)
+            ChainRulesOverloadGeneration.clear_new_rule_hooks!(rrule)
 
             old_frule_history = copy(frule_history)
             old_rrule_history = copy(rrule_history)
@@ -34,9 +34,9 @@
             @test old_rrule_history == rrule_history
             @test old_frule_history == frule_history
         end
-
     end
 
+
     @testset "_primal_sig" begin
         _primal_sig = ChainRulesOverloadGeneration._primal_sig
         @testset "frule" begin
@@ -69,4 +69,10 @@
             )
         end
     end
+
+    @testset "_is_fallback" begin
+        _is_fallback = ChainRulesOverloadGeneration._is_fallback
+        @test _is_fallback(rrule, only(methods(rrule, (Nothing,))))
+        @test _is_fallback(frule, only(methods(frule, (Nothing,))))
+    end
 end