Tapir Safe Mode (#259)

* Respect request for safe mode

* Require ADTypes 1.2

* Make safety message appear

* Disable safe mode to avoid polluting test logs

* Formatting

Author: willtebbutt

DifferentiationInterface/Project.toml

@@ -43,7 +43,7 @@ DifferentiationInterfaceTrackerExt = "Tracker"
 DifferentiationInterfaceZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "1.0.0"
+ADTypes = "1.2.0"
 ChainRulesCore = "1.23.0"
 Compat = "3,4"
 Diffractor = "=0.2.6"

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/DifferentiationInterfaceTapirExt.jl

@@ -18,7 +18,8 @@ using Tapir:
     NoRData,
     fdata,
     rdata,
-    __value_and_pullback!!
+    __value_and_pullback!!,
+    TapirInterpreter
 
 DI.check_available(::AutoTapir) = true
 

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/onearg.jl

@@ -5,7 +5,12 @@ end
 
 function DI.prepare_pullback(f, backend::AutoTapir, x, dy)
     y = f(x)
-    rrule = build_rrule(f, x)
+    rrule = build_rrule(
+        TapirInterpreter(),
+        Tuple{typeof(f),typeof(x)};
+        safety_on=backend.safe_mode,
+        silence_safety_messages=false,
+    )
     extras = TapirOneArgPullbackExtras(y, rrule)
     DI.value_and_pullback(f, backend, x, dy, extras)  # warm up
     return extras

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/twoarg.jl

@@ -3,7 +3,12 @@ struct TapirTwoArgPullbackExtras{R} <: PullbackExtras
 end
 
 function DI.prepare_pullback(f!, y, backend::AutoTapir, x, dy)
-    rrule = build_rrule(f!, y, x)
+    rrule = build_rrule(
+        TapirInterpreter(),
+        Tuple{typeof(f!),typeof(y),typeof(x)};
+        safety_on=backend.safe_mode,
+        silence_safety_messages=false,
+    )
     extras = TapirTwoArgPullbackExtras(rrule)
     DI.value_and_pullback(f!, y, backend, x, dy, extras)  # warm up
     return extras

DifferentiationInterface/test/Single/Tapir.jl

@@ -1,10 +1,11 @@
 using DifferentiationInterface, DifferentiationInterfaceTest
 using Tapir: Tapir
 
-for backend in [AutoTapir()]
+for backend in [AutoTapir(; safe_mode=false)]
     @test check_available(backend)
     @test check_twoarg(backend)
     @test !check_hessian(backend; verbose=false)
 end
 
-test_differentiation(AutoTapir(); second_order=false, logging=LOGGING);
+# Safe mode switched off to avoid polluting the test suite with 
+test_differentiation(AutoTapir(; safe_mode=false); second_order=false, logging=LOGGING);