Verify Safe Rule (#150)

* Ensure safe rrule works with value_and_pullback etc

* Bump patch

Project.toml

@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.10"
+version = "0.2.11"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/interface.jl

@@ -23,6 +23,8 @@ function __verify_sig(::DerivedRule{<:OpaqueClosure{sig}}, ::Tfx) where {sig, Tf
     end
 end
 
+__verify_sig(rule::SafeRRule, fx) = __verify_sig(rule.rule, fx)
+
 # rrule!! doesn't specify specific argument types which must be used, so there's nothing to
 # check here.
 __verify_sig(::typeof(rrule!!), fx::Tuple) = nothing

test/interface.jl

@@ -8,12 +8,14 @@ end
         ([1.0, 1.0], x -> [sin(x), sin(2x)], 3.0),
         (1.0, x -> sum(5x), [5.0, 2.0]),
     ]
-        rule = build_rrule(f, x...)
-        v, (df, dx...) = value_and_pullback!!(rule, ȳ, f, x...)
-        @test v ≈ f(x...)
-        @test df isa tangent_type(typeof(f))
-        for (_dx, _x) in zip(dx, x)
-            @test _dx isa tangent_type(typeof(_x))
+        @testset "safe_mode=$safe_mode" for safe_mode in Bool[false, true]
+            rule = build_rrule(f, x...; safety_on=safe_mode)
+            v, (df, dx...) = value_and_pullback!!(rule, ȳ, f, x...)
+            @test v ≈ f(x...)
+            @test df isa tangent_type(typeof(f))
+            for (_dx, _x) in zip(dx, x)
+                @test _dx isa tangent_type(typeof(_x))
+            end
         end
     end
     @testset "sensible error when CoDuals are passed to `value_and_pullback!!" begin