Test other tangents (#176)

Author: mzgubic

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.7.9"
+version = "0.7.10"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/ChainRulesTestUtils.jl

@@ -13,6 +13,7 @@ import FiniteDifferences: rand_tangent
 
 const _fdm = central_fdm(5, 1; max_range=1e-2)
 const TEST_INFERRED = Ref(true)
+const TRANSFORMS_TO_ALT_TANGENTS = Function[] # e.g. [x -> @thunk(x), _ -> ZeroTangent(), x -> rebasis(x)]
 
 export TestIterator
 export test_approx, test_scalar, test_frule, test_rrule, generate_well_conditioned_matrix

src/check_result.jl

@@ -33,6 +33,7 @@ end
 test_approx(::AbstractZero, x, msg=""; kwargs...) = test_approx(zero(x), x, msg; kwargs...)
 test_approx(x, ::AbstractZero, msg=""; kwargs...) = test_approx(x, zero(x), msg; kwargs...)
 test_approx(x::ZeroTangent, y::ZeroTangent, msg=""; kwargs...) = @test true
+test_approx(x::NoTangent, y::NoTangent, msg=""; kwargs...) = @test true
 
 # remove once https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/113
 test_approx(x::NoTangent, y::Nothing, msg=""; kwargs...) = @test true

src/testers.jl

@@ -80,6 +80,12 @@ end
 # Keyword Arguments
    - `output_tangent` tangent to test accumulation of derivatives against
      should be a differential for the output of `f`. Is set automatically if not provided.
+   - `tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS`: a vector of functions that
+      transform the passed argument tangents into alternative tangents that should be tested.
+      Note that the alternative tangents are only tested for not erroring when passed to
+      frule. Testing for correctness using finite differencing can be done using a
+      separate `test_frule` call, e.g. for testing a `ZeroTangent()` for correctness:
+      `test_frule(f, x ⊢ ZeroTangent(); tangent_transforms=[])`.
    - `fdm::FiniteDifferenceMethod`: the finite differencing method to use.
    - `frule_f=frule`: Function with an `frule`-like API that is tested (defaults to
      `frule`). Used for testing gradients from AD systems.
@@ -98,6 +104,7 @@ function test_frule(
     f,
     args...;
     output_tangent=Auto(),
+    tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS,
     fdm=_fdm,
     frule_f=ChainRulesCore.frule,
     check_inferred::Bool=true,
@@ -122,7 +129,7 @@ function test_frule(
             _test_inferred(frule_f, deepcopy(config), deepcopy(tangents), deepcopy(primals)...; deepcopy(fkwargs)...)
         end
 
-        res = frule_f(deepcopy(config), deepcopy(tangents), deepcopy(primals)...; deepcopy(fkwargs)...)
+        res = call_on_copy(frule_f, config, tangents, primals...)
         res === nothing && throw(MethodError(frule_f, typeof(primals)))
         @test_msg "The frule should return (y, ∂y), not $res." res isa Tuple{Any,Any}
         Ω_ad, dΩ_ad = res
@@ -144,10 +151,26 @@ function test_frule(
         test_approx(dΩ_ad, dΩ_fd; isapprox_kwargs...)
 
         acc = output_tangent isa Auto ? rand_tangent(Ω) : output_tangent
-        _test_add!!_behaviour(acc, dΩ_ad; rtol=rtol, atol=atol, kwargs...)
+        _test_add!!_behaviour(acc, dΩ_ad; isapprox_kwargs...)
+
+        # test that rules work for other tangents
+        _test_frule_alt_tangents(
+            call_on_copy, frule_f, config, tangent_transforms, tangents, primals, acc;
+            isapprox_kwargs...
+        )
     end  # top-level testset
 end
 
+function _test_frule_alt_tangents(
+    call, frule_f, config, tangent_transforms, tangents, primals, acc;
+    isapprox_kwargs...
+)
+    @testset "ȧrgs = $(tsf.(tangents))" for tsf in tangent_transforms
+        _, dΩ = call(frule_f, config, tsf.(tangents), primals...)
+        _test_add!!_behaviour(acc, dΩ; isapprox_kwargs...)
+    end
+end
+
 """
     test_rrule([config::RuleConfig,] f, args...; kwargs...)
 
@@ -162,6 +185,12 @@ end
 # Keyword Arguments
  - `output_tangent` the seed to propagate backward for testing (technically a cotangent).
    should be a differential for the output of `f`. Is set automatically if not provided.
+- `tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS`: a vector of functions that
+   transform the passed `output_tangent` into alternative tangents that should be tested.
+   Note that the alternative tangents are only tested for not erroring when passed to
+   rrule. Testing for correctness using finite differencing can be done using a
+   separate `test_rrule` call, e.g. for testing a `ZeroTangent()` for correctness:
+   `test_rrule(f, args...; output_tangent=ZeroTangent(), tangent_transforms=[])`.
  - `fdm::FiniteDifferenceMethod`: the finite differencing method to use.
  - `rrule_f=rrule`: Function with an `rrule`-like API that is tested (defaults to `rrule`).
    Used for testing gradients from AD systems.
@@ -180,6 +209,7 @@ function test_rrule(
     f,
     args...;
     output_tangent=Auto(),
+    tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS,
     fdm=_fdm,
     rrule_f=ChainRulesCore.rrule,
     check_inferred::Bool=true,
@@ -249,9 +279,24 @@ function test_rrule(
                 _test_add!!_behaviour(accum_cotangent, ad_cotangent; isapprox_kwargs...)
             end
         end
+
+        # test other tangents don't error when passed to the pullback
+        _test_rrule_alt_tangents(pullback, tangent_transforms, ȳ, accum_cotangents)
     end  # top-level testset
 end
 
+function _test_rrule_alt_tangents(
+    pullback, tangent_transforms, ȳ, accum_cotangents;
+    isapprox_kwargs...
+)
+    @testset "ȳ = $(tsf(ȳ))" for tsf in tangent_transforms
+        ad_cotangents = pullback(tsf(ȳ))
+        for (accum_cotangent, ad_cotangent) in zip(accum_cotangents, ad_cotangents)
+            _test_add!!_behaviour(accum_cotangent, ad_cotangent; isapprox_kwargs...)
+        end
+    end
+end
+
 """
     @maybe_inferred [Type] f(...)
 

test/testers.jl

@@ -570,6 +570,46 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
         end
     end
 
+    @testset "tangent_transforms frule" begin
+         others_work(x) = 2x
+         function ChainRulesCore.frule((Δd, Δx), ::typeof(others_work), x)
+             return others_work(x), 2Δx
+         end
+
+         others_nowork(x) = 2x
+         function ChainRulesCore.frule((Δd, Δx), ::typeof(others_nowork), x)
+             return others_nowork(x), error("nope")
+         end
+
+         test_frule(others_work, rand(); tangent_transforms=[identity, x -> @thunk(x)])
+         @test errors("nope") do
+             test_frule(others_nowork, 2.3; tangent_transforms=[x -> @thunk(x)])
+         end
+     end
+
+     @testset "tangent_transforms rrule" begin
+         others_work(x) = 2x
+         function ChainRulesCore.rrule(::typeof(others_work), x)
+             y = others_work(x)
+             others_work_pullback(ȳ) = return (NoTangent(), 2ȳ)
+             return y, others_work_pullback
+         end
+
+         others_nowork(x) = [x, x]
+         function ChainRulesCore.rrule(::typeof(others_nowork), x)
+             y = others_nowork(x)
+             others_nowork_pullback(ȳ) = return (NoTangent(), error("nope"))
+             return y, others_nowork_pullback
+         end
+
+         test_rrule(others_work, 2.3; tangent_transforms=[_ -> ZeroTangent()])
+         test_rrule(others_work, 2.3; tangent_transforms=[x -> @thunk(x)])
+
+         @test errors("nope") do
+             test_rrule(others_nowork, 2.3; tangent_transforms=[x -> @thunk(x)])
+         end
+     end
+
     @testset "Tuple primal that is not equal to differential backing" begin
         # https://github.com/JuliaMath/SpecialFunctions.jl/issues/288
         forwards_trouble(x) = (1, 2.0 * x)