Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add dual number–based second derivatives for ForwardDiff #310

Merged
merged 16 commits into from
Jun 7, 2024
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ using DifferentiationInterface:
HessianExtras,
JacobianExtras,
NoDerivativeExtras,
NoSecondDerivativeExtras,
PushforwardExtras
using ForwardDiff.DiffResults: DiffResults, DiffResult, GradientResult, MutableDiffResult
using ForwardDiff:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,58 @@ function DI.pushforward!(
return dy
end

## Second derivative

function DI.prepare_second_derivative(f::F, backend::AutoForwardDiff, x) where {F}
return NoSecondDerivativeExtras()
end

function DI.second_derivative(
f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
) where {F}
T = tag_type(f, backend, x)
xdual = make_dual(T, x, one(x))
T2 = tag_type(f, backend, xdual)
ydual = f(make_dual(T2, xdual, one(xdual)))
return myderivative(T, myderivative(T2, ydual))
end

function DI.second_derivative!(
f::F, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
) where {F}
T = tag_type(f, backend, x)
xdual = make_dual(T, x, one(x))
T2 = tag_type(f, backend, xdual)
ydual = f(make_dual(T2, xdual, one(xdual)))
return myderivative!(T, der2, myderivative(T2, ydual))
end

function DI.value_derivative_and_second_derivative(
f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
) where {F}
T = tag_type(f, backend, x)
xdual = make_dual(T, x, one(x))
T2 = tag_type(f, backend, xdual)
ydual = f(make_dual(T2, xdual, one(xdual)))
y = myvalue(T, myvalue(T2, ydual))
der = myderivative(T, myvalue(T2, ydual))
der2 = myderivative(T, myderivative(T2, ydual))
return y, der, der2
end

function DI.value_derivative_and_second_derivative!(
f::F, der, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
) where {F}
T = tag_type(f, backend, x)
xdual = make_dual(T, x, one(x))
T2 = tag_type(f, backend, xdual)
ydual = f(make_dual(T2, xdual, one(xdual)))
y = myvalue(T, myvalue(T2, ydual))
myderivative!(T, der, myvalue(T2, ydual))
myderivative!(T, der2, myderivative(T2, ydual))
return y, der, der2
end

## Gradient

struct ForwardDiffGradientExtras{C} <: GradientExtras
Expand Down
Loading