Catch inference problems for thunked cotangents (#215)

Author: mzgubic

Project.toml

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

src/testers.jl

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

test/testers.jl

@@ -57,7 +57,6 @@ end
 abstract type MySpecialTrait end
 struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
 
-
 @testset "testers.jl" begin
     @testset "test_scalar" begin
         @testset "Ensure correct rules succeed" begin
@@ -711,4 +710,20 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
 
         ChainRulesTestUtils.TEST_INFERRED[] = true
     end
+
+    @testset "inference of thunked cotangents" begin
+        my_id(x) = x
+        function ChainRulesCore.rrule(::typeof(my_id), x)
+            my_id_pb(ȳ) = (NoTangent(), ȳ)
+            function my_id_pb(ȳ::AbstractThunk)
+                precision = rand() > 0.5 ? Float64 : Float32
+                return (NoTangent(), precision(unthunk(ȳ)))
+            end
+            return x, my_id_pb
+        end
+
+        @test errors(() -> test_rrule(my_id, 2.0))
+        test_rrule(my_id, 2.0; check_inferred=false)
+        test_rrule(my_id, 2.0; check_thunked_output_tangent=false)
+    end
 end