Add more building blocks

Author: avik-pal

Project.toml

@@ -12,6 +12,7 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 
 [compat]
 ChainRulesCore = "1"

src/LuxNeuralOperators.jl

@@ -4,12 +4,12 @@ import PrecompileTools: @recompile_invalidations
 import Reexport: @reexport
 
 @recompile_invalidations begin
-    using ArrayInterface, FFTW, Lux, Random
+    using ArrayInterface, FFTW, Lux, Random, SciMLBase
 
     import ChainRulesCore as CRC
     import Lux.Experimental: @compact
-    import LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer,
-        initialparameters, initialstates
+    import LuxCore: AbstractExplicitLayer,
+        AbstractExplicitContainerLayer, initialparameters, initialstates
     import Random: AbstractRNG
 end
 
@@ -23,6 +23,18 @@ const False = Val(false)
 include("transform.jl")
 include("layers.jl")
 include("fno.jl")
+include("deq.jl")
+
+# Pass `rng` if user doesn't pass it
+for f in (:BasicBlock, :StackedBasicBlock, :OperatorConv, :OperatorKernel,
+    :FourierNeuralOperator)
+    @eval begin
+        $(f)(args...; kwargs...) = $(f)(__default_rng(), args...; kwargs...)
+    end
+end
+
+__destructure(x::Tuple) = x
+__destructure(x) = x, zero(eltype(x))
 
 export FourierTransform
 export SpectralConv, OperatorConv

src/deq.jl

@@ -0,0 +1 @@
+# Basic DEQ Implementation

src/fno.jl

@@ -1,6 +1,6 @@
 """
-    FourierNeuralOperator(; chs = (2, 64, 64, 64, 64, 64, 128, 1), modes = (16,),
-        σ = gelu, rng = nothing, permuted::Val = Val(false), kwargs...)
+    FourierNeuralOperator([rng = __default_rng()]; chs = (2, 64, 64, 64, 64, 64, 128, 1),
+        modes = (16,), σ = gelu, permuted::Val = Val(false), kwargs...)
 
 Fourier neural operator is a operator learning model that uses Fourier kernel to perform
 spectral convolutions. It is a promising way for surrogate methods, and can be regarded as
@@ -16,7 +16,6 @@ kernels, and two `Dense` layers to project data back to the scalar field of inte
   - `modes`: The modes to be preserved. A tuple of length `d`, where `d` is the dimension
     of data.
   - `σ`: Activation function for all layers in the model.
-  - `rng`: Random number generator. If provided, it is forwarded to the Operator Layers.
   - `permuted`: Whether the dim is permuted. If `permuted = Val(true)`, the layer accepts
     data in the order of `(ch, x_1, ... , x_d , batch)`. Otherwise the order is
     `(x_1, ... , x_d, ch, batch)`.
@@ -35,28 +34,28 @@ Chain(
             l1 = Dense(64 => 64),       # 4_160 parameters
             l2 = OperatorConv{FourierTransform}(64 => 64, (16,); permuted = false)(),  # 65_536 parameters, plus 2
             activation = σ,
-        ) do x::(AbstractArray{<:Real, M} where M) 
+        ) do x::(AbstractArray{<:Real, M} where M)
             return activation.(l1(x) .+ l2(x))
         end,
         layer_2 = @compact(
             l1 = Dense(64 => 64),       # 4_160 parameters
             l2 = OperatorConv{FourierTransform}(64 => 64, (16,); permuted = false)(),  # 65_536 parameters, plus 2
             activation = σ,
-        ) do x::(AbstractArray{<:Real, M} where M) 
+        ) do x::(AbstractArray{<:Real, M} where M)
             return activation.(l1(x) .+ l2(x))
         end,
         layer_3 = @compact(
             l1 = Dense(64 => 64),       # 4_160 parameters
             l2 = OperatorConv{FourierTransform}(64 => 64, (16,); permuted = false)(),  # 65_536 parameters, plus 2
             activation = σ,
-        ) do x::(AbstractArray{<:Real, M} where M) 
+        ) do x::(AbstractArray{<:Real, M} where M)
             return activation.(l1(x) .+ l2(x))
         end,
         layer_4 = @compact(
             l1 = Dense(64 => 64),       # 4_160 parameters
             l2 = OperatorConv{FourierTransform}(64 => 64, (16,); permuted = false)(),  # 65_536 parameters, plus 2
             activation = σ,
-        ) do x::(AbstractArray{<:Real, M} where M) 
+        ) do x::(AbstractArray{<:Real, M} where M)
             return activation.(l1(x) .+ l2(x))
         end,
     ),
@@ -68,10 +67,9 @@ Chain(
           #        plus 12 states.
 ```
 """
-function FourierNeuralOperator(; chs=(2, 64, 64, 64, 64, 64, 128, 1), modes=(16,),
-        σ=gelu, rng=nothing, permuted::Val{P}=False, kwargs...) where {P}
+function FourierNeuralOperator(rng::AbstractRNG; chs=(2, 64, 64, 64, 64, 64, 128, 1),
+        modes=(16,), σ=gelu, permuted::Val{P}=False, kwargs...) where {P}
     @assert length(chs) ≥ 5
-    rng === nothing && (rng = __default_rng())
 
     map₁ = chs[1] => chs[2]
     map₂ = chs[end - 2] => chs[end - 1]

src/layers.jl

@@ -1,6 +1,6 @@
 """
-    OperatorConv([rng::AbstractRNG = __defautl_rng()], ch::Pair{<:Integer, <:Integer},
-        _modes::NTuple{N, <:Integer}, ::Type{TR}; init_weight = glorot_uniform,
+    OperatorConv([rng::AbstractRNG = __default_rng()], ch::Pair{<:Integer, <:Integer},
+        modes::NTuple{N, <:Integer}, ::Type{TR}; init_weight = glorot_uniform,
         T::Type{TP} = ComplexF32,
         permuted::Val{P} = Val(false)) where {N, TR <: AbstractTransform, TP, P}
 
@@ -40,13 +40,13 @@ function OperatorConv(rng::AbstractRNG, ch::Pair{<:Integer, <:Integer},
     name = "OperatorConv{$TR}($in_chs => $out_chs, $modes; permuted = $permuted)"
 
     if permuted === True
-        return @compact(; modes, weights, transform,
+        return @compact(; modes, weights, transform, dispatch=:OperatorConv,
             name) do x::AbstractArray{<:Real, M} where {M}
             y = __operator_conv(x, transform, weights)
             return y
         end
     else
-        return @compact(; modes, weights, transform,
+        return @compact(; modes, weights, transform, dispatch=:OperatorConv,
             name) do x::AbstractArray{<:Real, M} where {M}
             N_ = ndims(transform)
             xᵀ = permutedims(x, (ntuple(i -> i + 1, N_)..., 1, N_ + 2))
@@ -57,8 +57,6 @@ function OperatorConv(rng::AbstractRNG, ch::Pair{<:Integer, <:Integer},
     end
 end
 
-OperatorConv(args...; kwargs...) = OperatorConv(__default_rng(), args...; kwargs...)
-
 """
     SpectralConv(args...; kwargs...)
 
@@ -78,7 +76,7 @@ OperatorConv{FourierTransform}(2 => 5, (16,); permuted = true)()  # 160 paramete
 SpectralConv(args...; kwargs...) = OperatorConv(args..., FourierTransform; kwargs...)
 
 """
-    OperatorKernel([rng::AbstractRNG = __defautl_rng()], ch::Pair{<:Integer, <:Integer},
+    OperatorKernel([rng::AbstractRNG = __default_rng()], ch::Pair{<:Integer, <:Integer},
         modes::NTuple{N, <:Integer}, transform::Type{TR}; σ = identity,
         permuted::Val{P} = Val(false), kwargs...) where {N, TR <: AbstractTransform, P}
 
@@ -128,7 +126,8 @@ function OperatorKernel(rng::AbstractRNG, ch::Pair{<:Integer, <:Integer},
     l₁ = permuted === True ? Conv(map(_ -> 1, modes), ch) : Dense(ch)
     l₂ = OperatorConv(rng, ch, modes, transform; permuted, kwargs...)
 
-    return @compact(; l₁, l₂, activation=σ) do x::AbstractArray{<:Real, M} where {M}
+    return @compact(; l₁, l₂, activation=σ,
+        dispatch=:OperatorKernel) do x::AbstractArray{<:Real, M} where {M}
         return activation.(l₁(x) .+ l₂(x))
     end
 end
@@ -165,7 +164,60 @@ end       # Total: 175 parameters,
 """
 SpectralKernel(args...; kwargs...) = OperatorKernel(args..., FourierTransform; kwargs...)
 
-OperatorKernel(args...; kwargs...) = OperatorKernel(__default_rng(), args...; kwargs...)
+# Building Blocks
+function BasicBlock(rng::AbstractRNG, ch::Integer, modes::NTuple{N, <:Integer}, args...;
+        add_mlp::Val=False, normalize::Val=False, σ=gelu, kwargs...) where {N}
+    conv1 = SpectralConv(rng, ch => ch, modes, args...; kwargs...)
+    conv2 = SpectralConv(rng, ch => ch, modes, args...; kwargs...)
+    conv3 = SpectralConv(rng, ch => ch, modes, args...; kwargs...)
+
+    kernel_size = map(_ -> 1, modes)
+
+    if add_mlp === True
+        mlp1 = Chain(Conv(kernel_size, ch => ch, σ), Conv(kernel_size, ch => ch))
+        mlp2 = Chain(Conv(kernel_size, ch => ch, σ), Conv(kernel_size, ch => ch))
+        mlp3 = Chain(Conv(kernel_size, ch => ch, σ), Conv(kernel_size, ch => ch))
+    else
+        mlp1 = NoOpLayer()
+        mlp2 = NoOpLayer()
+        mlp3 = NoOpLayer()
+    end
+
+    norm = normalize === True ? InstanceNorm(ch; affine = false) : NoOpLayer()
+
+    w1 = Conv(kernel_size, ch => ch)
+    w2 = Conv(kernel_size, ch => ch)
+    w3 = Conv(kernel_size, ch => ch)
+
+    return @compact(; conv1, conv2, conv3, mlp1, mlp2, mlp3, w1, w2, w3, norm, σ,
+        dispatch=:BasicBlock) do (inp,)
+        x, injection = __destructure(inp)
+
+        x = norm(x)
+        x1 = norm(mlp1(norm(conv1(x))))
+        x2 = norm(w1(x))
+        x = σ.(x1 .+ x2 .+ injection)
+
+        x1 = norm(mlp2(norm(conv2(x))))
+        x2 = norm(w2(x))
+        x = σ.(x1 .+ x2 .+ injection)
+
+        x1 = norm(mlp3(norm(conv3(x))))
+        x2 = norm(w3(x))
+        return σ.(x1 .+ x2 .+ injection)
+    end
+end
+
+function StackedBasicBlock(rng::AbstractRNG, args...; depth::Val{N} = Val(1),
+        kwargs...) where {N}
+    blocks = ntuple(i -> BasicBlock(rng, args...; kwargs...), N)
+    block = NamedTuple{ntuple(i -> Symbol("block_", i), N)}(blocks)
+
+    return @compact(; block, dispatch=:StackedBasicBlock) do (inp,)
+        x, injection = __destructure(inp)
+        return __applyblock(block, x, injection)
+    end
+end
 
 # Functional Versions
 @inline function __operator_conv(x, transform, weights)
@@ -191,6 +243,13 @@ end
     return permutedims(res, (3, 1, 2))     # m x o x b
 end
 
+@inline @generated function __applyblock(block::NamedTuple{names}, x, inj) where {names}
+    calls = [:(x = block.$name((x, inj))) for name in names]
+    return quote
+        $(calls...)
+    end
+end
+
 @inline __pad_modes(x, dims::Integer...) = __pad_modes(x, dims)
 @inline __pad_modes(x, dims::NTuple) = __pad_modes!(similar(x, dims), x)