Implement gradient_and_hvp (#588)

* Implement `gradient_and_hvp`

* Use inner

* Da fixes

* Add symbolic backends

* Fix

* Fix JuliaFormatter to v1

Author: gdalle

DifferentiationInterface/Project.toml

@@ -53,6 +53,7 @@ FastDifferentiation = "0.4.1"
 FiniteDiff = "2.23.1"
 FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36"
+JuliaFormatter = "1"
 LinearAlgebra = "<0.0.1,1"
 Mooncake = "0.4.0"
 PolyesterForwardDiff = "0.1.2"

DifferentiationInterface/docs/src/api.md

@@ -93,6 +93,8 @@ prepare_hvp
 prepare_hvp_same_point
 hvp
 hvp!
+gradient_and_hvp
+gradient_and_hvp!
 ```
 
 ### Hessian

DifferentiationInterface/docs/src/explanation/operators.md

@@ -58,7 +58,7 @@ Several variants of each operator are defined:
 | [`jacobian`](@ref)          | [`jacobian!`](@ref)          | [`value_and_jacobian`](@ref)                     | [`value_and_jacobian!`](@ref)                     |
 | [`pushforward`](@ref)       | [`pushforward!`](@ref)       | [`value_and_pushforward`](@ref)                  | [`value_and_pushforward!`](@ref)                  |
 | [`pullback`](@ref)          | [`pullback!`](@ref)          | [`value_and_pullback`](@ref)                     | [`value_and_pullback!`](@ref)                     |
-| [`hvp`](@ref)               | [`hvp!`](@ref)               | -                                                | -                                                 |
+| [`hvp`](@ref)               | [`hvp!`](@ref)               | [`gradient_and_hvp`](@ref)                       | [`gradient_and_hvp!`](@ref)                       |
 
 ## Mutation and signatures
 

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl

@@ -396,20 +396,23 @@ end
 
 ## HVP
 
-struct FastDifferentiationHVPPrep{E2,E2!} <: HVPPrep
+struct FastDifferentiationHVPPrep{E2,E2!,E1} <: HVPPrep
     hvp_exe::E2
     hvp_exe!::E2!
+    gradient_prep::E1
 end
 
-function DI.prepare_hvp(f, ::AutoFastDifferentiation, x, tx::NTuple)
+function DI.prepare_hvp(f, backend::AutoFastDifferentiation, x, tx::NTuple)
     x_var = make_variables(:x, size(x)...)
     y_var = f(x_var)
 
     x_vec_var = vec(x_var)
     hv_vec_var, v_vec_var = hessian_times_v(y_var, x_vec_var)
     hvp_exe = make_function(hv_vec_var, vcat(x_vec_var, v_vec_var); in_place=false)
     hvp_exe! = make_function(hv_vec_var, vcat(x_vec_var, v_vec_var); in_place=true)
-    return FastDifferentiationHVPPrep(hvp_exe, hvp_exe!)
+
+    gradient_prep = DI.prepare_gradient(f, backend, x)
+    return FastDifferentiationHVPPrep(hvp_exe, hvp_exe!, gradient_prep)
 end
 
 function DI.hvp(
@@ -439,6 +442,28 @@ function DI.hvp!(
     return tg
 end
 
+function DI.gradient_and_hvp(
+    f, prep::FastDifferentiationHVPPrep, backend::AutoFastDifferentiation, x, tx::NTuple
+)
+    tg = DI.hvp(f, prep, backend, x, tx)
+    grad = DI.gradient(f, prep.gradient_prep, backend, x)
+    return grad, tg
+end
+
+function DI.gradient_and_hvp!(
+    f,
+    grad,
+    tg::NTuple,
+    prep::FastDifferentiationHVPPrep,
+    backend::AutoFastDifferentiation,
+    x,
+    tx::NTuple,
+)
+    DI.hvp!(f, tg, prep, backend, x, tx)
+    DI.gradient!(f, grad, prep.gradient_prep, backend, x)
+    return grad, tg
+end
+
 ## Hessian
 
 struct FastDifferentiationHessianPrep{G,E2,E2!} <: HessianPrep

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -541,6 +541,32 @@ function DI.hvp!(
     return DI.hvp!(f, tg, prep, SecondOrder(backend, backend), x, tx, contexts...)
 end
 
+function DI.gradient_and_hvp(
+    f::F,
+    prep::HVPPrep,
+    backend::AutoForwardDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    return DI.gradient_and_hvp(f, prep, SecondOrder(backend, backend), x, tx, contexts...)
+end
+
+function DI.gradient_and_hvp!(
+    f::F,
+    grad,
+    tg::NTuple,
+    prep::HVPPrep,
+    backend::AutoForwardDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    return DI.gradient_and_hvp!(
+        f, grad, tg, prep, SecondOrder(backend, backend), x, tx, contexts...
+    )
+end
+
 ## Hessian
 
 ### Unprepared, only when chunk size and tag are not specified

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/secondorder.jl

@@ -32,6 +32,7 @@ function DI.prepare_hvp(
     T = tag_type(f, tagged_outer_backend, x)
     xdual = make_dual(T, x, tx)
     gradient_prep = DI.prepare_gradient(f, inner(backend), xdual, contexts...)
+    # TODO: get rid of closure?
     function inner_gradient(x, unannotated_contexts...)
         annotated_contexts = rewrap(unannotated_contexts...)
         return DI.gradient(f, gradient_prep, inner(backend), x, annotated_contexts...)
@@ -73,3 +74,34 @@ function DI.hvp!(
     )
     return tg
 end
+
+function DI.gradient_and_hvp(
+    f::F,
+    prep::ForwardDiffOverSomethingHVPPrep,
+    ::SecondOrder{<:AutoForwardDiff},
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    (; tagged_outer_backend, inner_gradient, outer_pushforward_prep) = prep
+    return DI.value_and_pushforward(
+        inner_gradient, outer_pushforward_prep, tagged_outer_backend, x, tx, contexts...
+    )
+end
+
+function DI.gradient_and_hvp!(
+    f::F,
+    grad,
+    tg::NTuple,
+    prep::ForwardDiffOverSomethingHVPPrep,
+    ::SecondOrder{<:AutoForwardDiff},
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    (; tagged_outer_backend, inner_gradient, outer_pushforward_prep) = prep
+    new_grad, _ = DI.value_and_pushforward!(
+        inner_gradient, tg, outer_pushforward_prep, tagged_outer_backend, x, tx, contexts...
+    )
+    return copyto!(grad, new_grad), tg
+end

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl

@@ -309,6 +309,32 @@ function DI.hvp!(
     return DI.hvp!(f, tg, prep, single_threaded(backend), x, tx, contexts...)
 end
 
+function DI.gradient_and_hvp(
+    f,
+    prep::HVPPrep,
+    backend::AutoPolyesterForwardDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {C}
+    return DI.gradient_and_hvp(f, prep, single_threaded(backend), x, tx, contexts...)
+end
+
+function DI.gradient_and_hvp!(
+    f,
+    grad,
+    tg::NTuple,
+    prep::HVPPrep,
+    backend::AutoPolyesterForwardDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {C}
+    return DI.gradient_and_hvp!(
+        f, grad, tg, prep, single_threaded(backend), x, tx, contexts...
+    )
+end
+
 ## Second derivative
 
 function DI.prepare_second_derivative(

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl

@@ -317,6 +317,22 @@ function DI.hvp!(
     return tg
 end
 
+function DI.gradient_and_hvp(
+    f, prep::SymbolicsOneArgHVPPrep, backend::AutoSymbolics, x, tx::NTuple
+)
+    tg = DI.hvp(f, prep, backend, x, tx)
+    grad = DI.gradient(f, prep.gradient_prep, backend, x)
+    return grad, tg
+end
+
+function DI.gradient_and_hvp!(
+    f, grad, tg::NTuple, prep::SymbolicsOneArgHVPPrep, backend::AutoSymbolics, x, tx::NTuple
+)
+    DI.hvp!(f, tg, prep, backend, x, tx)
+    DI.gradient!(f, grad, prep.gradient_prep, backend, x)
+    return grad, tg
+end
+
 ## Second derivative
 
 struct SymbolicsOneArgSecondDerivativePrep{D,E1,E1!} <: SecondDerivativePrep

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -171,6 +171,29 @@ function DI.hvp!(
     return DI.hvp!(f, tg, prep, SecondOrder(AutoForwardDiff(), backend), x, tx, contexts...)
 end
 
+function DI.gradient_and_hvp(
+    f, prep::HVPPrep, backend::AutoZygote, x, tx::NTuple, contexts::Vararg{Constant,C}
+) where {C}
+    return DI.gradient_and_hvp(
+        f, prep, SecondOrder(AutoForwardDiff(), backend), x, tx, contexts...
+    )
+end
+
+function DI.gradient_and_hvp!(
+    f,
+    grad,
+    tg::NTuple,
+    prep::HVPPrep,
+    backend::AutoZygote,
+    x,
+    tx::NTuple,
+    contexts::Vararg{Constant,C},
+) where {C}
+    return DI.gradient_and_hvp!(
+        f, grad, tg, prep, SecondOrder(AutoForwardDiff(), backend), x, tx, contexts...
+    )
+end
+
 ## Hessian
 
 function DI.prepare_hessian(f, ::AutoZygote, x, contexts::Vararg{Constant,C}) where {C}

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -84,7 +84,7 @@ export jacobian!, jacobian
 
 export second_derivative!, second_derivative
 export value_derivative_and_second_derivative, value_derivative_and_second_derivative!
-export hvp!, hvp
+export hvp!, hvp, gradient_and_hvp, gradient_and_hvp!
 export hessian!, hessian
 export value_gradient_and_hessian, value_gradient_and_hessian!