Remove fallback backend -> SecondOrder(backend, backend) (#459)

* Remove fallback `backend -> SecondOrder(backend, backend)`

* Add hvp mode

Author: gdalle

DifferentiationInterface/src/second_order/hvp.jl

@@ -58,123 +58,126 @@ struct ReverseOverReverseHVPExtras{G<:Gradient,E<:PullbackExtras} <: HVPExtras
 end
 
 function prepare_hvp(f::F, backend::AbstractADType, x, tx::Tangents) where {F}
-    return prepare_hvp(f, SecondOrder(backend, backend), x, tx)
-end
-
-function prepare_hvp(f::F, backend::SecondOrder, x, tx::Tangents) where {F}
     return _prepare_hvp_aux(f, backend, x, tx, hvp_mode(backend))
 end
 
 function _prepare_hvp_aux(
-    f::F, backend::SecondOrder, x, tx::Tangents, ::ForwardOverForward
+    f::F, backend::AbstractADType, x, tx::Tangents, ::ForwardOverForward
 ) where {F}
     # pushforward of many pushforwards in theory, but pushforward of gradient in practice
-    inner_gradient = Gradient(f, nested(inner(backend)))
-    outer_pushforward_extras = prepare_pushforward(inner_gradient, outer(backend), x, tx)
+    inner_gradient = Gradient(f, nested(maybe_inner(backend)))
+    outer_pushforward_extras = prepare_pushforward(
+        inner_gradient, maybe_outer(backend), x, tx
+    )
     return ForwardOverForwardHVPExtras(inner_gradient, outer_pushforward_extras)
 end
 
 function _prepare_hvp_aux(
-    f::F, backend::SecondOrder, x, tx::Tangents, ::ForwardOverReverse
+    f::F, backend::AbstractADType, x, tx::Tangents, ::ForwardOverReverse
 ) where {F}
     # pushforward of gradient
-    inner_gradient = Gradient(f, nested(inner(backend)))
-    outer_pushforward_extras = prepare_pushforward(inner_gradient, outer(backend), x, tx)
+    inner_gradient = Gradient(f, nested(maybe_inner(backend)))
+    outer_pushforward_extras = prepare_pushforward(
+        inner_gradient, maybe_outer(backend), x, tx
+    )
     return ForwardOverReverseHVPExtras(inner_gradient, outer_pushforward_extras)
 end
 
 function _prepare_hvp_aux(
-    f::F, backend::SecondOrder, x, tx::Tangents, ::ReverseOverForward
+    f::F, backend::AbstractADType, x, tx::Tangents, ::ReverseOverForward
 ) where {F}
     # gradient of pushforward
     # uses dx in the closure so it can't be prepared
     return ReverseOverForwardHVPExtras()
 end
 
 function _prepare_hvp_aux(
-    f::F, backend::SecondOrder, x, tx::Tangents, ::ReverseOverReverse
+    f::F, backend::AbstractADType, x, tx::Tangents, ::ReverseOverReverse
 ) where {F}
     # pullback of gradient
-    inner_gradient = Gradient(f, nested(inner(backend)))
-    outer_pullback_extras = prepare_pullback(inner_gradient, outer(backend), x, tx)
+    inner_gradient = Gradient(f, nested(maybe_inner(backend)))
+    outer_pullback_extras = prepare_pullback(inner_gradient, maybe_outer(backend), x, tx)
     return ReverseOverReverseHVPExtras(inner_gradient, outer_pullback_extras)
 end
 
 ## One argument
 
-function hvp(f::F, extras::HVPExtras, backend::AbstractADType, x, tx::Tangents) where {F}
-    return hvp(f, extras, SecondOrder(backend, backend), x, tx)
-end
-
 function hvp(
-    f::F, extras::ForwardOverForwardHVPExtras, backend::SecondOrder, x, tx::Tangents
+    f::F, extras::ForwardOverForwardHVPExtras, backend::AbstractADType, x, tx::Tangents
 ) where {F}
     @compat (; inner_gradient, outer_pushforward_extras) = extras
-    return pushforward(inner_gradient, outer_pushforward_extras, outer(backend), x, tx)
+    return pushforward(
+        inner_gradient, outer_pushforward_extras, maybe_outer(backend), x, tx
+    )
 end
 
 function hvp(
-    f::F, extras::ForwardOverReverseHVPExtras, backend::SecondOrder, x, tx::Tangents
+    f::F, extras::ForwardOverReverseHVPExtras, backend::AbstractADType, x, tx::Tangents
 ) where {F}
     @compat (; inner_gradient, outer_pushforward_extras) = extras
-    return pushforward(inner_gradient, outer_pushforward_extras, outer(backend), x, tx)
+    return pushforward(
+        inner_gradient, outer_pushforward_extras, maybe_outer(backend), x, tx
+    )
 end
 
 function hvp(
-    f::F, ::ReverseOverForwardHVPExtras, backend::SecondOrder, x, tx::Tangents
+    f::F, ::ReverseOverForwardHVPExtras, backend::AbstractADType, x, tx::Tangents
 ) where {F}
     dgs = map(tx.d) do dx
-        inner_pushforward = PushforwardFixedSeed(f, nested(inner(backend)), Tangents(dx))
-        gradient(only ∘ inner_pushforward, outer(backend), x)
+        inner_pushforward = PushforwardFixedSeed(f, nested(maybe_inner(backend)), Tangents(dx))
+        gradient(only ∘ inner_pushforward, maybe_outer(backend), x)
     end
     return Tangents(dgs...)
 end
 
 function hvp(
-    f::F, extras::ReverseOverReverseHVPExtras, backend::SecondOrder, x, tx::Tangents
+    f::F, extras::ReverseOverReverseHVPExtras, backend::AbstractADType, x, tx::Tangents
 ) where {F}
     @compat (; inner_gradient, outer_pullback_extras) = extras
-    return pullback(inner_gradient, outer_pullback_extras, outer(backend), x, tx)
-end
-
-function hvp!(
-    f::F, tg::Tangents, extras::HVPExtras, backend::AbstractADType, x, tx::Tangents
-) where {F}
-    return hvp!(f, tg, extras, SecondOrder(backend, backend), x, tx)
+    return pullback(inner_gradient, outer_pullback_extras, maybe_outer(backend), x, tx)
 end
 
 function hvp!(
     f::F,
     tg::Tangents,
     extras::ForwardOverForwardHVPExtras,
-    backend::SecondOrder,
+    backend::AbstractADType,
     x,
     tx::Tangents,
 ) where {F}
     @compat (; inner_gradient, outer_pushforward_extras) = extras
-    return pushforward!(inner_gradient, tg, outer_pushforward_extras, outer(backend), x, tx)
+    return pushforward!(
+        inner_gradient, tg, outer_pushforward_extras, maybe_outer(backend), x, tx
+    )
 end
 
 function hvp!(
     f::F,
     tg::Tangents,
     extras::ForwardOverReverseHVPExtras,
-    backend::SecondOrder,
+    backend::AbstractADType,
     x,
     tx::Tangents,
 ) where {F}
     @compat (; inner_gradient, outer_pushforward_extras) = extras
-    return pushforward!(inner_gradient, tg, outer_pushforward_extras, outer(backend), x, tx)
+    return pushforward!(
+        inner_gradient, tg, outer_pushforward_extras, maybe_outer(backend), x, tx
+    )
 end
 
 function hvp!(
-    f::F, tg::Tangents, ::ReverseOverForwardHVPExtras, backend::SecondOrder, x, tx::Tangents
+    f::F,
+    tg::Tangents,
+    ::ReverseOverForwardHVPExtras,
+    backend::AbstractADType,
+    x,
+    tx::Tangents,
 ) where {F}
     for b in eachindex(tx.d, tg.d)
         inner_pushforward = PushforwardFixedSeed(
-            f, nested(inner(backend)), Tangents(tx.d[b])
+            f, nested(maybe_inner(backend)), Tangents(tx.d[b])
         )
-        gradient!(only ∘ inner_pushforward, tg.d[b], outer(backend), x)
+        gradient!(only ∘ inner_pushforward, tg.d[b], maybe_outer(backend), x)
     end
     return tg
 end
@@ -183,10 +186,10 @@ function hvp!(
     f::F,
     tg::Tangents,
     extras::ReverseOverReverseHVPExtras,
-    backend::SecondOrder,
+    backend::AbstractADType,
     x,
     tx::Tangents,
 ) where {F}
     @compat (; inner_gradient, outer_pullback_extras) = extras
-    return pullback!(inner_gradient, tg, outer_pullback_extras, outer(backend), x, tx)
+    return pullback!(inner_gradient, tg, outer_pullback_extras, maybe_outer(backend), x, tx)
 end

DifferentiationInterface/src/second_order/second_derivative.jl

@@ -65,77 +65,47 @@ struct ClosureSecondDerivativeExtras{ID<:InnerDerivative,E<:DerivativeExtras} <:
 end
 
 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_derivative = InnerDerivative(f, nested(inner(backend)))
-    outer_derivative_extras = prepare_derivative(inner_derivative, outer(backend), x)
+    inner_derivative = InnerDerivative(f, nested(maybe_inner(backend)))
+    outer_derivative_extras = prepare_derivative(inner_derivative, maybe_outer(backend), x)
     return ClosureSecondDerivativeExtras(inner_derivative, outer_derivative_extras)
 end
 
 ## One argument
 
 function second_derivative(
-    f::F, extras::SecondDerivativeExtras, backend::AbstractADType, x
-) where {F}
-    return second_derivative(f, extras, SecondOrder(backend, backend), x)
-end
-
-function second_derivative(
-    f::F, extras::ClosureSecondDerivativeExtras, backend::SecondOrder, x
+    f::F, extras::ClosureSecondDerivativeExtras, backend::AbstractADType, x
 ) where {F}
     @compat (; inner_derivative, outer_derivative_extras) = extras
-    return derivative(inner_derivative, outer_derivative_extras, outer(backend), x)
-end
-
-function value_derivative_and_second_derivative(
-    f::F, extras::SecondDerivativeExtras, backend::AbstractADType, x
-) where {F}
-    return value_derivative_and_second_derivative(
-        f, extras, SecondOrder(backend, backend), x
-    )
+    return derivative(inner_derivative, outer_derivative_extras, maybe_outer(backend), x)
 end
 
 function value_derivative_and_second_derivative(
-    f::F, extras::ClosureSecondDerivativeExtras, backend::SecondOrder, x
+    f::F, extras::ClosureSecondDerivativeExtras, backend::AbstractADType, x
 ) where {F}
     @compat (; inner_derivative, outer_derivative_extras) = extras
     y = f(x)
     der, der2 = value_and_derivative(
-        inner_derivative, outer_derivative_extras, outer(backend), x
+        inner_derivative, outer_derivative_extras, maybe_outer(backend), x
     )
     return y, der, der2
 end
 
 function second_derivative!(
     f::F, der2, extras::SecondDerivativeExtras, backend::AbstractADType, x
-) where {F}
-    return second_derivative!(f, der2, extras, SecondOrder(backend, backend), x)
-end
-
-function second_derivative!(
-    f::F, der2, extras::SecondDerivativeExtras, backend::SecondOrder, x
 ) where {F}
     @compat (; inner_derivative, outer_derivative_extras) = extras
-    return derivative!(inner_derivative, der2, outer_derivative_extras, outer(backend), x)
-end
-
-function value_derivative_and_second_derivative!(
-    f::F, der, der2, extras::SecondDerivativeExtras, backend::AbstractADType, x
-) where {F}
-    return value_derivative_and_second_derivative!(
-        f, der, der2, extras, SecondOrder(backend, backend), x
+    return derivative!(
+        inner_derivative, der2, outer_derivative_extras, maybe_outer(backend), x
     )
 end
 
 function value_derivative_and_second_derivative!(
-    f::F, der, der2, extras::SecondDerivativeExtras, backend::SecondOrder, x
+    f::F, der, der2, extras::SecondDerivativeExtras, backend::AbstractADType, x
 ) where {F}
     @compat (; inner_derivative, outer_derivative_extras) = extras
     y = f(x)
     new_der, _ = value_and_derivative!(
-        inner_derivative, der2, outer_derivative_extras, outer(backend), x
+        inner_derivative, der2, outer_derivative_extras, maybe_outer(backend), x
     )
     return y, copyto!(der, new_der), der2
 end

DifferentiationInterface/src/utils/traits.jl

@@ -125,6 +125,8 @@ Traits identifying second-order backends that compute HVPs in forward over forwa
 """
 struct ForwardOverForward <: HVPMode end
 
+hvp_mode(backend::AbstractADType) = hvp_mode(SecondOrder(backend, backend))
+
 function hvp_mode(ba::SecondOrder)
     if Bool(pushforward_performance(outer(ba))) && Bool(pullback_performance(inner(ba)))
         return ForwardOverReverse()