OptimizationBase: fix Enzyme out-of-place grad/fg wiring (SciML#1284)

Two pre-existing bugs in the `OptimizationFunction{false}` method of
`OptimizationEnzymeExt.instantiate_function`:

- The fallback that forwards a user-supplied gradient was gated on `fg == true`
  rather than `g == true`, so `g = true, fg = false` with a supplied `f.grad`
  produced `grad === nothing`. Every other extension, and the in-place Enzyme
  method, gate on `g`.
- The generated `fg!` differentiated into `res_fg` but returned `res`, the
  buffer belonging to the `grad` branch. With `g = false` that name is never
  assigned and `fg!` throws `UndefVarError: res not defined`; with `g = true`
  it silently returns the untouched gradient buffer — `[0.0, 0.0]` on a
  first call.



Claude-Session: https://claude.ai/code/session_014Uj4Pu1LKHYqLSTLfPYDPo

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: 3 people

lib/OptimizationBase/Project.toml

@@ -1,6 +1,6 @@
 name = "OptimizationBase"
 uuid = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
-version = "5.2.3"
+version = "5.2.4"
 authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com> and contributors"]
 
 [deps]

lib/OptimizationBase/ext/OptimizationEnzymeExt.jl

@@ -524,7 +524,7 @@ function OptimizationBase.instantiate_function(
             )
             return res
         end
-    elseif fg == true
+    elseif g == true
         grad = (θ, p = p) -> f.grad(θ, p)
     else
         grad = nothing
@@ -545,7 +545,7 @@ function OptimizationBase.instantiate_function(
                 Enzyme.Duplicated(θ, res_fg),
                 Const(p)
             )[2]
-            return y, res
+            return y, res_fg
         end
     elseif fg == true
         fg! = (θ, p = p) -> f.fg(θ, p)

lib/OptimizationBase/test/AD/adtests.jl

@@ -1423,3 +1423,32 @@ end
         @test Gad ≈ Gref
     end
 end
+
+@testset "Enzyme out-of-place grad/fg wiring" begin
+    rosen(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
+    rosen_grad(x, p = nothing) = [
+        -2 * (1 - x[1]) - 400 * x[1] * (x[2] - x[1]^2),
+        200 * (x[2] - x[1]^2),
+    ]
+    ad = AutoEnzyme()
+    z = [0.5, 0.7]
+    gref = rosen_grad(z)
+
+    # `g` alone must honour a supplied gradient; it used to be gated on `fg`.
+    optf = OptimizationBase.instantiate_function(
+        OptimizationFunction{false}(rosen, ad; grad = rosen_grad),
+        zeros(2), ad, nothing; g = true, fg = false
+    )
+    @test optf.grad !== nothing
+    @test optf.grad(z) ≈ gref
+
+    # The AD `fg!` must return the buffer it differentiated into, with or without `g`.
+    for g in (false, true)
+        optf = OptimizationBase.instantiate_function(
+            OptimizationFunction{false}(rosen, ad), zeros(2), ad, nothing; g = g, fg = true
+        )
+        y, G = optf.fg(z)
+        @test y ≈ rosen(z)
+        @test G ≈ gref
+    end
+end