Get rid of functors for second order (#460)

* Get rid of functors for second order

* Types

* Fix

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/secondorder.jl

@@ -18,9 +18,8 @@ function tag_backend_hvp(f, backend::AutoForwardDiff, x)
     return backend
 end
 
-struct ForwardDiffOverSomethingHVPExtras{
-    B<:AutoForwardDiff,G<:DI.Gradient,E<:PushforwardExtras
-} <: HVPExtras
+struct ForwardDiffOverSomethingHVPExtras{B<:AutoForwardDiff,G,E<:PushforwardExtras} <:
+       HVPExtras
     tagged_outer_backend::B
     inner_gradient::G
     outer_pushforward_extras::E
@@ -33,7 +32,7 @@ function DI.prepare_hvp(
     T = tag_type(f, tagged_outer_backend, x)
     xdual = make_dual(T, x, tx)
     gradient_extras = DI.prepare_gradient(f, inner(backend), xdual)
-    inner_gradient = DI.Gradient(f, inner(backend), gradient_extras)
+    inner_gradient(x) = DI.gradient(f, gradient_extras, inner(backend), x)
     outer_pushforward_extras = DI.prepare_pushforward(
         inner_gradient, tagged_outer_backend, x, tx
     )

DifferentiationInterface/src/first_order/gradient.jl

@@ -88,53 +88,3 @@ function gradient!(
     pullback!(f, Tangents(grad), extras.pullback_extras, backend, x, Tangents(true))
     return grad
 end
-
-## Functors
-
-"""
-    Gradient
-
-Functor computing the gradient of `f` with a fixed `backend`.
-
-!!! warning
-    This type is not part of the public API.
-
-# Constructor
-
-    Gradient(f, backend, extras=nothing)
-
-If `extras` is provided, the gradient closure will skip preparation.
-
-# Example
-
-```jldoctest
-using DifferentiationInterface
-import Zygote
-
-g = DifferentiationInterface.Gradient(x -> sum(abs2, x), AutoZygote())
-g([2.0, 3.0])
-
-# output
-
-2-element Vector{Float64}:
- 4.0
- 6.0
-```
-"""
-struct Gradient{F,B,E}
-    f::F
-    backend::B
-    extras::E
-end
-
-Gradient(f, backend::AbstractADType) = Gradient(f, backend, nothing)
-
-function (g::Gradient{F,B,Nothing})(x) where {F,B}
-    @compat (; f, backend) = g
-    return gradient(f, backend, x)
-end
-
-function (g::Gradient{F,B,<:GradientExtras})(x) where {F,B}
-    @compat (; f, backend, extras) = g
-    return gradient(f, extras, backend, x)
-end

DifferentiationInterface/src/first_order/pushforward.jl

@@ -242,22 +242,3 @@ function pushforward!(
 ) where {F}
     return value_and_pushforward!(f!, y, ty, extras, backend, x, tx)[2]
 end
-
-## Functors
-
-struct PushforwardFixedSeed{F,B,TX,E}
-    f::F
-    backend::B
-    tx::TX
-    extras::E
-end
-
-function PushforwardFixedSeed(f, backend::AbstractADType, tx)
-    return PushforwardFixedSeed(f, backend, tx, nothing)
-end
-
-# not covered but don't remove, Enzyme messes with code coverage
-function (pfs::PushforwardFixedSeed{F,B,TX,Nothing})(x) where {F,B,TX}
-    @compat (; f, backend, tx) = pfs
-    return pushforward(f, backend, x, tx)
-end

DifferentiationInterface/src/second_order/hvp.jl

@@ -40,19 +40,19 @@ function hvp! end
 
 ## Preparation
 
-struct ForwardOverForwardHVPExtras{G<:Gradient,E<:PushforwardExtras} <: HVPExtras
+struct ForwardOverForwardHVPExtras{G,E<:PushforwardExtras} <: HVPExtras
     inner_gradient::G
     outer_pushforward_extras::E
 end
 
-struct ForwardOverReverseHVPExtras{G<:Gradient,E<:PushforwardExtras} <: HVPExtras
+struct ForwardOverReverseHVPExtras{G,E<:PushforwardExtras} <: HVPExtras
     inner_gradient::G
     outer_pushforward_extras::E
 end
 
 struct ReverseOverForwardHVPExtras <: HVPExtras end
 
-struct ReverseOverReverseHVPExtras{G<:Gradient,E<:PullbackExtras} <: HVPExtras
+struct ReverseOverReverseHVPExtras{G,E<:PullbackExtras} <: HVPExtras
     inner_gradient::G
     outer_pullback_extras::E
 end
@@ -65,7 +65,7 @@ function _prepare_hvp_aux(
     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(maybe_inner(backend)))
+    inner_gradient(x) = gradient(f, nested(maybe_inner(backend)), x)
     outer_pushforward_extras = prepare_pushforward(
         inner_gradient, maybe_outer(backend), x, tx
     )
@@ -76,7 +76,7 @@ function _prepare_hvp_aux(
     f::F, backend::AbstractADType, x, tx::Tangents, ::ForwardOverReverse
 ) where {F}
     # pushforward of gradient
-    inner_gradient = Gradient(f, nested(maybe_inner(backend)))
+    inner_gradient(x) = gradient(f, nested(maybe_inner(backend)), x)
     outer_pushforward_extras = prepare_pushforward(
         inner_gradient, maybe_outer(backend), x, tx
     )
@@ -95,7 +95,7 @@ function _prepare_hvp_aux(
     f::F, backend::AbstractADType, x, tx::Tangents, ::ReverseOverReverse
 ) where {F}
     # pullback of gradient
-    inner_gradient = Gradient(f, nested(maybe_inner(backend)))
+    inner_gradient(x) = gradient(f, nested(maybe_inner(backend)), x)
     outer_pullback_extras = prepare_pullback(inner_gradient, maybe_outer(backend), x, tx)
     return ReverseOverReverseHVPExtras(inner_gradient, outer_pullback_extras)
 end
@@ -123,11 +123,13 @@ end
 function hvp(
     f::F, ::ReverseOverForwardHVPExtras, backend::AbstractADType, x, tx::Tangents
 ) where {F}
-    dgs = map(tx.d) do dx
-        inner_pushforward = PushforwardFixedSeed(f, nested(maybe_inner(backend)), Tangents(dx))
+    tg = map(tx) do dx
+        function inner_pushforward(x)
+            return only(pushforward(f, nested(maybe_inner(backend)), x, Tangents(dx)))
+        end
         gradient(only ∘ inner_pushforward, maybe_outer(backend), x)
     end
-    return Tangents(dgs...)
+    return tg
 end
 
 function hvp(
@@ -174,9 +176,9 @@ function hvp!(
     tx::Tangents,
 ) where {F}
     for b in eachindex(tx.d, tg.d)
-        inner_pushforward = PushforwardFixedSeed(
-            f, nested(maybe_inner(backend)), Tangents(tx.d[b])
-        )
+        function inner_pushforward(x)
+            return only(pushforward(f, nested(maybe_inner(backend)), x, Tangents(tx.d[b])))
+        end
         gradient!(only ∘ inner_pushforward, tg.d[b], maybe_outer(backend), x)
     end
     return tg

DifferentiationInterface/src/second_order/second_derivative.jl

@@ -48,24 +48,13 @@ function value_derivative_and_second_derivative! end
 
 ## Preparation
 
-struct InnerDerivative{F,B}
-    f::F
-    backend::B
-end
-
-function (id::InnerDerivative)(x)
-    @compat (; f, backend) = id
-    return derivative(f, backend, x)
-end
-
-struct ClosureSecondDerivativeExtras{ID<:InnerDerivative,E<:DerivativeExtras} <:
-       SecondDerivativeExtras
+struct ClosureSecondDerivativeExtras{ID,E<:DerivativeExtras} <: SecondDerivativeExtras
     inner_derivative::ID
     outer_derivative_extras::E
 end
 
 function prepare_second_derivative(f::F, backend::AbstractADType, x) where {F}
-    inner_derivative = InnerDerivative(f, nested(maybe_inner(backend)))
+    inner_derivative(x) = derivative(f, nested(maybe_inner(backend)), x)
     outer_derivative_extras = prepare_derivative(inner_derivative, maybe_outer(backend), x)
     return ClosureSecondDerivativeExtras(inner_derivative, outer_derivative_extras)
 end