Warm up for Tapir's rrule (#254)

* Warm up for Tapir's rrule

* Almost recursive warm up

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/onearg.jl

@@ -3,9 +3,12 @@ struct TapirOneArgPullbackExtras{Y,R} <: PullbackExtras
     rrule::R
 end
 
-function DI.prepare_pullback(f, ::AutoTapir, x, dy)
+function DI.prepare_pullback(f, backend::AutoTapir, x, dy)
     y = f(x)
-    return TapirOneArgPullbackExtras(y, build_rrule(f, x))
+    rrule = build_rrule(f, x)
+    extras = TapirOneArgPullbackExtras(y, rrule)
+    DI.value_and_pullback(f, backend, x, dy, extras)  # warm up
+    return extras
 end
 
 function DI.value_and_pullback(

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/twoarg.jl

@@ -2,8 +2,11 @@ struct TapirTwoArgPullbackExtras{R} <: PullbackExtras
     rrule::R
 end
 
-function DI.prepare_pullback(f!, y, ::AutoTapir, x, dy)
-    return TapirTwoArgPullbackExtras(build_rrule(f!, y, x))
+function DI.prepare_pullback(f!, y, backend::AutoTapir, x, dy)
+    rrule = build_rrule(f!, y, x)
+    extras = TapirTwoArgPullbackExtras(rrule)
+    DI.value_and_pullback(f!, y, backend, x, dy, extras)  # warm up
+    return extras
 end
 
 # see https://github.com/withbayes/Tapir.jl/issues/113#issuecomment-2036718992