Avoid wide Enzyme batches for Lagrangian Hessians

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

@@ -456,23 +456,29 @@ function OptimizationBase.instantiate_function(
             lag_vdbθ = Tuple((copy(r) for r in eachrow(f.hess_prototype)))
         end
 
-        function lag_h!(h, θ, σ, μ, p = p)
-            Enzyme.make_zero!(lag_bθ)
-            Enzyme.make_zero!.(lag_vdbθ)
+        function fill_lag_hessian!(θ, σ, μ, p)
+            for i in eachindex(θ)
+                Enzyme.make_zero!(lag_bθ)
+                Enzyme.make_zero!(lag_vdbθ[i])
+                Enzyme.autodiff(
+                    fmode,
+                    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(μ)
+                )
+            end
+            return nothing
+        end
 
-            Enzyme.autodiff(
-                fmode,
-                lag_grad,
-                Const(rmode),
-                Enzyme.BatchDuplicated(θ, lag_vdθ),
-                Enzyme.BatchDuplicatedNoNeed(lag_bθ, lag_vdbθ),
-                Const(lagrangian),
-                Const(f.f),
-                Const(f.cons),
-                Const(p),
-                Const(σ),
-                Const(μ)
-            )
+        function lag_h!(h, θ, σ, μ, p = p)
+            fill_lag_hessian!(θ, σ, μ, p)
             k = 0
 
             for i in eachindex(θ)
@@ -485,22 +491,7 @@ function OptimizationBase.instantiate_function(
 
         function lag_h!(H::AbstractMatrix, θ, σ, μ, p = p)
             Enzyme.make_zero!(H)
-            Enzyme.make_zero!(lag_bθ)
-            Enzyme.make_zero!.(lag_vdbθ)
-
-            Enzyme.autodiff(
-                fmode,
-                lag_grad,
-                Const(rmode),
-                Enzyme.BatchDuplicated(θ, lag_vdθ),
-                Enzyme.BatchDuplicatedNoNeed(lag_bθ, lag_vdbθ),
-                Const(lagrangian),
-                Const(f.f),
-                Const(f.cons),
-                Const(p),
-                Const(σ),
-                Const(μ)
-            )
+            fill_lag_hessian!(θ, σ, μ, p)
 
             for i in eachindex(θ)
                 H[i, :] .= lag_vdbθ[i]

lib/OptimizationBase/test/AD/enzyme_lagrangian_hessian.jl

@@ -0,0 +1,50 @@
+using ChainRulesCore, Enzyme, OptimizationBase, Test
+
+function check_lagrangian_hessian(N)
+    h = 1 / N
+    alpha = 350
+    x_offset = N + 1
+    u_offset = 2(N + 1)
+    function objective(x, p)
+        return sum(
+            0.5 * h * (x[u_offset + i + 1]^2 + x[u_offset + i]^2) +
+                0.5 * alpha * h * (cos(x[i + 1]) + cos(x[i])) for i in 1:N
+        ) + x[1] * x[2]
+    end
+    function constraint!(res, x, p)
+        for i in 1:N
+            res[i] = x[x_offset + i + 1] - x[x_offset + i] -
+                0.5 * h * (sin(x[i + 1]) + sin(x[i]))
+            res[N + i] = x[i + 1] - x[i] -
+                0.5 * h * (x[u_offset + i + 1] + x[u_offset + i])
+        end
+        return nothing
+    end
+
+    x = zeros(3(N + 1))
+    f = OptimizationFunction(objective, AutoEnzyme(); cons = constraint!)
+    instantiated = OptimizationBase.instantiate_function(
+        f, x, AutoEnzyme(), nothing, 2N; lag_h = true
+    )
+    multipliers = ones(2N)
+    expected = zeros(length(x), length(x))
+    for i in 1:(N + 1)
+        expected[i, i] = (i == 1 || i == N + 1) ? -0.5alpha * h : -alpha * h
+        expected[u_offset + i, u_offset + i] =
+            (i == 1 || i == N + 1) ? h : 2h
+    end
+    expected[1, 2] = expected[2, 1] = 1
+
+    packed = zeros(length(x) * (length(x) + 1) ÷ 2)
+    instantiated.lag_h(packed, x, 1.0, multipliers)
+    @test packed ≈ [expected[i, j] for i in axes(expected, 1) for j in 1:i]
+
+    dense = zeros(length(x), length(x))
+    instantiated.lag_h(dense, x, 1.0, multipliers)
+    @test dense ≈ expected
+    return nothing
+end
+
+@testset "Enzyme Lagrangian Hessian" begin
+    foreach(check_lagrangian_hessian, (1:10..., 20, 40, 60))
+end

lib/OptimizationBase/test/AD/tests.jl

@@ -1,6 +1,7 @@
 using Test
 
 @testset "OptimizationBase AD" begin
+    include("enzyme_lagrangian_hessian.jl")
     include("adtests.jl")
     include("dual_tolerant_tests.jl")
     include("cvxtest.jl")