fix: use PolyesterForwardDiff's inner gradient in its HVP (#610)

* fix: use PolyesterForwardDiff's inner gradient in its HVP * SecondOrder
JuliaDiff · Nov 5, 2024 · fb16fab · fb16fab · gdalle · Nov 5, 2024
1 parent 3b978b8
commit fb16fab
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 6 deletions.
diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml
@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.19"
+version = "0.6.20"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/...iationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl b/...iationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl
@@ -15,6 +15,7 @@ using DifferentiationInterface:
     NoJacobianPrep,
     PushforwardPrep,
     SecondDerivativePrep,
+    SecondOrder,
     unwrap,
     with_contexts
 using LinearAlgebra: mul!

diff --git a/DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl
@@ -283,7 +283,9 @@ end
 function DI.prepare_hvp(
     f, backend::AutoPolyesterForwardDiff, x, tx::NTuple, contexts::Vararg{Context,C}
 ) where {C}
-    return DI.prepare_hvp(f, single_threaded(backend), x, tx, contexts...)
+    return DI.prepare_hvp(
+        f, SecondOrder(single_threaded(backend), backend), x, tx, contexts...
+    )
 end
 
 function DI.hvp(
@@ -294,7 +296,9 @@ function DI.hvp(
     tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {C}
-    return DI.hvp(f, prep, single_threaded(backend), x, tx, contexts...)
+    return DI.hvp(
+        f, prep, SecondOrder(single_threaded(backend), backend), x, tx, contexts...
+    )
 end
 
 function DI.hvp!(
@@ -306,7 +310,9 @@ function DI.hvp!(
     tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {C}
-    return DI.hvp!(f, tg, prep, single_threaded(backend), x, tx, contexts...)
+    return DI.hvp!(
+        f, tg, prep, SecondOrder(single_threaded(backend), backend), x, tx, contexts...
+    )
 end
 
 function DI.gradient_and_hvp(
@@ -317,7 +323,9 @@ function DI.gradient_and_hvp(
     tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {C}
-    return DI.gradient_and_hvp(f, prep, single_threaded(backend), x, tx, contexts...)
+    return DI.gradient_and_hvp(
+        f, prep, SecondOrder(single_threaded(backend), backend), x, tx, contexts...
+    )
 end
 
 function DI.gradient_and_hvp!(
@@ -331,7 +339,14 @@ function DI.gradient_and_hvp!(
     contexts::Vararg{Context,C},
 ) where {C}
     return DI.gradient_and_hvp!(
-        f, grad, tg, prep, single_threaded(backend), x, tx, contexts...
+        f,
+        grad,
+        tg,
+        prep,
+        SecondOrder(single_threaded(backend), backend),
+        x,
+        tx,
+        contexts...,
     )
 end