test: more tests fixed

LuxDL · Dec 31, 2024 · 4dc08d1 · 4dc08d1
1 parent 9826a2d
commit 4dc08d1
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 21 deletions.
diff --git a/Project.toml b/Project.toml
@@ -89,7 +89,7 @@ ComponentArrays = "0.15.18"
 ConcreteStructs = "0.2.3"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 FastClosures = "0.3.2"
 Flux = "0.15, 0.16"
 ForwardDiff = "0.10.36"

diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -66,7 +66,7 @@ Compat = "4.16"
 CpuId = "0.3"
 DispatchDoctor = "0.4.12"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"

diff --git a/lib/LuxLib/test/Project.toml b/lib/LuxLib/test/Project.toml
@@ -39,7 +39,7 @@ BenchmarkTools = "1.5"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.18"
 Enzyme = "0.13.16"
-EnzymeCore = "0.8.6"
+EnzymeCore = "0.8.8"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"

diff --git a/test/layers/normalize_tests.jl b/test/layers/normalize_tests.jl
@@ -154,8 +154,6 @@ end
             @jet __f(z)
         end
 
-        broken_backends = VERSION ≥ v"1.11-" ? Any[AutoEnzyme()] : []
-
         @testset "Conv" begin
             c = Conv((3, 3), 3 => 3; init_bias=Lux.ones32)
 
@@ -165,35 +163,31 @@ end
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight, :bias), (2, 2))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(c, (:weight,), (2,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 3, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
         end
 
         @testset "Dense" begin
@@ -205,35 +199,31 @@ end
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight, :bias), (2, 2))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
 
             wn = WeightNorm(d, (:weight,), (2,))
             display(wn)
             ps, st = Lux.setup(rng, wn) |> dev
             x = randn(rng, Float32, 3, 1) |> aType
 
             @jet wn(x, ps, st)
-            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3,
-                broken_backends)
+            @test_gradients(sumabs2first, wn, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
         end
 
         # See https://github.com/LuxDL/Lux.jl/issues/95