Apply suggestions from code review

Co-authored-by: Vaibhav Kumar Dixit <[email protected]>

Author: wsmoses

ext/OptimizationEnzymeExt.jl

@@ -19,6 +19,7 @@ function Optimization.instantiate_function(f::OptimizationFunction{true}, x,
             res .= zero(eltype(res))
             Enzyme.autodiff(Enzyme.Reverse,
                 Const(firstapply),
+                Active,
                 Const(f.f),
                 Enzyme.Duplicated(θ, res),
                 Const(p),
@@ -32,6 +33,7 @@ function Optimization.instantiate_function(f::OptimizationFunction{true}, x,
         function g(θ, bθ, f, p, args...)
             Enzyme.autodiff_deferred(Enzyme.Reverse,
                 Const(firstapply),
+                Active,
                 Const(f),
                 Enzyme.Duplicated(θ, bθ),
                 Const(p),
@@ -65,6 +67,7 @@ function Optimization.instantiate_function(f::OptimizationFunction{true}, x,
             dx = zeros(length(x))
             Enzyme.autodiff_deferred(Enzyme.Reverse,
                 firstapply,
+                Active,
                 f,
                 Enzyme.Duplicated(x, dx),
                 Const(p),
@@ -155,6 +158,7 @@ function Optimization.instantiate_function(f::OptimizationFunction{true},
             res .= zero(eltype(res))
             Enzyme.autodiff(Enzyme.Reverse,
                 Const(firstapply),
+                Active,
                 Const(f.f),
                 Enzyme.Duplicated(θ, res),
                 Const(p),
@@ -197,7 +201,7 @@ function Optimization.instantiate_function(f::OptimizationFunction{true},
     if f.hv === nothing
         function f2(x, f, p, args...)
             dx = zeros(length(x))
-            Enzyme.autodiff_deferred(Enzyme.Reverse, firstapply,
+            Enzyme.autodiff_deferred(Enzyme.Reverse, firstapply, Active,
                 f,
                 Enzyme.Duplicated(x, dx),
                 Const(p),