[AI] OptimizationBase: use a supplied f.grad when synthesizing fg!

lib/OptimizationBase/Project.toml

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

lib/OptimizationBase/ext/OptimizationEnzymeExt.jl

@@ -134,7 +134,11 @@ function OptimizationBase.instantiate_function(
         grad = nothing
     end
 
-    if fg == true && f.fg === nothing
+    if fg == true && f.fg === nothing && f.grad !== nothing
+        # A user-supplied gradient is authoritative; building an AD `fg!` off `f.f` would
+        # silently discard it for every value+gradient evaluation.
+        fg! = (res, θ, p = p) -> (f.grad(res, θ, p); f.f(θ, p))
+    elseif fg == true && f.fg === nothing
         function fg!(res, θ, p = p)
             Enzyme.make_zero!(res)
             y = Enzyme.autodiff(
@@ -526,7 +530,10 @@ function OptimizationBase.instantiate_function(
         grad = nothing
     end
 
-    if fg == true && f.fg === nothing
+    if fg == true && f.fg === nothing && f.grad !== nothing
+        # A user-supplied gradient is authoritative; see the in-place path above.
+        fg! = (θ, p = p) -> (f.f(θ, p), f.grad(θ, p))
+    elseif fg == true && f.fg === nothing
         res_fg = zeros(eltype(x), size(x))
         function fg!(θ, p = p)
             Enzyme.make_zero!(res_fg)

lib/OptimizationBase/ext/OptimizationZygoteExt.jl

@@ -49,11 +49,13 @@ function OptimizationBase.instantiate_function(
         grad = nothing
     end
 
-    if fg == true && f.fg === nothing
-        # Prepare a gradient prep for fg unless we already have one from the
-        # `g == true && f.grad === nothing` branch above. When the user supplied
-        # `f.grad`, line 37 didn't run, so `prep_grad` was never assigned.
-        if !(g == true && f.grad === nothing)
+    if fg == true && f.fg === nothing && f.grad !== nothing
+        # A user-supplied gradient is authoritative; building an AD `fg!` off `f.f` would
+        # silently discard it for every value+gradient evaluation.
+        fg! = (res, θ, p = p) -> (f.grad(res, θ, p); f.f(θ, p))
+    elseif fg == true && f.fg === nothing
+        # `f.grad === nothing` here, so `g == true` means the branch above already built the prep.
+        if g == false
             prep_grad = prepare_gradient(f.f, adtype, x, Constant(p), strict = Val(false))
         end
         function fg!(res, θ)
@@ -406,7 +408,10 @@ function OptimizationBase.instantiate_function(
         grad = nothing
     end
 
-    if fg == true && f.fg === nothing
+    if fg == true && f.fg === nothing && f.grad !== nothing
+        # A user-supplied gradient is authoritative; see the dense path above.
+        fg! = (res, θ, p = p) -> (f.grad(res, θ, p); f.f(θ, p))
+    elseif fg == true && f.fg === nothing
         if g == false
             extras_grad = prepare_gradient(
                 f.f, adtype.dense_ad, x, Constant(p), strict = Val(false)

lib/OptimizationBase/src/OptimizationDIExt.jl

@@ -98,14 +98,18 @@ function instantiate_function(
         nothing
     end
 
-    # Create fg! closures - need separate prep if g was false
-    fg! = if fg == true && f.fg === nothing
-        _prep_grad_fg = if g == false
-            prepare_gradient(f.f, adtype, x, Constant(p))
-        else
-            # Reuse the prep from gradient if available
-            prepare_gradient(f.f, adtype, x, Constant(p))
+    # A user-supplied `f.grad` is authoritative: building an AD `fg!` off `f.f` would silently
+    # discard it (and any tuned preparation behind it) for every value+gradient evaluation.
+    fg! = if fg == true && f.fg === nothing && f.grad !== nothing
+        let f = f, p = p
+            function (res, θ, p = p)
+                f.grad(res, θ, p)
+                return f.f(θ, p)
+            end
         end
+    elseif fg == true && f.fg === nothing
+        # Reuse the gradient prep when `grad` built one; they are the same DI operator.
+        _prep_grad_fg = g == true ? _prep_grad : prepare_gradient(f.f, adtype, x, Constant(p))
         if p !== SciMLBase.NullParameters() && p !== nothing
             let _prep_grad_fg = _prep_grad_fg, f = f, adtype = adtype
                 function (res, θ, p = p)
@@ -520,7 +524,11 @@ function instantiate_function(
     end
 
     # Create fg! closures
-    fg! = if fg == true && f.fg === nothing
+    fg! = if fg == true && f.fg === nothing && f.grad !== nothing
+        let f = f, p = p
+            (θ, p = p) -> (f.f(θ, p), f.grad(θ, p))
+        end
+    elseif fg == true && f.fg === nothing
         _prep_grad_fg = prepare_gradient(f.f, adtype, x, Constant(p))
         if p !== SciMLBase.NullParameters() && p !== nothing
             let _prep_grad_fg = _prep_grad_fg, f = f, adtype = adtype

lib/OptimizationBase/test/AD/adtests.jl

@@ -1379,3 +1379,47 @@ using MLUtils
     end
     @test G0 ≈ sum(stochgrads) rtol = 1.0e-1
 end
+
+@testset "user-supplied grad is used by the generated fg!" begin
+    # A supplied `f.grad` must drive `fg!` rather than being replaced by a fresh AD
+    # preparation: rebuilding it silently discards whatever tuning the caller did
+    # (SciML/Optimization.jl#1282).
+    ncalls = Ref(0)
+    rosen(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
+    function rosen_grad!(res, x, p = nothing)
+        ncalls[] += 1
+        res[1] = -2 * (1 - x[1]) - 400 * x[1] * (x[2] - x[1]^2)
+        res[2] = 200 * (x[2] - x[1]^2)
+        return res
+    end
+
+    Gref = zeros(2)
+    rosen_grad!(Gref, [0.5, 0.7])
+
+    adtypes = (
+        AutoForwardDiff(), AutoReverseDiff(), AutoZygote(), AutoEnzyme(),
+        AutoSparse(AutoZygote()),
+    )
+    for adtype in adtypes, g in (false, true)
+        optf = OptimizationBase.instantiate_function(
+            OptimizationFunction(rosen, adtype; grad = rosen_grad!),
+            zeros(2), adtype, nothing; g = g, fg = true
+        )
+        optf.fg === nothing && continue
+        ncalls[] = 0
+        G = zeros(2)
+        y = optf.fg(G, [0.5, 0.7])
+        @test ncalls[] == 1
+        @test y ≈ rosen([0.5, 0.7])
+        @test G ≈ Gref
+
+        # Without a user gradient the AD path still has to build its own preparation,
+        # including when `fg` is requested but `g` is not.
+        optf_ad = OptimizationBase.instantiate_function(
+            OptimizationFunction(rosen, adtype), zeros(2), adtype, nothing; g = g, fg = true
+        )
+        Gad = zeros(2)
+        @test optf_ad.fg(Gad, [0.5, 0.7]) ≈ rosen([0.5, 0.7])
+        @test Gad ≈ Gref
+    end
+end