Change to just testing output tangent thunks in the test_rrule only

Author: oxinabox

src/testers.jl

@@ -80,12 +80,6 @@ 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.
@@ -104,7 +98,6 @@ function test_frule(
     f,
     args...;
     output_tangent=Auto(),
-    tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS,
     fdm=_fdm,
     frule_f=ChainRulesCore.frule,
     check_inferred::Bool=true,
@@ -143,25 +136,9 @@ function test_frule(
 
         acc = output_tangent isa Auto ? rand_tangent(Ω) : output_tangent
         _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 = $(_string_typeof(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...)
 
@@ -176,12 +153,8 @@ 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=[])`.
+- `check_thunked_output_tangent=true`: also checks that passing a thunked version of the 
+    output tangent to the pullback returns the same result.
  - `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.
@@ -200,7 +173,7 @@ function test_rrule(
     f,
     args...;
     output_tangent=Auto(),
-    tangent_transforms=TRANSFORMS_TO_ALT_TANGENTS,
+    check_thunked_output_tangent=true,
     fdm=_fdm,
     rrule_f=ChainRulesCore.rrule,
     check_inferred::Bool=true,
@@ -267,21 +240,10 @@ function test_rrule(
             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 "ȳ = $(_string_typeof(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...)
+        if check_thunked_output_tangent
+            test_approx(ad_cotangents, pullback(@thunk(ȳ)), "pulling back a thunk")
         end
-    end
+    end  # top-level testset
 end
 
 """

test/testers.jl

@@ -444,9 +444,7 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
         ẋ = [4.0, 5.0, 6.0]
         xcopy, ẋcopy = copy(x), copy(ẋ)
         y = [1, 2]
-        # Don't test tangent transforms, we do not support thunks for mutating frules
-        # TODO: Should we disable testing thunks for frules in general
-        test_frule(finplace!, x ⊢ ẋ; fkwargs=(y=y,), tangent_transforms=[])
+        test_frule(finplace!, x ⊢ ẋ; fkwargs=(y=y,))
         @test x == xcopy
         @test ẋ == ẋcopy
         @test y == [1, 2]
@@ -585,45 +583,6 @@ 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