test: conv and pool with more parameters

Author: avik-pal

test/nn/flux.jl

@@ -67,50 +67,3 @@ optim # parameters, momenta and output have all changed
 out2 = model(noisy)  # first row is prob. of true, second row p(false)
 
 mean((out2[1, :] .> 0.5) .== truth)  # accuracy 94% so far!
-
-@testset "conv: groups $groups" for groups in (1, 2, 4)
-    nn_conv = Conv(
-        randn(Float32, 10, 10, 8 ÷ groups, groups), randn(Float32, groups); groups
-    )
-    conv_reactant = Conv(
-        Reactant.ConcreteRArray(nn_conv.weight),
-        Reactant.ConcreteRArray(nn_conv.bias);
-        groups,
-    )
-
-    img = randn(Float32, 224, 224, 8, 2)
-    img_reactant = Reactant.ConcreteRArray(img)
-
-    comp_conv = @compile conv_reactant(img_reactant)
-
-    res_reactant = Array{Float32,4}(comp_conv(img_reactant))
-
-    @test res_reactant ≈ nn_conv(img)
-end
-
-@testset "conv 1d: flip" begin
-    x = [1; 2; 3;;;]
-    W = [1; 2; 3;;;]
-
-    xx = Reactant.ConcreteRArray(x)
-    WW = Reactant.ConcreteRArray(W)
-
-    conv_noflip(x, W) = NNlib.conv(x, W; pad=1, flipped=true)
-    conv_flip(x, W) = NNlib.conv(x, W; pad=1, flipped=false)
-
-    @test @compile(conv_noflip(xx, WW))(xx, WW) ==
-        [0*1+1*2+2*3; 1*1+2*2+3*3; 1*2+2*3+3*0;;;]
-    @test @compile(conv_flip(xx, WW))(xx, WW) == [3*0+2*1+1*2; 3*1+2*2+1*3; 3*2+2*3+1*0;;;]
-end
-
-@testset "$f" for f in (NNlib.meanpool, NNlib.maxpool)
-    img = randn(Float32, 224, 224, 3, 2)
-    img_reactant = Reactant.ConcreteRArray(img)
-
-    f_reactant = @compile f(img_reactant, (3, 3))
-
-    res_reactant = f_reactant(img_reactant, (3, 3))
-    res = f(img, (3, 3))
-
-    @test res_reactant ≈ res
-end

test/nn/luxlib.jl

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

test/nn/nnlib.jl

@@ -32,7 +32,61 @@ using NNlib, Reactant, Enzyme
 end
 
 @testset "Pooling" begin
+    @testset for f in (NNlib.meanpool, NNlib.maxpool)
+        x = randn(Float32, 32, 32, 3, 2)
+        x_reactant = Reactant.ConcreteRArray(x)
+
+        @testset for window in ((2, 2), (3, 3), (4, 4)),
+            stride in ((1, 1), (2, 2)),
+            padding in ((0, 0), (1, 1), (2, 2), (0, 2), (2, 0))
+
+            pool_dims = PoolDims(x, window; stride, padding)
+
+            f_reactant = Reactant.compile(f, (x_reactant, pool_dims))
+
+            broken = any(==(2), padding) && f === NNlib.maxpool && window == (2, 2)
+            @test f_reactant(x_reactant, pool_dims) ≈ f(x, pool_dims) broken = broken
+        end
+
+        # TODO: test for gradients
+    end
 end
 
 @testset "Convolution" begin
+    @testset for groups in (1, 2, 4)
+        weight = randn(Float32, 10, 10, 8 ÷ groups, groups)
+        x = randn(Float32, 16, 16, 8, 2)
+
+        weight_reactant = Reactant.ConcreteRArray(weight)
+        x_reactant = Reactant.ConcreteRArray(x)
+
+        @testset for stride in ((1, 1), (2, 2), (3, 3)), padding in ((0, 0), (1, 1), (2, 2))
+            conv_dims = DenseConvDims(x, weight; stride, padding, dilation=1, groups)
+
+            conv_compiled = Reactant.compile(
+                NNlib.conv, (x_reactant, weight_reactant, conv_dims)
+            )
+
+            @test conv_compiled(x_reactant, weight_reactant, conv_dims) ≈
+                NNlib.conv(x, weight, conv_dims)
+        end
+
+        # TODO: test for gradients
+    end
+
+    @testset "conv 1d: flip" begin
+        x = [1; 2; 3;;;]
+        W = [1; 2; 3;;;]
+    
+        xx = Reactant.ConcreteRArray(x)
+        WW = Reactant.ConcreteRArray(W)
+    
+        conv_noflip(x, W) = NNlib.conv(x, W; pad=1, flipped=true)
+        conv_flip(x, W) = NNlib.conv(x, W; pad=1, flipped=false)
+    
+        @test Reactant.compile(conv_noflip, (xx, WW))(xx, WW) ==
+              [0*1+1*2+2*3; 1*1+2*2+3*3; 1*2+2*3+3*0;;;]
+        @test Reactant.compile(conv_flip, (xx, WW))(xx, WW) ==
+              [3*0+2*1+1*2; 3*1+2*2+1*3; 3*2+2*3+1*0;;;]
+    end
 end