test: move the tests around a bit

Author: avik-pal

test/bcast.jl

@@ -57,37 +57,6 @@ function test()
 end
 test()
 
-@testset "Activation Functions" begin
-    sumabs2(f, x) = sum(abs2, f.(x))
-
-    function ∇sumabs2(f, x)
-        dx = Enzyme.make_zero(x)
-        Enzyme.autodiff(Reverse, sumabs2, Active, Const(f), Duplicated(x, dx))
-        return dx
-    end
-
-    x_act = randn(Float32, 10, 10)
-    x_act_ca = Reactant.ConcreteRArray(x_act)
-
-    @testset "Activation: $act" for act in (
-        identity, relu, sigmoid, tanh, tanh_fast, sigmoid_fast, gelu, abs2
-    )
-        f_compile = @compile sumabs2(act, x_act)
-
-        y_simple = sumabs2(act, x_act)
-        y_compile = f_compile(act, x_act_ca)
-
-        ∂x_enz = Enzyme.make_zero(x_act)
-        Enzyme.autodiff(Reverse, sumabs2, Active, Const(act), Duplicated(x_act, ∂x_enz))
-
-        ∇sumabs2_compiled = @compile ∇sumabs2(act, x_act_ca)
-
-        ∂x_compile = ∇sumabs2_compiled(act, x_act_ca)
-
-        @test y_simple ≈ y_compile
-    end
-end
-
 @testset "ConcreteRArray broadcasting" begin
     x = ones(10, 10)
     y = ones(10, 10)

test/nn/luxlib.jl

@@ -0,0 +1,13 @@
+using LuxLib, Reactant, Enzyme
+
+@testset "Fused Dense" begin
+end
+
+@testset "Bias Activation" begin
+end
+
+@testset "Fast Activation" begin
+end
+
+@testset "Fused Conv" begin
+end

test/nn/nnlib.jl

@@ -0,0 +1,38 @@
+using NNlib, Reactant, Enzyme
+
+@testset "Activation Functions" begin
+    sumabs2(f, x) = sum(abs2, f.(x))
+
+    function ∇sumabs2(f, x)
+        dx = Enzyme.make_zero(x)
+        Enzyme.autodiff(Reverse, sumabs2, Active, Const(f), Duplicated(x, dx))
+        return dx
+    end
+
+    x_act = randn(Float32, 10, 10)
+    x_act_ca = Reactant.ConcreteRArray(x_act)
+
+    @testset "Activation: $act" for act in (
+        identity, relu, sigmoid, tanh, tanh_fast, sigmoid_fast, gelu, abs2
+    )
+        f_compile = Reactant.compile(sumabs2, (act, x_act))
+
+        y_simple = sumabs2(act, x_act)
+        y_compile = f_compile(act, x_act_ca)
+
+        ∂x_enz = Enzyme.make_zero(x_act)
+        Enzyme.autodiff(Reverse, sumabs2, Active, Const(act), Duplicated(x_act, ∂x_enz))
+
+        ∇sumabs2_compiled = Reactant.compile(∇sumabs2, (act, x_act_ca))
+
+        ∂x_compile = ∇sumabs2_compiled(act, x_act_ca)
+
+        @test y_simple ≈ y_compile
+    end
+end
+
+@testset "Pooling" begin
+end
+
+@testset "Convolution" begin
+end