Improve permuted

src/fourier.jl

@@ -3,23 +3,21 @@ export
     FourierOperator
 
 struct SpectralConv{P, N, T, S, F}
-    permuted::Bool
     weight::T
     in_channel::S
     out_channel::S
     modes::NTuple{N, S}
     σ::F
 end
 
-function SpectralConv(
-    permuted::Bool,
+function SpectralConv{P}(
     weight::T,
     in_channel::S,
     out_channel::S,
     modes::NTuple{N, S},
     σ::F
-) where {N, T, S, F}
-    return SpectralConv{permuted, N, T, S, F}(permuted, weight, in_channel, out_channel, modes, σ)
+) where {P, N, T, S, F}
+    return SpectralConv{P, N, T, S, F}(weight, in_channel, out_channel, modes, σ)
 end
 
 """
@@ -63,13 +61,16 @@ function SpectralConv(
     scale = one(T) / (in_chs * out_chs)
     weights = scale * init(out_chs, in_chs, prod(modes))
 
-    return SpectralConv(permuted, weights, in_chs, out_chs, modes, σ)
+    return SpectralConv{permuted}(weights, in_chs, out_chs, modes, σ)
 end
 
-Flux.@functor SpectralConv
+Flux.@functor SpectralConv{true}
+Flux.@functor SpectralConv{false}
 
 Base.ndims(::SpectralConv{P, N}) where {P, N} = N
 
+ispermuted(::SpectralConv{P}) where {P} = P
+
 function Base.show(io::IO, l::SpectralConv{P}) where {P}
     print(io, "SpectralConv($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)), permuted=$P)")
 end
@@ -156,7 +157,7 @@ function Base.show(io::IO, l::FourierOperator)
             "$(l.conv.in_channel) => $(l.conv.out_channel), " *
             "$(l.conv.modes), " *
             "σ=$(string(l.σ)), " *
-            "permuted=$(l.conv.permuted)" *
+            "permuted=$(ispermuted(l.conv))" *
         ")"
     )
 end