add RuleConfig

Author: mcabbott

src/runtime.jl

@@ -1,4 +1,5 @@
 using ChainRulesCore
+struct DiffractorRuleConfig <: RuleConfig{Union{HasReverseMode,HasForwardsMode}} end
 
 @Base.constprop :aggressive accum(a, b) = a + b
 @Base.constprop :aggressive accum(a::Tuple, b::Tuple) = map(accum, a, b)

src/stage1/forward.jl

@@ -13,7 +13,7 @@ first_partial(x::CompositeBundle) = map(first_partial, getfield(x, :tup))
 
 # TODO: Which version do we want in ChainRules?
 function my_frule(args::ATB{1}...)
-    frule(map(first_partial, args), map(primal, args)...)
+    frule(DiffractorRuleConfig(), map(first_partial, args), map(primal, args)...)
 end
 
 # Fast path for some hot cases
@@ -118,6 +118,8 @@ function (::∂☆internal{1})(args::AbstractTangentBundle{1}...)
     end
 end
 
+ChainRulesCore.frule_via_ad(::DiffractorRuleConfig, args...) = ∂☆internal{1}()(args...)
+
 function (::∂☆shuffle{N})(args::AbstractTangentBundle{N}...) where {N}
     ∂☆p = ∂☆{minus1(N)}()
     ∂☆p(ZeroBundle{minus1(N)}(my_frule), map(shuffle_down, args)...)

src/stage1/generated.jl

@@ -210,7 +210,7 @@ function (::∂⃖{N})(f::T, args...) where {T, N}
     if N == 1
         # Base case (inlined to avoid ambiguities with manually specified
         # higher order rules)
-        z = rrule(f, args...)
+        z = rrule(DiffractorRuleConfig(), f, args...)
         if z === nothing
             return ∂⃖recurse{1}()(f, args...)
         end
@@ -226,6 +226,10 @@ function (::∂⃖{N})(f::T, args...) where {T, N}
     end
 end
 
+function ChainRulesCore.rrule_via_ad(::DiffractorRuleConfig, f::T, args...) where {T}
+    ∂⃖{1}()(f, args...)
+end
+
 @Base.pure function (::∂⃖{1})(::typeof(Core.apply_type), head, args...)
     return rrule(Core.apply_type, head, args...)
 end