feat(DI): add value_gradient_and_hvp and deprecate gradient_and_hvp

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.21...main)
 
+### Added
+
+- New operators `value_gradient_and_hvp` and `value_gradient_and_hvp!`, which also return the primal value ([#1086](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1086))
+
+### Deprecated
+
+- `gradient_and_hvp` and `gradient_and_hvp!` are deprecated in favor of `value_gradient_and_hvp` and `value_gradient_and_hvp!` ([#1086](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1086))
+
 ### Fixed
 
 - Speed up `pick_batchsize` when the batch size is only known at run time ([#1076](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1076))

DifferentiationInterface/Project.toml

@@ -1,6 +1,6 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-version = "0.7.21"
+version = "0.7.22"
 authors = ["Guillaume Dalle", "Adrian Hill"]
 
 [deps]

DifferentiationInterface/docs/src/api.md

@@ -94,6 +94,8 @@ prepare_hvp
 prepare_hvp_same_point
 hvp
 hvp!
+value_gradient_and_hvp
+value_gradient_and_hvp!
 gradient_and_hvp
 gradient_and_hvp!
 ```

DifferentiationInterface/docs/src/explanation/operators.md

@@ -60,7 +60,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)               | [`gradient_and_hvp`](@ref)                       | [`gradient_and_hvp!`](@ref)                       |
+| [`hvp`](@ref)               | [`hvp!`](@ref)               | [`value_gradient_and_hvp`](@ref)                 | [`value_gradient_and_hvp!`](@ref)                 |
 
 ## Mutation and signatures
 

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl

@@ -599,7 +599,7 @@ function DI.hvp!(
     return tg
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f,
         prep::FastDifferentiationHVPPrep,
         backend::AutoFastDifferentiation,
@@ -609,11 +609,11 @@ function DI.gradient_and_hvp(
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     tg = DI.hvp(f, prep, backend, x, tx, contexts...)
-    grad = DI.gradient(f, prep.gradient_prep, backend, x, contexts...)
-    return grad, tg
+    y, grad = DI.value_and_gradient(f, prep.gradient_prep, backend, x, contexts...)
+    return y, grad, tg
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f,
         grad,
         tg::NTuple,
@@ -625,8 +625,8 @@ function DI.gradient_and_hvp!(
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     DI.hvp!(f, tg, prep, backend, x, tx, contexts...)
-    DI.gradient!(f, grad, prep.gradient_prep, backend, x, contexts...)
-    return grad, tg
+    y, _ = DI.value_and_gradient!(f, grad, prep.gradient_prep, backend, x, contexts...)
+    return y, grad, tg
 end
 
 ## Hessian

DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl

@@ -609,7 +609,7 @@ function DI.hvp!(
     return tg
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f,
         prep::GTPSAOneArgHVPPrep,
         backend::AutoGTPSA{D},
@@ -619,7 +619,7 @@ function DI.gradient_and_hvp(
     ) where {D, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     grad = similar(x, eltype(prep.hess))
-    DI.value_gradient_and_hessian!(
+    y, _, _ = DI.value_gradient_and_hessian!(
         f, grad, prep.hess, prep.hessprep, backend, x, contexts...
     )
     tg = map(tx) do dx
@@ -634,10 +634,10 @@ function DI.gradient_and_hvp(
         end
         return dg
     end
-    return grad, tg
+    return y, grad, tg
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f,
         grad,
         tg,
@@ -648,7 +648,7 @@ function DI.gradient_and_hvp!(
         contexts::Vararg{DI.Constant, C},
     ) where {D, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    DI.value_gradient_and_hessian!(
+    y, _, _ = DI.value_gradient_and_hessian!(
         f, grad, prep.hess, prep.hessprep, backend, x, contexts...
     )
     for b in eachindex(tg, tx)
@@ -662,5 +662,5 @@ function DI.gradient_and_hvp!(
             j += 1
         end
     end
-    return grad, tg
+    return y, grad, tg
 end

DifferentiationInterface/ext/DifferentiationInterfaceHyperHessiansExt/DifferentiationInterfaceHyperHessiansExt.jl

@@ -283,7 +283,7 @@ function DI.hvp!(
     return hvp!(tg, fc, x, tx, prep.cfg)
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f,
         prep::HyperHessiansHVPPrep,
         backend::DI.AutoHyperHessians,
@@ -295,10 +295,10 @@ function DI.gradient_and_hvp(
     contexts_prepared = translate_prepared(contexts, prep.contexts_prepared)
     fc = DI.fix_tail(f, contexts_prepared...)
     res = hvp_gradient_value(fc, x, tx, prep.cfg)
-    return res.gradient, res.hvp
+    return res.value, res.gradient, res.hvp
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f,
         grad,
         tg::NTuple,
@@ -311,8 +311,8 @@ function DI.gradient_and_hvp!(
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     contexts_prepared = translate_prepared(contexts, prep.contexts_prepared)
     fc = DI.fix_tail(f, contexts_prepared...)
-    hvp_gradient_value!(tg, grad, fc, x, tx, prep.cfg)
-    return grad, tg
+    y = hvp_gradient_value!(tg, grad, fc, x, tx, prep.cfg)
+    return y, grad, tg
 end
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/second_order.jl

@@ -13,7 +13,7 @@ function DI.prepare_hvp_nokwarg(
     return prep
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f::F,
         prep::MooncakeHVPPrep{Y},
         backend::AutoMooncakeForwardOverReverse,
@@ -24,19 +24,19 @@ function DI.gradient_and_hvp(
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     dx = only(tx)
     if isempty(contexts)
-        _, new_g, new_dg = value_and_hvp!!(
+        y, new_g, new_dg = value_and_hvp!!(
             prep.cache, f, (dx,), x
         )
     else
         dall = (dx, map(zero_tangent_unwrap, contexts)...)
-        _, (new_g,), (new_dg,) = value_and_hvp!!(
+        y, (new_g,), (new_dg,) = value_and_hvp!!(
             prep.cache, f, dall, x, map(DI.unwrap, contexts)...
         )
     end
-    return _copy_output(new_g), (_copy_output(new_dg),)
+    return y, _copy_output(new_g), (_copy_output(new_dg),)
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f::F,
         prep::MooncakeHVPPrep{Y},
         backend::AutoMooncakeForwardOverReverse,
@@ -45,25 +45,25 @@ function DI.gradient_and_hvp(
         contexts::Vararg{DI.Context, C},
     ) where {F, Y, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    gs_and_tg = map(tx) do dx
+    ys_gs_and_tg = map(tx) do dx
         if isempty(contexts)
-            _, new_g, new_dg = value_and_hvp!!(
+            y, new_g, new_dg = value_and_hvp!!(
                 prep.cache, f, dx, x
             )
         else
             dall = (dx, map(zero_tangent_unwrap, contexts)...)
-            _, (new_g,), (new_dg,) = value_and_hvp!!(
+            y, (new_g,), (new_dg,) = value_and_hvp!!(
                 prep.cache, f, dall, x, map(DI.unwrap, contexts)...
             )
         end
-        _copy_output(new_g), _copy_output(new_dg)
+        y, _copy_output(new_g), _copy_output(new_dg)
     end
-    g = first(gs_and_tg[1])
-    tg = map(last, gs_and_tg)
-    return g, tg
+    y, g, _ = ys_gs_and_tg[1]
+    tg = map(last, ys_gs_and_tg)
+    return y, g, tg
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f::F,
         g,
         tg::NTuple,
@@ -74,10 +74,10 @@ function DI.gradient_and_hvp!(
         contexts::Vararg{DI.Context, C},
     ) where {F, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    new_g, new_tg = DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)
+    y, new_g, new_tg = DI.value_gradient_and_hvp(f, prep, backend, x, tx, contexts...)
     copyto!(g, new_g)
     foreach(copyto!, tg, new_tg)
-    return g, tg
+    return y, g, tg
 end
 
 function DI.hvp(
@@ -89,7 +89,7 @@ function DI.hvp(
         contexts::Vararg{DI.Context, C},
     ) where {F, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    return DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)[2]
+    return DI.value_gradient_and_hvp(f, prep, backend, x, tx, contexts...)[3]
 end
 
 function DI.hvp!(
@@ -102,7 +102,7 @@ function DI.hvp!(
         contexts::Vararg{DI.Context, C},
     ) where {F, C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    _, new_tg = DI.gradient_and_hvp(f, prep, backend, x, tx, contexts...)
+    _, _, new_tg = DI.value_gradient_and_hvp(f, prep, backend, x, tx, contexts...)
     foreach(copyto!, tg, new_tg)
     return tg
 end

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl

@@ -483,7 +483,7 @@ function DI.hvp!(
     return tg
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f,
         prep::SymbolicsOneArgHVPPrep,
         backend::AutoSymbolics,
@@ -493,11 +493,11 @@ function DI.gradient_and_hvp(
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     tg = DI.hvp(f, prep, backend, x, tx, contexts...)
-    grad = DI.gradient(f, prep.gradient_prep, backend, x, contexts...)
-    return grad, tg
+    y, grad = DI.value_and_gradient(f, prep.gradient_prep, backend, x, contexts...)
+    return y, grad, tg
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f,
         grad,
         tg::NTuple,
@@ -509,8 +509,8 @@ function DI.gradient_and_hvp!(
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
     DI.hvp!(f, tg, prep, backend, x, tx, contexts...)
-    DI.gradient!(f, grad, prep.gradient_prep, backend, x, contexts...)
-    return grad, tg
+    y, _ = DI.value_and_gradient!(f, grad, prep.gradient_prep, backend, x, contexts...)
+    return y, grad, tg
 end
 
 ## Second derivative

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -229,7 +229,7 @@ function DI.hvp!(
     )
 end
 
-function DI.gradient_and_hvp(
+function DI.value_gradient_and_hvp(
         f,
         prep::ZygoteHVPPrep,
         backend::AutoZygote,
@@ -238,12 +238,12 @@ function DI.gradient_and_hvp(
         contexts::Vararg{DI.Context, C},
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    return DI.gradient_and_hvp(
+    return DI.value_gradient_and_hvp(
         f, prep.fd_prep, DI.SecondOrder(AutoForwardDiff(), backend), x, tx, contexts...
     )
 end
 
-function DI.gradient_and_hvp!(
+function DI.value_gradient_and_hvp!(
         f,
         grad,
         tg::NTuple,
@@ -254,7 +254,7 @@ function DI.gradient_and_hvp!(
         contexts::Vararg{DI.Context, C},
     ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    return DI.gradient_and_hvp!(
+    return DI.value_gradient_and_hvp!(
         f,
         grad,
         tg,

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -92,7 +92,8 @@ export jacobian!, jacobian
 
 export second_derivative!, second_derivative
 export value_derivative_and_second_derivative, value_derivative_and_second_derivative!
-export hvp!, hvp, gradient_and_hvp, gradient_and_hvp!
+export hvp!, hvp, value_gradient_and_hvp, value_gradient_and_hvp!
+export gradient_and_hvp, gradient_and_hvp!
 export hessian!, hessian
 export value_gradient_and_hessian, value_gradient_and_hessian!