Implement ADviaFDConfig (#243)

Author: mzgubic

src/ChainRulesTestUtils.jl

@@ -33,5 +33,6 @@ include("rule_config.jl")
 include("finite_difference_calls.jl")
 include("testers.jl")
 
+include("deprecated.jl")
 include("global_checks.jl")
 end # module

src/deprecated.jl

@@ -0,0 +1 @@
+Base.@deprecate_binding ADviaRuleConfig TestConfig false

src/rule_config.jl

@@ -1,16 +1,44 @@
-# For testing this config re-dispatches Xrule_via_ad to Xrule without config argument
-struct ADviaRuleConfig <: RuleConfig{Union{HasReverseMode, HasForwardsMode}} end
+# For testing this config uses finite differences to evaluate the frule and rrule
+struct TestConfig <: RuleConfig{Union{HasReverseMode, HasForwardsMode}}
+    fdm
+end
+TestConfig() = TestConfig(central_fdm(5, 1))
+
+function ChainRulesCore.frule_via_ad(config::TestConfig, ȧrgs, f, args...; kws...)
 
-function ChainRulesCore.frule_via_ad(config::ADviaRuleConfig, ȧrgs, f, args...; kws...)
+    # try using a rule
     ret = frule(config, ȧrgs, f, args...; kws...)
-    # we don't support actually doing AD: the rule has to exist. lets give helpfulish error
-    ret === nothing && throw(MethodError(frule, (ȧrgs, f, args...)))
-    return ret
+    ret === nothing || return ret
+
+    # but if the rule doesn't exist, use finite differencing instead
+    call_on_copy(f, xs...) = deepcopy(f)(deepcopy(xs)...; deepcopy(kws)...)
+
+    primals = (f, args...)
+    is_ignored = isa.(ȧrgs, NoTangent)
+
+    Ω = call_on_copy(f, args...)
+    ΔΩ = _make_jvp_call(config.fdm, call_on_copy, Ω, primals, ȧrgs, is_ignored)
+
+    return Ω, ΔΩ
 end
 
-function ChainRulesCore.rrule_via_ad(config::ADviaRuleConfig, f, args...; kws...)
+function ChainRulesCore.rrule_via_ad(config::TestConfig, f, args...; kws...)
+
+    # try using a rule
     ret = rrule(config, f, args...; kws...)
-    # we don't support actually doing AD: the rule has to exist. lets give helpfulish error
-    ret === nothing && throw(MethodError(rrule, (f, args...)))
-    return ret
+    ret === nothing || return ret
+
+    # but if the rule doesn't exist, use finite differencing instead
+    call(f, xs...) = f(xs...; kws...)
+
+    # this block is here just to work out which tangents should be ignored
+    primals = (f, args...)
+    primals_and_tangents = auto_primal_and_tangent.(primals)
+    is_ignored = isa.(tangent.(primals_and_tangents), NoTangent)
+
+    function f_pb(ȳ)
+        return _make_j′vp_call(config.fdm, call, ȳ, primals, is_ignored)
+    end
+
+    return call(f, args...), f_pb
 end

src/testers.jl

@@ -72,7 +72,7 @@ end
     test_frule([config::RuleConfig,] f, args..; kwargs...)
 
 # Arguments
-- `config`: defaults to `ChainRulesTestUtils.ADviaRuleConfig`.
+- `config`: defaults to `ChainRulesTestUtils.TestConfig`.
 - `f`: function for which the `frule` should be tested. Its tangent can be provided using `f ⊢ ḟ`.
   (You can enter `⊢` via `\\vdash` + tab in the Julia REPL and supporting editors.)
 - `args...`: either the primal args `x`, or primals and their tangents: `x ⊢ ẋ`
@@ -92,8 +92,7 @@ end
    - All remaining keyword arguments are passed to `isapprox`.
 """
 function test_frule(args...; kwargs...)
-    config = ChainRulesTestUtils.ADviaRuleConfig()
-    test_frule(config, args...; kwargs...)
+    test_frule(TestConfig(), args...; kwargs...)
 end
 
 function test_frule(
@@ -146,7 +145,7 @@ end
     test_rrule([config::RuleConfig,] f, args...; kwargs...)
 
 # Arguments
-- `config`: defaults to `ChainRulesTestUtils.ADviaRuleConfig`.
+- `config`: defaults to `ChainRulesTestUtils.TestConfig`.
 - `f`: function for which the `rrule` should be tested. Its tangent can be provided using `f ⊢ f̄`.
   (You can enter `⊢` via `\\vdash` + tab in the Julia REPL and supporting editors.)
 - `args...`: either the primal args `x`, or primals and their tangents: `x ⊢ x̄`
@@ -168,8 +167,7 @@ end
  - All remaining keyword arguments are passed to `isapprox`.
 """
 function test_rrule(args...; kwargs...)
-    config = ChainRulesTestUtils.ADviaRuleConfig()
-    test_rrule(config, args...; kwargs...)
+    test_rrule(TestConfig(), args...; kwargs...)
 end
 
 function test_rrule(

test/rule_config.jl

@@ -49,4 +49,66 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
         errors(() -> test_rrule(has_config, rand()), "no method matching rrule")
         errors(() -> test_rrule(has_trait, rand()), "no method matching rrule")
     end
+
+    @testset "TestConfig direct" begin
+        poly(x) = x^2 + 3.2x
+
+        x = 2.1
+        config = ChainRulesTestUtils.TestConfig()
+
+        @testset "rrule" begin
+            y, pb = rrule_via_ad(config, poly, x)
+            @test y == poly(x)
+            test_approx(pb(1.0), (NoTangent(), (2*x + 3.2) * 1.0))
+            # and automatically
+            test_rrule(config, poly, rand(); rrule_f=rrule_via_ad, check_inferred=false)
+        end
+
+        @testset "frule" begin
+            ḟ, ẋ = (NoTangent(), rand())
+            Ω, ΔΩ = frule_via_ad(config, (ḟ, ẋ), poly, x)
+            @test Ω == poly(x)
+            test_approx(ΔΩ, (2*x + 3.2) * ẋ)
+            # and automatically
+            test_frule(config, poly, x; frule_f=frule_via_ad, check_inferred=false)
+        end
+
+        # more functions
+        simo(x) = (x, 2x, 3x)
+        miso(x, y, z) = x+y
+        test_rrule(config, simo, x; rrule_f=rrule_via_ad, check_inferred=false)
+        test_rrule(config, miso, x, 2x, "s"; rrule_f=rrule_via_ad, check_inferred=false)
+
+        test_frule(config, simo, x; frule_f=frule_via_ad, check_inferred=false)
+        test_frule(config, miso, x, x, "s"; frule_f=frule_via_ad, check_inferred=false)
+    end
+
+    @testset "TestConfig in a rule" begin
+        inner(x, y) = x^2 + 2*y + 3
+        outer(f, x) = 2 * f(x, 3.2)
+
+        function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(outer), f, x)
+            fx, pb_f = rrule_via_ad(config, f, x, 3.2)
+            outer_pb(ȳ) = (NoTangent(), pb_f(2 * ȳ)[1:2]...)
+            return outer(f, x), outer_pb
+        end
+
+        function ChainRulesCore.frule(config::RuleConfig{>:HasForwardsMode}, (ȯuter, ḟ, ẋ), ::typeof(outer), f, x)
+            inner, inner_dot = frule_via_ad(config, (ḟ, ẋ, NoTangent()), f, x, 3.2)
+            return 2 * inner, 2 * inner_dot
+        end
+
+        config = ChainRulesTestUtils.TestConfig()
+        test_rrule(config, outer, inner, rand(); rrule_f=rrule_via_ad, check_inferred=false)
+        test_frule(config, outer, inner, rand(); frule_f=frule_via_ad, check_inferred=false)
+    end
+
+    @testset "Catch incorrect rules" begin
+        myid(x) = x
+        function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(myid), x)
+            wrong_pb(dy) = (NoTangent(), 8dy)
+            return x, wrong_pb
+        end
+        @test fails(() -> test_rrule(myid, 3.0; rrule_f=rrule_via_ad, check_inferred=false))
+    end
 end

test/testers.jl

@@ -285,11 +285,11 @@ end
 
     @testset "single input, multiple output" begin
         simo(x) = (x, 2x)
-        function ChainRulesCore.rrule(simo, x)
+        function ChainRulesCore.rrule(::typeof(simo), x)
             simo_pullback((a, b)) = (NoTangent(), a .+ 2 .* b)
             return simo(x), simo_pullback
         end
-        function ChainRulesCore.frule((_, ẋ), simo, x)
+        function ChainRulesCore.frule((_, ẋ), ::typeof(simo), x)
             y = simo(x)
             return y, Tangent{typeof(y)}(ẋ, 2ẋ)
         end