From 2e3ce047024b02a36b6e5d2761419cec5474d045 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 4 Dec 2024 15:22:07 +0530
Subject: [PATCH 1/2] test: try fixing enzyme test

---
 test/layers/basic_tests.jl | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/test/layers/basic_tests.jl b/test/layers/basic_tests.jl
index 3adea7323d..0d8926764f 100644
--- a/test/layers/basic_tests.jl
+++ b/test/layers/basic_tests.jl
@@ -165,12 +165,15 @@ end
         )
         x = randn(SVector{N, Float64})
 
-        fun = let d = d, x = x
-            ps -> sum(d(x, ps, (;))[1])
+        grad1 = ForwardDiff.gradient(ComponentArray(ps)) do ps
+            sumabs2first(d, x, ps, (;))
         end
-        grad1 = ForwardDiff.gradient(fun, ComponentVector(ps))
-        grad2 = Enzyme.gradient(Enzyme.Reverse, fun, ps)[1]
-        @test maximum(abs, grad1 .- ComponentVector(grad2)) < 1e-6
+
+        grad2 = Enzyme.gradient(
+            Enzyme.Reverse, sumabs2first, Const(d), Const(x), ps, Const((;))
+        )[3]
+
+        @test maximum(abs, grad1 .- ComponentArray(grad2)) < 1e-6
     end
 end
 

From 47b9d11a49c7ce20fe5ee7e4557a2e399acbfd09 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Wed, 4 Dec 2024 18:29:00 +0530
Subject: [PATCH 2/2] test: temporarily mark the test as broken

---
 test/enzyme_tests.jl       |  8 ++------
 test/layers/basic_tests.jl | 18 +++++++++++-------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/test/enzyme_tests.jl b/test/enzyme_tests.jl
index 8fbf085fc5..7141e7d098 100644
--- a/test/enzyme_tests.jl
+++ b/test/enzyme_tests.jl
@@ -83,12 +83,10 @@ export generic_loss_function, compute_enzyme_gradient, compute_zygote_gradient,
 end
 
 @testitem "Enzyme Integration" setup=[EnzymeTestSetup, SharedTestSetup] tags=[
-    :autodiff, :enzyme] begin
+    :autodiff, :enzyme] timeout=3600 begin
     rng = StableRNG(12345)
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
-        # TODO: Currently all the tests are run on CPU. We should eventually add tests for
-        #       CUDA and AMDGPU.
         ongpu && continue
 
         @testset "[$(i)] $(nameof(typeof(model)))" for (i, (model, x)) in enumerate(MODELS_LIST)
@@ -106,15 +104,13 @@ end
     end
 end
 
-@testitem "Enzyme Integration ComponentArray" setup=[EnzymeTestSetup, SharedTestSetup] tags=[
+@testitem "Enzyme Integration ComponentArray" setup=[EnzymeTestSetup, SharedTestSetup] timeout=3600 tags=[
     :autodiff, :enzyme] begin
     using ComponentArrays
 
     rng = StableRNG(12345)
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
-        # TODO: Currently all the tests are run on CPU. We should eventually add tests for
-        #       CUDA and AMDGPU.
         ongpu && continue
 
         @testset "[$(i)] $(nameof(typeof(model)))" for (i, (model, x)) in enumerate(MODELS_LIST)
diff --git a/test/layers/basic_tests.jl b/test/layers/basic_tests.jl
index 0d8926764f..1e36b9899b 100644
--- a/test/layers/basic_tests.jl
+++ b/test/layers/basic_tests.jl
@@ -165,15 +165,19 @@ end
         )
         x = randn(SVector{N, Float64})
 
-        grad1 = ForwardDiff.gradient(ComponentArray(ps)) do ps
-            sumabs2first(d, x, ps, (;))
-        end
+        broken = pkgversion(Enzyme) ≥ v"0.13.18"
+
+        @test begin
+            grad1 = ForwardDiff.gradient(ComponentArray(ps)) do ps
+                sumabs2first(d, x, ps, (;))
+            end
 
-        grad2 = Enzyme.gradient(
-            Enzyme.Reverse, sumabs2first, Const(d), Const(x), ps, Const((;))
-        )[3]
+            grad2 = Enzyme.gradient(
+                Enzyme.Reverse, sumabs2first, Const(d), Const(x), ps, Const((;))
+            )[3]
 
-        @test maximum(abs, grad1 .- ComponentArray(grad2)) < 1e-6
+            maximum(abs, grad1 .- ComponentArray(grad2)) < 1e-6
+        end broken=broken
     end
 end