feat: more support for complex

Author: avik-pal

Project.toml

@@ -13,6 +13,12 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
+[extensions]
+NeuralOperatorsReactantExt = "Reactant"
+
+[weakdeps]
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+
 [compat]
 AbstractFFTs = "1.5.0"
 ConcreteStructs = "0.2.3"
@@ -21,5 +27,6 @@ LuxCore = "1.2"
 LuxLib = "1.8"
 NNlib = "0.9.30"
 Random = "1.10"
+Reactant = "0.2.129"
 WeightInitializers = "1"
 julia = "1.10"

ext/NeuralOperatorsReactantExt.jl

@@ -0,0 +1,8 @@
+module NeuralOperatorsReactantExt
+
+using NeuralOperators: NeuralOperators
+using Reactant: Reactant, TracedRNumber
+
+NeuralOperators.unwrapped_eltype(x::TracedRNumber) = Reactant.unwrapped_eltype(x)
+
+end

src/models/fno.jl

@@ -64,3 +64,22 @@ function FourierNeuralOperator(
         ),
     )
 end
+
+function FourierNeuralOperator(
+    modes::Dims,
+    in_channels::Integer,
+    out_channels::Integer,
+    hidden_channels::Integer;
+    num_layers::Integer=4,
+    lifting_channel_ratio::Integer=2,
+    projection_channel_ratio::Integer=2,
+    positional_embedding::Union{Symbol,AbstractLuxLayer}=:grid, # :grid | :none
+    activation=gelu,
+    use_channel_mlp::Bool=true,
+    channel_mlp_dropout_rate::Real=0.0,
+    channel_mlp_expansion::Real=0.5,
+    channel_mlp_skip::Symbol=:soft_gating,
+    fno_skip::Symbol=:linear,
+)
+    return nothing
+end

src/utils.jl

@@ -39,3 +39,18 @@ function meshgrid(args::AbstractVector...)
         end
     end
 end
+
+unwrapped_eltype(x) = eltype(x)
+
+function decomposed_activation(f::F, x::Number) where {F}
+    unwrapped_eltype(x) <: Complex && return Complex(f(real(x)), f(imag(x)))
+    return f(x)
+end
+
+apply_complex((rfn, ifn), x::Number) = apply_complex(rfn, ifn, x)
+function apply_complex(rfn, ifn, x::Number)
+    @assert unwrapped_eltype(x) <: Complex "Expected a complex number, got \
+                                            $(unwrapped_eltype(x))"
+    rl, img = real(x), imag(x)
+    return Complex(rfn(rl) - ifn(img), rfn(img) + ifn(rl))
+end