fix tests

Author: tknopp

test/Transform/fourier_transform.jl

@@ -5,10 +5,12 @@
 
     ft = FourierTransform((3, 4, 5))
 
-    @test size(transform(ft, 𝐱)) == (30, 40, 50, ch, batch)
+    @test size(transform(ft, 𝐱)) == (16, 40, 50, ch, batch)
     @test size(truncate_modes(ft, transform(ft, 𝐱))) == (3, 4, 5, ch, batch)
-    @test size(inverse(ft, truncate_modes(ft, transform(ft, 𝐱)), size(𝐱))) == (3, 4, 5, ch, batch)
+    @test size(inverse(ft, NeuralOperators.pad_modes( truncate_modes(ft, transform(ft, 𝐱)), size(transform(ft, 𝐱)) ), 
+                         size(𝐱))) == (30, 40, 50, ch, batch)
 
-    g = Zygote.gradient(x -> sum(inverse(ft, truncate_modes(ft, transform(ft, x)), size(𝐱))), 𝐱)
+    g = Zygote.gradient(x -> sum(inverse(ft, NeuralOperators.pad_modes( truncate_modes(ft, transform(ft, x)), 
+                         (16, 40, 50, ch, batch) ), (30, 40, 50, ch, batch) )), 𝐱)
     @test size(g[1]) == (30, 40, 50, ch, batch)
-end
+end

test/operator_kernel.jl

@@ -71,7 +71,7 @@ end
 end
 
 @testset "2D OperatorConv" begin
-    modes = (16, 16)
+    modes = (10, 10)
     ch = 64 => 64
 
     m = Chain(Dense(1, 64),
@@ -87,7 +87,7 @@ end
 end
 
 @testset "permuted 2D OperatorConv" begin
-    modes = (16, 16)
+    modes = (10, 10)
     ch = 64 => 64
 
     m = Chain(Conv((1, 1), 1 => 64),
@@ -104,7 +104,7 @@ end
 end
 
 @testset "2D OperatorKernel" begin
-    modes = (16, 16)
+    modes = (10, 10)
     ch = 64 => 64
 
     m = Chain(Dense(1, 64),
@@ -119,7 +119,7 @@ end
 end
 
 @testset "permuted 2D OperatorKernel" begin
-    modes = (16, 16)
+    modes = (10, 10)
     ch = 64 => 64
 
     m = Chain(Conv((1, 1), 1 => 64),