Merge pull request #4 from ayushinav/main

Minor fixes

Author: avik-pal

src/fno.jl

@@ -16,7 +16,7 @@ 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.
-  - `permuted`: Whether the dim is permuted. If `permuted = Val(true)`, the layer accepts
+  - `permuted`: Whether the dim is permuted. If `permuted = Val(false)`, the layer accepts
     data in the order of `(ch, x_1, ... , x_d , batch)`. Otherwise the order is
     `(x_1, ... , x_d, ch, batch)`.
 

src/layers.jl

@@ -15,7 +15,7 @@
 ## Keyword Arguments
 
   - `init_weight`: Initial function to initialize parameters.
-  - `permuted`: Whether the dim is permuted. If `permuted = Val(true)`, the layer accepts
+  - `permuted`: Whether the dim is permuted. If `permuted = Val(false)`, the layer accepts
     data in the order of `(ch, x_1, ... , x_d , batch)`. Otherwise the order is
     `(x_1, ... , x_d, ch, batch)`.
   - `T`: Datatype of parameters.
@@ -86,7 +86,7 @@ SpectralConv(args...; kwargs...) = OperatorConv(args..., FourierTransform; kwarg
   - `ch`: A `Pair` of input and output channel size `ch_in => ch_out`, e.g. `64 => 64`.
   - `modes`: The modes to be preserved. A tuple of length `d`, where `d` is the dimension of
     data.
-  - `::Type{TR}`: The traform to operate the transformation.
+  - `::Type{TR}`: The transform to operate the transformation.
 
 ## Keyword Arguments
 
@@ -237,7 +237,7 @@ end
 end
 
 @inline function __apply_pattern_batched_mul(x, y)
-    # Use permutedims to guarantee contiguous memory
+    # Use permutedims to guarantee contiguous memory 
     x_ = permutedims(x, (2, 3, 1))         # i x b x m
     res = batched_mul(y, x_)               # o x b x m
     return permutedims(res, (3, 1, 2))     # m x o x b

test/fno.jl

@@ -12,7 +12,7 @@ include("test_utils.jl")
             y_size=(1024, 1, 5), permuted=Val(true))]
 
     @testset "$(length(setup.modes))D: permuted = $(setup.permuted)" for setup in setups
-        fno = FourierNeuralOperator(; rng, setup.chs, setup.modes, setup.permuted)
+        fno = FourierNeuralOperator(rng; setup.chs, setup.modes, setup.permuted)
 
         x = rand(rng, Float32, setup.x_size...)
         y = rand(rng, Float32, setup.y_size...)

test/runtest.jl renamed to test/runtests.jl

@@ -40,7 +40,7 @@ default_loss_function(model, ps, x, y) = mean(abs2, y .- model(x, ps))
 train!(args...; kwargs...) = train!(default_loss_function, args...; kwargs...)
 
 function train!(loss, model, ps, st, data; epochs = 10)
-    m = Lux.Experimental.StatefulLuxLayer(model, ps, st)
+    m = Lux.StatefulLuxLayer(model, ps, st)
 
     l1 = loss(m, ps, first(data)...)
     st_opt = Optimisers.setup(Adam(0.01f0), ps)