test: update tests for enzyme (#1552)

* test: update tests for enzyme

* Update src/helpers/training.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

src/helpers/training.jl

@@ -358,7 +358,7 @@ function maybe_wrap_adtype(
                          Enzyme.jl (`AutoEnzyme`)."))
 end
 
-function generate_wrappers(::F, m, ps, st, data, ::False) where {F}
+function generate_wrappers(::F, m, ps, st, data, ::False, ::StaticBool) where {F}
     @warn "Detected function wrapper generation with function being updated between calls. \
            This will generate type-unstable code. A possible reason for this is \
            `TrainState` was compiled (first call to `compute_gradients`) with function \
@@ -367,16 +367,31 @@ function generate_wrappers(::F, m, ps, st, data, ::False) where {F}
     return Ref{Any}(), Ref{NamedTuple}()
 end
 
+function generate_wrappers(
+    objective_function::F, m, ps, st, data, ::True, ::False
+) where {F}
+    _, stₙ, statsₙ = objective_function(m, ps, st, data)
+    return Ref{typeof(stₙ)}(stₙ), Ref{NamedTuple}()  # State type is not preserved
+end
+
 # Run the code when trying to compile the function for the first time.
-function generate_wrappers(objective_function::F, m, ps, st, data, ::True) where {F}
+function generate_wrappers(objective_function::F, m, ps, st, data, ::True, ::True) where {F}
     _, stₙ, statsₙ = objective_function(m, ps, st, data)
     return Ref{typeof(stₙ)}(stₙ), Ref{typeof(statsₙ)}(statsₙ)
 end
 
 function wrap_objective_function(
     objective_function::F, m, ps, st, data, first_try::StaticBool
 ) where {F}
-    st_updated, stats = generate_wrappers(objective_function, m, ps, st, data, first_try)
+    st_updated, stats = generate_wrappers(
+        objective_function,
+        m,
+        ps,
+        st,
+        data,
+        first_try,
+        static(LuxCore.preserves_state_type(m)),
+    )
 
     wrapped_objective_function = @closure (model, ps, st, data) -> begin
         loss, st_, stats_ = objective_function(model, ps, st, data)

test/helpers/training_tests.jl

@@ -260,9 +260,7 @@ end
 
     tstate = Training.TrainState(model, ps, st, opt)
 
-    _, _, _, tstate_new = @inferred Training.compute_gradients(
-        AutoEnzyme(), mse, (x, x), tstate
-    )
+    _, _, _, tstate_new = Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate)
 
     @test tstate_new.states !== tstate.states
 
@@ -271,14 +269,12 @@ end
 
     tstate = Training.TrainState(model, ps, st, opt)
 
-    _, _, _, tstate_new = @inferred Training.compute_gradients(
-        AutoEnzyme(), mse, (x, x), tstate
-    )
+    _, _, _, tstate_new = Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate)
 
     @test @inferred(Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate_new)) isa
         Any
 
-    _, _, _, tstate_new2 = @inferred Training.compute_gradients(
+    _, _, _, tstate_new2 = Training.compute_gradients(
         AutoEnzyme(), mse2, (x, x), tstate_new
     )
     @test hasfield(typeof(tstate_new2.cache.extras), :forward)
@@ -293,9 +289,7 @@ end
 
     tstate = Training.TrainState(model, ps, st, opt)
 
-    _, _, _, tstate_new = @inferred Training.compute_gradients(
-        AutoEnzyme(), mse, (x, x), tstate
-    )
+    _, _, _, tstate_new = Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate)
 
     @test tstate_new.states !== tstate.states
 
@@ -304,14 +298,12 @@ end
 
     tstate = Training.TrainState(model, ps, st, opt)
 
-    _, _, _, tstate_new = @inferred Training.compute_gradients(
-        AutoEnzyme(), mse, (x, x), tstate
-    )
+    _, _, _, tstate_new = Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate)
 
     @test @inferred(Training.compute_gradients(AutoEnzyme(), mse, (x, x), tstate_new)) isa
         Any
 
-    _, _, _, tstate_new2 = @inferred Training.compute_gradients(
+    _, _, _, tstate_new2 = Training.compute_gradients(
         AutoEnzyme(), mse2, (x, x), tstate_new
     )
     @test hasfield(typeof(tstate_new2.cache.extras), :forward)