Apply suggestions from code review

SciML · Oct 15, 2024 · b13c812 · b13c812
1 parent 97ec67b
commit b13c812
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/Project.toml b/Project.toml
@@ -3,7 +3,6 @@ uuid = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
 authors = ["Vaibhav Dixit <vaibhavyashdixit@gmail.com> and contributors"]
 version = "2.3.0"
 
-
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

diff --git a/ext/OptimizationEnzymeExt.jl b/ext/OptimizationEnzymeExt.jl
@@ -91,13 +91,13 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
     rmode = if adtype.mode isa Nothing
         Enzyme.Reverse
     else
-        set_runtime_activity2(Enzyme.Reverse)
+        set_runtime_activity2(Enzyme.Reverse, adtype.mode)
     end
 
     fmode = if adtype.mode isa Nothing
         Enzyme.Forward
     else
-        set_runtime_activity2(Enzyme.Forward)
+        set_runtime_activity2(Enzyme.Forward. adtype.mode)
     end
 
     if g == true && f.grad === nothing
@@ -423,13 +423,13 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{false}, x
     rmode = if adtype.mode isa Nothing
         Enzyme.Reverse
     else
-        set_runtime_activity2(Enzyme.Reverse)
+        set_runtime_activity2(Enzyme.Reverse, adtype.mode)
     end
 
     fmode = if adtype.mode isa Nothing
         Enzyme.Forward
     else
-        set_runtime_activity2(Enzyme.Forward)
+        set_runtime_activity2(Enzyme.Forward, adtype.mode)
     end
 
     if g == true && f.grad === nothing

diff --git a/test/adtests.jl b/test/adtests.jl
@@ -1,6 +1,6 @@
 using OptimizationBase, Test, DifferentiationInterface, SparseArrays, Symbolics
 using ForwardDiff, Zygote, ReverseDiff, FiniteDiff, Tracker
-using ModelingToolkit, Enzyme, Random
+using Enzyme, Random
 
 x0 = zeros(2)
 rosenbrock(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
@@ -1174,8 +1174,8 @@ using MLUtils
     optf = OptimizationBase.instantiate_function(
         optf, rand(3), AutoEnzyme(), iterate(data)[1], g = true, fg = true)
     G0 = zeros(3)
-    @test_broken optf.grad(G0, ones(3), (x, y))
-    stochgrads = []
+    @test_broken optf.grad(G0, ones(3), (x0, y0))
+    # stochgrads = []
     # for (x,y) in data
     #     G = zeros(3)
     #     optf.grad(G, ones(3), (x,y))