Better preparation for second order with single backend (#287)

Author: gdalle

DifferentiationInterface/docs/src/operators.md

@@ -135,9 +135,6 @@ Some backends natively support a set of second-order operators (typically only t
 In that case, it can be advantageous to use the backend on its own.
 If the operator is not supported natively, we will fall back on `SecondOrder(backend, backend)` (see below).
 
-!!! warning
-    Whenever the fallback on `SecondOrder(backend, backend)` occurs, the results of any preparation will be discarded.
-
 ### Combining backends
 
 In general, you can use [`SecondOrder`](@ref) to combine different backends.

DifferentiationInterface/src/second_order/hessian.jl

@@ -40,6 +40,10 @@ struct HVPHessianExtras{E<:HVPExtras} <: HessianExtras
 end
 
 function prepare_hessian(f::F, backend::AbstractADType, x) where {F}
+    return prepare_hessian(f, SecondOrder(backend, backend), x)
+end
+
+function prepare_hessian(f::F, backend::SecondOrder, x) where {F}
     v = basis(backend, x, first(CartesianIndices(x)))
     hvp_extras = prepare_hvp(f, backend, x, v)
     return HVPHessianExtras(hvp_extras)
@@ -50,9 +54,7 @@ end
 function hessian(
     f::F, backend::AbstractADType, x, extras::HessianExtras=prepare_hessian(f, backend, x)
 ) where {F}
-    new_backend = SecondOrder(backend, backend)
-    new_extras = prepare_hessian(f, new_backend, x)
-    return hessian(f, new_backend, x, new_extras)
+    return hessian(f, SecondOrder(backend, backend), x, extras)
 end
 
 function hessian(

DifferentiationInterface/src/second_order/hvp.jl

@@ -75,7 +75,9 @@ struct ReverseOverReverseHVPExtras{C,E} <: HVPExtras
     outer_pullback_extras::E
 end
 
-prepare_hvp(f, ::AbstractADType, x, v) = NoHVPExtras()
+function prepare_hvp(f::F, backend::AbstractADType, x, v) where {F}
+    return prepare_hvp(f, SecondOrder(backend, backend), x, v)
+end
 
 function prepare_hvp(f::F, backend::SecondOrder, x, v) where {F}
     return prepare_hvp_aux(f, backend, x, v, hvp_mode(backend))
@@ -143,9 +145,7 @@ end
 function hvp(
     f::F, backend::AbstractADType, x, v, extras::HVPExtras=prepare_hvp(f, backend, x, v)
 ) where {F}
-    new_backend = SecondOrder(backend, backend)
-    new_extras = prepare_hvp(f, new_backend, x, v)
-    return hvp(f, new_backend, x, v, new_extras)
+    return hvp(f, SecondOrder(backend, backend), x, v, extras)
 end
 
 function hvp(

DifferentiationInterface/src/second_order/second_derivative.jl

@@ -40,7 +40,9 @@ struct ClosureSecondDerivativeExtras{C,E} <: SecondDerivativeExtras
     outer_derivative_extras::E
 end
 
-prepare_second_derivative(f::F, ::AbstractADType, x) where {F} = NoSecondDerivativeExtras()
+function prepare_second_derivative(f::F, backend::AbstractADType, x) where {F}
+    return prepare_second_derivative(f, SecondOrder(backend, backend), x)
+end
 
 function prepare_second_derivative(f::F, backend::SecondOrder, x) where {F}
     inner_backend = nested(inner(backend))
@@ -59,9 +61,7 @@ function second_derivative(
     x,
     extras::SecondDerivativeExtras=prepare_second_derivative(f, backend, x),
 ) where {F}
-    new_backend = SecondOrder(backend, backend)
-    new_extras = prepare_second_derivative(f, new_backend, x)
-    return second_derivative(f, new_backend, x, new_extras)
+    return second_derivative(f, SecondOrder(backend, backend), x, extras)
 end
 
 function second_derivative(