Methods with and without extras (#313)

* Methods with and without extras * Typos * Fix * SecondOrder for hvp * Polyester * Typo * Typo * Add correctness tests * Extras type * Typo * Other typo * Typo * Hessian doc
JuliaDiff · Jun 7, 2024 · 61c2741 · 61c2741
1 parent 0b7a2f9
commit 61c2741
Show file tree

Hide file tree

Showing 11 changed files with 772 additions and 765 deletions.
diff --git a/...iationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl b/...iationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl
@@ -11,7 +11,8 @@ using DifferentiationInterface:
     NoGradientExtras,
     NoHessianExtras,
     NoJacobianExtras,
-    PushforwardExtras
+    PushforwardExtras,
+    PushforwardDerivativeExtras
 using DocStringExtensions
 using LinearAlgebra: mul!
 using PolyesterForwardDiff: threaded_gradient!, threaded_jacobian!

diff --git a/DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl
@@ -31,28 +31,32 @@ end
 
 ## Derivative
 
-function DI.prepare_derivative(f, backend::AutoPolyesterForwardDiff, x)
+function DI.prepare_derivative(
+    f, backend::AutoPolyesterForwardDiff, x
+)::PushforwardDerivativeExtras
     return DI.prepare_derivative(f, single_threaded(backend), x)
 end
 
 function DI.value_and_derivative(
-    f, backend::AutoPolyesterForwardDiff, x, extras::DerivativeExtras
+    f, backend::AutoPolyesterForwardDiff, x, extras::PushforwardDerivativeExtras
 )
     return DI.value_and_derivative(f, single_threaded(backend), x, extras)
 end
 
 function DI.value_and_derivative!(
-    f, der, backend::AutoPolyesterForwardDiff, x, extras::DerivativeExtras
+    f, der, backend::AutoPolyesterForwardDiff, x, extras::PushforwardDerivativeExtras
 )
     return DI.value_and_derivative!(f, der, single_threaded(backend), x, extras)
 end
 
-function DI.derivative(f, backend::AutoPolyesterForwardDiff, x, extras::DerivativeExtras)
+function DI.derivative(
+    f, backend::AutoPolyesterForwardDiff, x, extras::PushforwardDerivativeExtras
+)
     return DI.derivative(f, single_threaded(backend), x, extras)
 end
 
 function DI.derivative!(
-    f, der, backend::AutoPolyesterForwardDiff, x, extras::DerivativeExtras
+    f, der, backend::AutoPolyesterForwardDiff, x, extras::PushforwardDerivativeExtras
 )
     return DI.derivative!(f, der, single_threaded(backend), x, extras)
 end

diff --git a/DifferentiationInterface/src/first_order/derivative.jl b/DifferentiationInterface/src/first_order/derivative.jl
@@ -72,84 +72,86 @@ end
 
 ## One argument
 
+function value_and_derivative(f::F, backend::AbstractADType, x) where {F}
+    return value_and_derivative(f, backend, x, prepare_derivative(f, backend, x))
+end
+
+function value_and_derivative!(f::F, der, backend::AbstractADType, x) where {F}
+    return value_and_derivative!(f, der, backend, x, prepare_derivative(f, backend, x))
+end
+
+function derivative(f::F, backend::AbstractADType, x) where {F}
+    return derivative(f, backend, x, prepare_derivative(f, backend, x))
+end
+
+function derivative!(f::F, der, backend::AbstractADType, x) where {F}
+    return derivative!(f, der, backend, x, prepare_derivative(f, backend, x))
+end
+
 function value_and_derivative(
-    f::F,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f, backend, x),
+    f::F, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return value_and_pushforward(f, backend, x, one(x), extras.pushforward_extras)
 end
 
 function value_and_derivative!(
-    f::F,
-    der,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f, backend, x),
+    f::F, der, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return value_and_pushforward!(f, der, backend, x, one(x), extras.pushforward_extras)
 end
 
 function derivative(
-    f::F,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f, backend, x),
+    f::F, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return pushforward(f, backend, x, one(x), extras.pushforward_extras)
 end
 
 function derivative!(
-    f::F,
-    der,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f, backend, x),
+    f::F, der, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return pushforward!(f, der, backend, x, one(x), extras.pushforward_extras)
 end
 
 ## Two arguments
 
+function value_and_derivative(f!::F, y, backend::AbstractADType, x) where {F}
+    return value_and_derivative(f!, y, backend, x, prepare_derivative(f!, y, backend, x))
+end
+
+function value_and_derivative!(f!::F, y, der, backend::AbstractADType, x) where {F}
+    return value_and_derivative!(
+        f!, y, der, backend, x, prepare_derivative(f!, y, backend, x)
+    )
+end
+
+function derivative(f!::F, y, backend::AbstractADType, x) where {F}
+    return derivative(f!, y, backend, x, prepare_derivative(f!, y, backend, x))
+end
+
+function derivative!(f!::F, y, der, backend::AbstractADType, x) where {F}
+    return derivative!(f!, y, der, backend, x, prepare_derivative(f!, y, backend, x))
+end
+
 function value_and_derivative(
-    f!::F,
-    y,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
+    f!::F, y, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return value_and_pushforward(f!, y, backend, x, one(x), extras.pushforward_extras)
 end
 
 function value_and_derivative!(
-    f!::F,
-    y,
-    der,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
+    f!::F, y, der, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return value_and_pushforward!(f!, y, der, backend, x, one(x), extras.pushforward_extras)
 end
 
 function derivative(
-    f!::F,
-    y,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
+    f!::F, y, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return pushforward(f!, y, backend, x, one(x), extras.pushforward_extras)
 end
 
 function derivative!(
-    f!::F,
-    y,
-    der,
-    backend::AbstractADType,
-    x,
-    extras::DerivativeExtras=prepare_derivative(f!, y, backend, x),
+    f!::F, y, der, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
     return pushforward!(f!, y, der, backend, x, one(x), extras.pushforward_extras)
 end
diff --git a/DifferentiationInterface/src/first_order/gradient.jl b/DifferentiationInterface/src/first_order/gradient.jl
@@ -62,34 +62,42 @@ end
 
 ## One argument
 
+function value_and_gradient(f::F, backend::AbstractADType, x) where {F}
+    return value_and_gradient(f, backend, x, prepare_gradient(f, backend, x))
+end
+
+function value_and_gradient!(f::F, der, backend::AbstractADType, x) where {F}
+    return value_and_gradient!(f, der, backend, x, prepare_gradient(f, backend, x))
+end
+
+function gradient(f::F, backend::AbstractADType, x) where {F}
+    return gradient(f, backend, x, prepare_gradient(f, backend, x))
+end
+
+function gradient!(f::F, der, backend::AbstractADType, x) where {F}
+    return gradient!(f, der, backend, x, prepare_gradient(f, backend, x))
+end
+
 function value_and_gradient(
-    f::F, backend::AbstractADType, x, extras::GradientExtras=prepare_gradient(f, backend, x)
+    f::F, backend::AbstractADType, x, extras::PullbackGradientExtras
 ) where {F}
     return value_and_pullback(f, backend, x, one(eltype(x)), extras.pullback_extras)
 end
 
 function value_and_gradient!(
-    f::F,
-    grad,
-    backend::AbstractADType,
-    x,
-    extras::GradientExtras=prepare_gradient(f, backend, x),
+    f::F, grad, backend::AbstractADType, x, extras::PullbackGradientExtras
 ) where {F}
     return value_and_pullback!(f, grad, backend, x, one(eltype(x)), extras.pullback_extras)
 end
 
 function gradient(
-    f::F, backend::AbstractADType, x, extras::GradientExtras=prepare_gradient(f, backend, x)
+    f::F, backend::AbstractADType, x, extras::PullbackGradientExtras
 ) where {F}
     return pullback(f, backend, x, one(eltype(x)), extras.pullback_extras)
 end
 
 function gradient!(
-    f::F,
-    grad,
-    backend::AbstractADType,
-    x,
-    extras::GradientExtras=prepare_gradient(f, backend, x),
+    f::F, grad, backend::AbstractADType, x, extras::PullbackGradientExtras
 ) where {F}
     return pullback!(f, grad, backend, x, one(eltype(x)), extras.pullback_extras)
 end