Bring back simple hvp for now

SciML · Sep 18, 2023 · 6a54f64 · 6a54f64
1 parent 5e92e40
commit 6a54f64
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 16 deletions.
diff --git a/ext/OptimizationReverseDiffExt.jl b/ext/OptimizationReverseDiffExt.jl
@@ -54,10 +54,13 @@ function Optimization.instantiate_function(f, x, adtype::AutoReverseDiff,
 
     if f.hv === nothing
         hv = function (H, θ, v, args...)
-            _θ = ForwardDiff.Dual.(θ, v)
-            res = similar(_θ)
-            grad(res, _θ, args...)
-            H .= getindex.(ForwardDiff.partials.(res), 1)
+            # _θ = ForwardDiff.Dual.(θ, v)
+            # res = similar(_θ)
+            # grad(res, _θ, args...)
+            # H .= getindex.(ForwardDiff.partials.(res), 1)
+            res = zeros(length(θ), length(θ))
+            hess(res, θ, args...)
+            H .= res * v
         end
     else
         hv = f.hv
@@ -171,10 +174,13 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
 
     if f.hv === nothing
         hv = function (H, θ, v, args...)
-            _θ = ForwardDiff.Dual.(θ, v)
-            res = similar(_θ)
-            grad(res, _θ, args...)
-            H .= getindex.(ForwardDiff.partials.(res), 1)
+            # _θ = ForwardDiff.Dual.(θ, v)
+            # res = similar(_θ)
+            # grad(res, θ, args...)
+            # H .= getindex.(ForwardDiff.partials.(res), 1)
+            res = zeros(length(θ), length(θ))
+            hess(res, θ, args...)
+            H .= res * v
         end
     else
         hv = f.hv

diff --git a/ext/OptimizationSparseDiffExt.jl b/ext/OptimizationSparseDiffExt.jl
@@ -537,10 +537,13 @@ function Optimization.instantiate_function(f, x, adtype::AutoSparseReverseDiff,
 
     if f.hv === nothing
         hv = function (H, θ, v, args...)
-            _θ = ForwardDiff.Dual.(θ, v)
-            res = similar(_θ)
-            grad(res, _θ, args...)
-            H .= getindex.(ForwardDiff.partials.(res), 1)
+            # _θ = ForwardDiff.Dual.(θ, v)
+            # res = similar(_θ)
+            # grad(res, _θ, args...)
+            # H .= getindex.(ForwardDiff.partials.(res), 1)
+            res = zeros(length(θ), length(θ))
+            hess(res, θ, args...)
+            H .= res * v
         end
     else
         hv = f.hv
@@ -671,10 +674,13 @@ function Optimization.instantiate_function(f, cache::Optimization.ReInitCache,
 
     if f.hv === nothing
         hv = function (H, θ, v, args...)
-            _θ = ForwardDiff.Dual.(θ, v)
-            res = similar(_θ)
-            grad(res, _θ, args...)
-            H .= getindex.(ForwardDiff.partials.(res), 1)
+            # _θ = ForwardDiff.Dual.(θ, v)
+            # res = similar(_θ)
+            # grad(res, _θ, args...)
+            # H .= getindex.(ForwardDiff.partials.(res), 1)
+            res = zeros(length(θ), length(θ))
+            hess(res, θ, args...)
+            H .= res * v
         end
     else
         hv = f.hv