confirm that check_thunked_output_tangent catches problems

Author: oxinabox

src/check_result.jl

@@ -143,7 +143,7 @@ function test_approx(actual::A, expected::E, msg=""; kwargs...) where {A,E}
         if (c_actual isa A) && (c_expected isa E)  # prevent stack-overflow
             throw(MethodError, test_approx, (actual, expected))
         end
-        test_approx(c_actual, c_expected; kwargs...)
+        test_approx(c_actual, c_expected, msg; kwargs...)
     end
 end
 

src/testers.jl

@@ -241,7 +241,7 @@ function test_rrule(
         end
 
         if check_thunked_output_tangent
-            test_approx(ad_cotangents, pullback(@thunk(ȳ)), "pulling back a thunk")
+            test_approx(ad_cotangents, pullback(@thunk(ȳ)), "pulling back a thunk:")
         end
     end  # top-level testset
 end

test/testers.jl

@@ -598,6 +598,35 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
         test_rrule(rev_trouble, (3, 3.0) ⊢ Tangent{Tuple{Int,Float64}}(ZeroTangent(), 1.0))
     end
 
+    @testset "check_thunked_output_tangent" begin
+        @testset "no method for thunk" begin
+            does_not_accept_thunk_id(x) = x
+            function ChainRulesCore.rrule(::typeof(does_not_accept_thunk_id), x)
+                does_not_accept_thunk_id_pullback(ȳ::AbstractArray) = (NoTangent() ,ȳ)
+                return does_not_accept_thunk_id(x), does_not_accept_thunk_id_pullback
+            end
+
+            test_rrule(
+                does_not_accept_thunk_id, [1.0, 2.0]; check_thunked_output_tangent=false
+            )
+            @test errors(r"MethodError.*Thunk") do 
+                test_rrule(does_not_accept_thunk_id, [1.0, 2.0])
+            end
+        end
+
+        @testset "Thunk wrong" begin
+            bad_thunk_id(x) = x
+            function ChainRulesCore.rrule(::typeof(bad_thunk_id), x)
+                bad_thunk_id_pullback(ȳ::AbstractArray) = (NoTangent(), ȳ)
+                bad_thunk_id_pullback(ȳ::AbstractThunk) = (NoTangent(), 2 * ȳ)
+                return bad_thunk_id(x), bad_thunk_id_pullback
+            end
+
+            test_rrule(bad_thunk_id, [1.0, 2.0]; check_thunked_output_tangent=false)
+            @test fails(()->test_rrule(bad_thunk_id, [1.0, 2.0]))
+        end
+    end
+
     @testset "error message about incorrectly using ZeroTangent()" begin
         foo(a, i) = a[i]
         function ChainRulesCore.rrule(::typeof(foo), a, i)