Operators that combine first and second order (#300)

* Update fallback for mutating second-order operators

* Enable operators that combine first and second order

* Docstring

* Scenarios

* Sparse scenarios

* Scenario intact

* Fix benchmarks

* Fix ref

* Logging

* Useless logging

* Typo

* Typo

* FastDiff fix

* Format

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.5.3"
+version = "0.5.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/docs/src/api.md

@@ -74,6 +74,8 @@ SecondOrder
 prepare_second_derivative
 second_derivative
 second_derivative!
+value_derivative_and_second_derivative
+value_derivative_and_second_derivative!
 ```
 
 ### Hessian-vector product

DifferentiationInterface/docs/src/operators.md

@@ -45,16 +45,16 @@ These operators are computed using the input `x` and a "seed" `v`, which lives e
 
 Several variants of each operator are defined.
 
-| out-of-place                | in-place                     | out-of-place + primal           | in-place + primal                |
-| :-------------------------- | :--------------------------- | :------------------------------ | :------------------------------- |
-| [`derivative`](@ref)        | [`derivative!`](@ref)        | [`value_and_derivative`](@ref)  | [`value_and_derivative!`](@ref)  |
-| [`second_derivative`](@ref) | [`second_derivative!`](@ref) | NA                              | NA                               |
-| [`gradient`](@ref)          | [`gradient!`](@ref)          | [`value_and_gradient`](@ref)    | [`value_and_gradient!`](@ref)    |
-| [`hessian`](@ref)           | [`hessian!`](@ref)           | NA                              | NA                               |
-| [`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)               | NA                              | NA                               |
+| out-of-place                | in-place                     | out-of-place + primal                            | in-place + primal                                |
+| :-------------------------- | :--------------------------- | :----------------------------------------------- | :----------------------------------------------- |
+| [`derivative`](@ref)        | [`derivative!`](@ref)        | [`value_and_derivative`](@ref)                   | [`value_and_derivative!`](@ref)                  |
+| [`second_derivative`](@ref) | [`second_derivative!`](@ref) | [`value_derivative_and_second_derivative`](@ref) | [`value_derivative_and_second_derivative!`](@ref) |
+| [`gradient`](@ref)          | [`gradient!`](@ref)          | [`value_and_gradient`](@ref)                     | [`value_and_gradient!`](@ref)                    |
+| [`hessian`](@ref)           | [`hessian!`](@ref)           | NA                                               | NA                                               |
+| [`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)               | NA                                               | NA                                               |
 
 ## Mutation and signatures
 

DifferentiationInterface/ext/DifferentiationInterfaceFastDifferentiationExt/onearg.jl

@@ -307,11 +307,13 @@ end
 
 ## Second derivative
 
-struct FastDifferentiationAllocatingSecondDerivativeExtras{Y,E1,E2} <:
+struct FastDifferentiationAllocatingSecondDerivativeExtras{Y,E1,E1!,E2,E2!} <:
        SecondDerivativeExtras
     y_prototype::Y
-    der2_exe::E1
-    der2_exe!::E2
+    der_exe::E1
+    der_exe!::E1!
+    der2_exe::E2
+    der2_exe!::E2!
 end
 
 function DI.prepare_second_derivative(f, ::AutoFastDifferentiation, x)
@@ -321,11 +323,18 @@ function DI.prepare_second_derivative(f, ::AutoFastDifferentiation, x)
 
     x_vec_var = monovec(x_var)
     y_vec_var = y_var isa Number ? monovec(y_var) : vec(y_var)
+
+    der_vec_var = derivative(y_vec_var, x_var)
     der2_vec_var = derivative(y_vec_var, x_var, x_var)
+
+    der_exe = make_function(der_vec_var, x_vec_var; in_place=false)
+    der_exe! = make_function(der_vec_var, x_vec_var; in_place=true)
+
     der2_exe = make_function(der2_vec_var, x_vec_var; in_place=false)
     der2_exe! = make_function(der2_vec_var, x_vec_var; in_place=true)
+
     return FastDifferentiationAllocatingSecondDerivativeExtras(
-        y_prototype, der2_exe, der2_exe!
+        y_prototype, der_exe, der_exe!, der2_exe, der2_exe!
     )
 end
 
@@ -353,6 +362,38 @@ function DI.second_derivative!(
     return der2
 end
 
+function DI.value_derivative_and_second_derivative(
+    f,
+    ::AutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationAllocatingSecondDerivativeExtras,
+)
+    y = f(x)
+    if extras.y_prototype isa Number
+        der = only(extras.der_exe(monovec(x)))
+        der2 = only(extras.der2_exe(monovec(x)))
+        return y, der, der2
+    else
+        der = reshape(extras.der_exe(monovec(x)), size(extras.y_prototype))
+        der2 = reshape(extras.der2_exe(monovec(x)), size(extras.y_prototype))
+        return y, der, der2
+    end
+end
+
+function DI.value_derivative_and_second_derivative!(
+    f,
+    der,
+    der2,
+    backend::AutoFastDifferentiation,
+    x,
+    extras::FastDifferentiationAllocatingSecondDerivativeExtras,
+)
+    y = f(x)
+    extras.der_exe!(vec(der), monovec(x))
+    extras.der2_exe!(vec(der2), monovec(x))
+    return y, der, der2
+end
+
 ## HVP
 
 struct FastDifferentiationHVPExtras{E1,E2} <: HVPExtras

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -96,6 +96,7 @@ export gradient!, gradient
 export jacobian!, jacobian
 
 export second_derivative!, second_derivative
+export value_derivative_and_second_derivative, value_derivative_and_second_derivative!
 export hvp!, hvp
 export hessian!, hessian
 

DifferentiationInterface/src/second_order/second_derivative.jl

@@ -24,6 +24,20 @@ Compute the second derivative of the function `f` at point `x`, overwriting `der
 """
 function second_derivative! end
 
+"""
+    value_derivative_and_second_derivative(f, backend, x, [extras]) -> (y, der, der2)
+
+Compute the value, first derivative and second derivative of the function `f` at point `x`.
+"""
+function value_derivative_and_second_derivative end
+
+"""
+    value_derivative_and_second_derivative!(f, der, der2, backend, x, [extras]) -> (y, der, der2)
+
+Compute the value, first derivative and second derivative of the function `f` at point `x`, overwriting `der` and `der2`.
+"""
+function value_derivative_and_second_derivative! end
+
 ## Preparation
 
 """
@@ -74,6 +88,31 @@ function second_derivative(
     return derivative(inner_derivative_closure, outer(backend), x, outer_derivative_extras)
 end
 
+function value_derivative_and_second_derivative(
+    f::F,
+    backend::AbstractADType,
+    x,
+    extras::SecondDerivativeExtras=prepare_second_derivative(f, backend, x),
+) where {F}
+    return value_derivative_and_second_derivative(
+        f, SecondOrder(backend, backend), x, extras
+    )
+end
+
+function value_derivative_and_second_derivative(
+    f::F,
+    backend::SecondOrder,
+    x,
+    extras::ClosureSecondDerivativeExtras=prepare_second_derivative(f, backend, x),
+) where {F}
+    @compat (; inner_derivative_closure, outer_derivative_extras) = extras
+    y = f(x)
+    der, der2 = value_and_derivative(
+        inner_derivative_closure, outer(backend), x, outer_derivative_extras
+    )
+    return y, der, der2
+end
+
 function second_derivative!(
     f::F,
     der2,
@@ -96,3 +135,32 @@ function second_derivative!(
         inner_derivative_closure, der2, outer(backend), x, outer_derivative_extras
     )
 end
+
+function value_derivative_and_second_derivative!(
+    f::F,
+    der,
+    der2,
+    backend::AbstractADType,
+    x,
+    extras::SecondDerivativeExtras=prepare_second_derivative(f, backend, x),
+) where {F}
+    return value_derivative_and_second_derivative!(
+        f, der, der2, SecondOrder(backend, backend), x, extras
+    )
+end
+
+function value_derivative_and_second_derivative!(
+    f::F,
+    der,
+    der2,
+    backend::SecondOrder,
+    x,
+    extras::SecondDerivativeExtras=prepare_second_derivative(f, backend, x),
+) where {F}
+    @compat (; inner_derivative_closure, outer_derivative_extras) = extras
+    y = f(x)
+    new_der, _ = value_and_derivative!(
+        inner_derivative_closure, der2, outer(backend), x, outer_derivative_extras
+    )
+    return y, copyto!(der, new_der), der2
+end

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterfaceTest/src/scenarios/default.jl

@@ -25,7 +25,11 @@ function num_to_num_scenarios_onearg(x::Number)
         PullbackScenario(num_to_num; x=x, ref=num_to_num_pullback, place=:outofplace),
         DerivativeScenario(num_to_num; x=x, ref=num_to_num_derivative, place=:outofplace),
         SecondDerivativeScenario(
-            num_to_num; x=x, ref=num_to_num_second_derivative, place=:outofplace
+            num_to_num;
+            x=x,
+            ref=num_to_num_second_derivative,
+            first_order_ref=num_to_num_derivative,
+            place=:outofplace,
         ),
     ]
 end
@@ -79,7 +83,11 @@ function num_to_arr_scenarios_onearg(x::Number, a::AbstractArray)
                     _num_to_arr(a); x=x, ref=_num_to_arr_derivative(a), place=place
                 ),
                 SecondDerivativeScenario(
-                    _num_to_arr(a); x=x, ref=_num_to_arr_second_derivative(a), place=place
+                    _num_to_arr(a);
+                    x=x,
+                    ref=_num_to_arr_second_derivative(a),
+                    first_order_ref=_num_to_arr_derivative(a),
+                    place=place,
                 ),
             ],
         )
@@ -205,8 +213,20 @@ function arr_to_num_scenarios_onearg(x::AbstractArray; linalg=true)
             [
                 PullbackScenario(arr_to_num; x=x, ref=arr_to_num_pullback, place=place),
                 GradientScenario(arr_to_num; x=x, ref=arr_to_num_gradient, place=place),
-                HVPScenario(arr_to_num; x=x, ref=arr_to_num_hvp, place=place),
-                HessianScenario(arr_to_num; x=x, ref=arr_to_num_hessian, place=place),
+                HVPScenario(
+                    arr_to_num;
+                    x=x,
+                    ref=arr_to_num_hvp,
+                    first_order_ref=arr_to_num_gradient,
+                    place=place,
+                ),
+                HessianScenario(
+                    arr_to_num;
+                    x=x,
+                    ref=arr_to_num_hessian,
+                    first_order_ref=arr_to_num_gradient,
+                    place=place,
+                ),
             ],
         )
     end