Stabilize Enzyme Lagrangian Hessians on Julia 1.12

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://chatgpt.com/codex/tasks/01a03a17-ad6f-7131-82fc-d0fd57ea6512

Author: ChrisRackauckas

lib/OptimizationBase/ext/OptimizationEnzymeExt.jl

@@ -80,11 +80,9 @@ function lagrangian(x, _f::Function, cons::Function, p, λ, σ = one(eltype(x)))
     return σ * _f(x, p) + dot(λ, res)
 end
 
-function lag_grad(mode, x, dx, lagrangian::Function, _f::Function, cons::Function, p, σ, λ)
-    Enzyme.autodiff_deferred(
-        mode, Const(lagrangian), Active, Enzyme.Duplicated(x, dx),
-        Const(_f), Const(cons), Const(p), Const(λ), Const(σ)
-    )
+function lag_grad(mode, x, dx, f)
+    Enzyme.make_zero!(dx)
+    Enzyme.autodiff(mode, Const(f), Active, Enzyme.Duplicated(x, dx))
     return nothing
 end
 
@@ -457,6 +455,7 @@ function OptimizationBase.instantiate_function(
         end
 
         function fill_lag_hessian!(θ, σ, μ, p)
+            lag = x -> lagrangian(x, f.f, f.cons, p, μ, σ)
             for i in eachindex(θ)
                 Enzyme.make_zero!(lag_bθ)
                 Enzyme.make_zero!(lag_vdbθ[i])
@@ -465,13 +464,8 @@ function OptimizationBase.instantiate_function(
                     lag_grad,
                     Const(rmode),
                     Enzyme.Duplicated(θ, lag_vdθ[i]),
-                    Enzyme.DuplicatedNoNeed(lag_bθ, lag_vdbθ[i]),
-                    Const(lagrangian),
-                    Const(f.f),
-                    Const(f.cons),
-                    Const(p),
-                    Const(σ),
-                    Const(μ)
+                    Enzyme.Duplicated(lag_bθ, lag_vdbθ[i]),
+                    Const(lag)
                 )
             end
             return nothing
@@ -494,7 +488,9 @@ function OptimizationBase.instantiate_function(
             fill_lag_hessian!(θ, σ, μ, p)
 
             for i in eachindex(θ)
-                H[i, :] .= lag_vdbθ[i]
+                vec_lagv = lag_vdbθ[i]
+                H[i, 1:i] .= @view(vec_lagv[1:i])
+                H[1:i, i] .= @view(vec_lagv[1:i])
             end
             return
         end

lib/OptimizationBase/test/AD/enzyme_lagrangian_hessian.jl

@@ -46,5 +46,7 @@ function check_lagrangian_hessian(N)
 end
 
 @testset "Enzyme Lagrangian Hessian" begin
-    foreach(check_lagrangian_hessian, (1:10..., 20, 40, 60))
+    @testset "N = $N" for N in (1:10..., 20, 40, 60)
+        check_lagrangian_hessian(N)
+    end
 end