Fix cx in Hessian

JuliaSmoothOptimizers · Dec 3, 2024 · 99dfac8 · 99dfac8
1 parent 3b6a666
commit 99dfac8
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/src/enzyme.jl b/src/enzyme.jl
@@ -449,31 +449,32 @@ end
           b.v[col] = 1
         end
 
-        function _gradient!(dx, f, x, y, obj_weight, cx)
+        function _gradient!(dx, ℓ, x, y, obj_weight, cx)
           Enzyme.make_zero!(dx)
+          dcx = make_zero(cx)
           res = Enzyme.autodiff(
             Enzyme.Reverse,
-            f,
+            ℓ,
             Enzyme.Active,
             Enzyme.Duplicated(x, dx),
             Enzyme.Const(y),
             Enzyme.Const(obj_weight),
-            Enzyme.Const(cx)
+            Enzyme.Duplicated(cx, dcx)
           )
           return nothing
         end
 
-        function _hvp!(res, f, x, v, y, obj_weight, cx)
-          # grad = Enzyme.make_zero(x)
+        function _hvp!(res, ℓ, x, v, y, obj_weight, cx)
+          dcx = make_zero(cx)
           Enzyme.autodiff(
               Enzyme.Forward,
               _gradient!,
               res,
-              Enzyme.Const(f),
+              Enzyme.Const(ℓ),
               Enzyme.Duplicated(x, v),
               Enzyme.Const(y),
               Enzyme.Const(obj_weight),
-              Enzyme.Const(cx),
+              Enzyme.Duplicated(cx, dcx),
           )
           return nothing
         end