test: check if extended test sets are useful

LuxDL · Jul 5, 2024 · 63642a7 · 63642a7
1 parent 31d03b6
commit 63642a7
Show file tree

Hide file tree

Showing 10 changed files with 12 additions and 11 deletions.
diff --git a/Project.toml b/Project.toml
@@ -84,9 +84,10 @@ ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [targets]
-test = ["AMDGPU", "Aqua", "CUDA", "ComponentArrays", "ExplicitImports", "ForwardDiff", "LuxCUDA", "LuxDeviceUtils", "LuxTestUtils", "ReTestItems", "ReverseDiff", "StableRNGs", "Test", "Tracker", "Zygote", "cuDNN"]
+test = ["AMDGPU", "Aqua", "CUDA", "ComponentArrays", "ExplicitImports", "ForwardDiff", "LuxCUDA", "LuxDeviceUtils", "LuxTestUtils", "ReTestItems", "ReverseDiff", "StableRNGs", "Test", "TestSetExtensions", "Tracker", "Zygote", "cuDNN"]
diff --git a/test/batchnorm_tests.jl b/test/batchnorm_tests.jl
@@ -15,7 +15,7 @@
         end
     end
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         @testset "eltype $T, size $sz, $act" for T in (Float16, Float32, Float64),
             sz in ((4, 4, 6, 2), (8, 2), (4, 4, 4, 3, 2)),
             training in (Val(true), Val(false)),

diff --git a/test/conv_tests.jl b/test/conv_tests.jl
@@ -18,7 +18,7 @@
 
     anonact = x -> gelu(x)
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         # These are not all possible combinations but rather a representative set to keep
         # CI timings under check
         # Most of the actual tests happen upstream in Lux

diff --git a/test/dense_tests.jl b/test/dense_tests.jl
@@ -1,7 +1,7 @@
 @testitem "Fused Dense Bias Activation" tags=[:common_ops] setup=[SharedTestSetup] begin
     rng = get_stable_rng(12345)
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         # These are not all possible combinations but rather a representative set to keep
         # CI timings under check
         @testset "$(Tw) x $(Tx)" for (Tw, Tx) in [

diff --git a/test/dropout_tests.jl b/test/dropout_tests.jl
@@ -134,7 +134,7 @@ end
 
     rng = get_stable_rng(12345)
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         for T in (Float16, Float32, Float64),
             x_shape in ((2, 3), (2, 2, 3), (2, 2, 3, 1), (2, 2, 1, 3, 1))
 

diff --git a/test/forwarddiff_tests.jl b/test/forwarddiff_tests.jl
@@ -37,7 +37,7 @@
         end
     end
 
-    @testset "$(mode): Jacobian Vector Products" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$(mode): Jacobian Vector Products" for (mode, aType, on_gpu) in MODES
         @testset "$(op)(; flipped = $flipped)" for flipped in (true, false),
             op in (depthwiseconv, conv)
 
@@ -96,7 +96,7 @@ end
 
     rng = get_stable_rng(12345)
 
-    @testset "$mode: dropout" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode: dropout" for (mode, aType, on_gpu) in MODES
         x = randn(rng, Float32, 10, 2) |> aType
         x_dual = ForwardDiff.Dual.(x)
 

diff --git a/test/groupnorm_tests.jl b/test/groupnorm_tests.jl
@@ -8,7 +8,7 @@
         return x, scale, bias
     end
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         @testset "eltype $T, size $sz, ngroups $groups, $act" for T in (
                 Float16, Float32, Float64),
             sz in ((4, 6, 2), (8, 8, 8, 6, 2), (3, 16, 16, 12, 2),

diff --git a/test/instancenorm_tests.jl b/test/instancenorm_tests.jl
@@ -10,7 +10,7 @@
         return x, scale, bias
     end
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         @testset "eltype $T, size $sz, $act" for T in (Float16, Float32, Float64),
             sz in ((4, 4, 6, 2), (3, 4, 2), (4, 4, 4, 3, 2)),
             training in (Val(true), Val(false)),

diff --git a/test/layernorm_tests.jl b/test/layernorm_tests.jl
@@ -12,7 +12,7 @@
         end
     end
 
-    @testset "$mode" for (mode, aType, on_gpu) in MODES
+    @testset ExtendedTestSet "$mode" for (mode, aType, on_gpu) in MODES
         @testset "eltype $T, size $x_shape, $act" for T in (Float16, Float32, Float64),
             x_shape in ((3, 3, 2, 1), (2, 2, 2, 1), (2, 3, 2, 2)),
             affine_shape in (nothing, x_shape[1:3], (1, 1, 1), (1, 1, x_shape[3])),

diff --git a/test/shared_testsetup.jl b/test/shared_testsetup.jl
@@ -3,7 +3,7 @@ import Reexport: @reexport
 
 using LuxLib, LuxCUDA, AMDGPU
 using LuxDeviceUtils
-@reexport using LuxTestUtils, StableRNGs, Test, Zygote
+@reexport using LuxTestUtils, StableRNGs, Test, TestSetExtensions, Zygote
 import LuxTestUtils: @jet, @test_gradients, check_approx
 
 const BACKEND_GROUP = get(ENV, "BACKEND_GROUP", "All")