mode in stoch ad tests for enzyme

Author: Vaibhavdixit02

ext/OptimizationEnzymeExt.jl

@@ -98,7 +98,7 @@ function OptimizationBase.instantiate_function(f::OptimizationFunction{true}, x,
     fmode = if adtype.mode isa Nothing
         Enzyme.Forward
     else
-        set_runtime_activity2(Enzyme.Forward.adtype.mode)
+        set_runtime_activity2(Enzyme.Forward, adtype.mode)
     end
 
     if g == true && f.grad === nothing

test/adtests.jl

@@ -1172,17 +1172,18 @@ using MLUtils
 
     optf = OptimizationFunction(loss, AutoEnzyme())
     optf = OptimizationBase.instantiate_function(
-        optf, rand(3), AutoEnzyme(), iterate(data)[1], g = true, fg = true)
+        optf, rand(3), AutoEnzyme(mode = set_runtime_activity(Reverse)),
+        iterate(data)[1], g = true, fg = true)
     G0 = zeros(3)
-    @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))
-    #     push!(stochgrads, copy(G))
-    #     G1 = zeros(3)
-    #     optf.fg(G1, ones(3), (x,y))
-    #     @test G ≈ G1 rtol=1e-6
-    # end
-    # @test G0 ≈ sum(stochgrads)/length(stochgrads) rtol=1e-1
+    optf.grad(G0, ones(3), (x0, y0))
+    stochgrads = []
+    for (x, y) in data
+        G = zeros(3)
+        optf.grad(G, ones(3), (x, y))
+        push!(stochgrads, copy(G))
+        G1 = zeros(3)
+        optf.fg(G1, ones(3), (x, y))
+        @test G≈G1 rtol=1e-6
+    end
+    @test G0≈sum(stochgrads) rtol=1e-1
 end