refactor SpectralConv1d

Author: foldfelis

src/NeuralOperators.jl

@@ -6,6 +6,9 @@ module NeuralOperators
     using Flux
     using FFTW
     using Tullio
+    using CUDA
+    using CUDAKernels
+    using KernelAbstractions
     using Zygote
 
     function __init__()

src/fourier.jl

@@ -1,5 +1,3 @@
-using CUDA, CUDAKernels, KernelAbstractions
-
 export
     SpectralConv1d,
     FourierOperator,
@@ -17,9 +15,6 @@ function c_glorot_uniform(dims...)
     return Flux.glorot_uniform(dims...) + Flux.glorot_uniform(dims...) * im
 end
 
-t(𝐱) = @tullio 𝐱ᵀ[a, b, c] := 𝐱[b, a, c]
-ein_mul(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m]
-
 function SpectralConv1d(
     ch::Pair{<:Integer, <:Integer},
     modes::Integer,
@@ -32,29 +27,32 @@ function SpectralConv1d(
     weights = scale * init(out_chs, in_chs, modes)
 
     return Chain(
-        t,
         x -> Zygote.hook(real, x),
         SpectralConv1d(weights, in_chs, out_chs, modes, σ),
-        t
     )
 end
 
 Flux.@functor SpectralConv1d
 
+t(𝐱) = @tullio 𝐱ᵀ[i, j, k] := 𝐱[j, i, k]
+ein_mul(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m]
+
 function (m::SpectralConv1d)(𝐱::AbstractArray)
-    𝐱_fft = fft(𝐱, 1) # [x, in_chs, batch]
+    𝐱ᵀ = t(𝐱) # [x, in_chs, batch] <- [in_chs, x, batch]
+    𝐱_fft = fft(𝐱ᵀ, 1) # [x, in_chs, batch]
     𝐱_selected = 𝐱_fft[1:m.modes, :, :] # [modes, in_chs, batch]
 
-    # [modes, out_chs, batch] <- [modes, in_chs, batch] [out_chs, in_chs, modes]
+    # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes]
     𝐱_weighted = ein_mul(𝐱_selected, m.weight)
 
     s = size(𝐱_weighted)[2:end]
-    d = size(𝐱, 1) - m.modes
+    d = size(𝐱ᵀ, 1) - m.modes
     𝐱_padded = cat(𝐱_weighted, zeros(ComplexF32, d, s...), dims=1)
 
     𝐱_out = ifft(𝐱_padded, 1) # [x, out_chs, batch]
+    𝐱_outᵀ = t(𝐱_out) # [out_chs, x, batch] <- [x, out_chs, batch]
 
-    return m.σ.(real(𝐱_out))
+    return m.σ.(real(𝐱_outᵀ))
 end
 
 function FourierOperator(