test: fft testing against FFTW

EnzymeAD · Nov 5, 2024 · 3e26602 · 3e26602
1 parent 34a1da2
commit 3e26602
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 2 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -30,6 +30,7 @@ jobs:
         test_group:
           - core
           - neural_networks
+          - integration
         arch:
           - x64
         assertions:

diff --git a/ext/ReactantAbstractFFTsExt.jl b/ext/ReactantAbstractFFTsExt.jl
@@ -33,6 +33,7 @@ end
 
 for op in (:rfft, :fft, :ifft)
     @eval function AbstractFFTs.$(op)(x::TracedRArray, dims)
+        @assert maximum(dims) ≤ ndims(x) "dims out of range"
         if dims isa Integer
             if dims != 1
                 pdims = compute_correct_pdims(x, dims)
@@ -54,6 +55,7 @@ end
 
 for op in (:irfft,)
     @eval function AbstractFFTs.$(op)(x::TracedRArray, d::Int, dims)
+        @assert maximum(dims) ≤ ndims(x) "dims out of range"
         if dims isa Integer
             if dims != 1
                 pdims = compute_correct_pdims(x, dims)

diff --git a/test/integration/fft.jl b/test/integration/fft.jl
@@ -0,0 +1,46 @@
+using FFTW, Reactant
+
+@testset "fft" begin
+    x = rand(ComplexF32, 2, 2, 3, 4)
+    x_ra = Reactant.ConcreteRArray(x)
+
+    @test_throws AssertionError @jit(fft(x_ra))
+
+    x = rand(ComplexF32, 2, 3, 4)
+    x_ra = Reactant.ConcreteRArray(x)
+
+    @test @jit(fft(x_ra)) ≈ fft(x)
+    @test @jit(fft(x_ra, (1, 2))) ≈ fft(x, (1, 2))
+    @test @jit(fft(x_ra, (1, 2, 3))) ≈ fft(x, (1, 2, 3))
+    @test @jit(fft(x_ra, (2, 3))) ≈ fft(x, (2, 3))
+    @test @jit(fft(x_ra, (1, 3))) ≈ fft(x, (1, 3))
+
+    @test_throws AssertionError @jit(fft(x_ra, (3, 2)))
+    @test_throws AssertionError @jit(fft(x_ra, (1, 4)))
+
+    y_ra = @jit(fft(x_ra))
+    @test @jit(ifft(y_ra)) ≈ x
+end
+
+@testset "rfft" begin
+    x = rand(2, 2, 3, 4)
+    x_ra = Reactant.ConcreteRArray(x)
+
+    @test_throws AssertionError @jit(rfft(x_ra))
+
+    x = rand(2, 3, 4)
+    x_ra = Reactant.ConcreteRArray(x)
+
+    @test @jit(rfft(x_ra)) ≈ rfft(x)
+    @test @jit(rfft(x_ra, (1, 2))) ≈ rfft(x, (1, 2))
+    @test @jit(rfft(x_ra, (1, 2, 3))) ≈ rfft(x, (1, 2, 3))
+    @test @jit(rfft(x_ra, (2, 3))) ≈ rfft(x, (2, 3))
+    @test @jit(rfft(x_ra, (1, 3))) ≈ rfft(x, (1, 3))
+
+    @test_throws AssertionError @jit(rfft(x_ra, (3, 2)))
+    @test_throws AssertionError @jit(rfft(x_ra, (1, 4)))
+
+    y_ra = @jit(rfft(x_ra))
+    @test @jit(irfft(y_ra, 2)) ≈ x
+    @test @jit(irfft(y_ra, 3)) ≈ irfft(rfft(x), 3)
+end
diff --git a/test/nn/nnlib.jl b/test/nn/nnlib.jl
@@ -131,8 +131,7 @@ end
     @test @jit(NNlib.pad_constant(x_ra, (1, 0))) ≈ NNlib.pad_constant(x, (1, 0))
 
     # Symmetric Padding with value (test type-casting)
-    @test @jit(NNlib.pad_constant(x_ra, (1, 1), 2)) ≈
-        NNlib.pad_constant(x, (1, 1), 2)
+    @test @jit(NNlib.pad_constant(x_ra, (1, 1), 2)) ≈ NNlib.pad_constant(x, (1, 1), 2)
     @test @jit(NNlib.pad_constant(x_ra, (1, 1, 1, 1), 2)) ≈
         NNlib.pad_constant(x, (1, 1, 1, 1), 2)
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -56,6 +56,10 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Control Flow" include("control_flow.jl")
     end
 
+    if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"
+        @safetestset "AbstractFFTs" include("integration/fft.jl")
+    end
+
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "neural_networks"
         @testset "Neural Networks" begin
             @safetestset "NNlib Primitives" include("nn/nnlib.jl")