Merge pull request #60 from SciML/comp

Update comp

Author: foldfelis

Project.toml

@@ -20,8 +20,8 @@ CUDA = "3.8"
 CUDAKernels = "0.3, 0.4"
 ChainRulesCore = "1.13"
 FFTW = "1.4"
-Flux = "0.12"
-GeometricFlux = "0.10"
+Flux = "0.13"
+GeometricFlux = "0.11"
 KernelAbstractions = "0.7, 0.8"
 Tullio = "0.3"
 Zygote = "0.6"

src/model.jl

@@ -27,7 +27,7 @@ function FourierNeuralOperator(;
         OperatorKernel(ch[5]=>ch[6], modes, Transform),
         Dense(ch[6], ch[7], σ),
         Dense(ch[7], ch[8]),
-        flatten
+        Flux.flatten
     )
 end
 
@@ -48,7 +48,7 @@ function MarkovNeuralOperator(;
     σ=gelu
 )
     Transform = FourierTransform
-    
+
     return Chain(
         Dense(ch[1], ch[2]),
         OperatorKernel(ch[2]=>ch[3], modes, Transform, σ),

test/model.jl

@@ -6,7 +6,7 @@
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
     data = [(𝐱, 𝐲)]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "MarkovNeuralOperator" begin
@@ -17,5 +17,5 @@ end
 
     loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
     data = [(𝐱, 𝐲)]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end

test/operator_kernel.jl

@@ -14,7 +14,7 @@
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 128, 1024, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "permuted 1D OperatorConv" begin
@@ -34,7 +34,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 1024, 128, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "1D OperatorKernel" begin
@@ -52,7 +52,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 128, 1024, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "permuted 1D OperatorKernel" begin
@@ -71,7 +71,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 1024, 128, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "2D OperatorConv" begin
@@ -89,7 +89,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 64, 22, 22, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "permuted 2D OperatorConv" begin
@@ -108,7 +108,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 22, 22, 64, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "2D OperatorKernel" begin
@@ -125,7 +125,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 64, 22, 22, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "permuted 2D OperatorKernel" begin
@@ -143,7 +143,7 @@ end
 
     loss(x, y) = Flux.mse(m(x), y)
     data = [(𝐱, rand(Float32, 22, 22, 64, 5))]
-    Flux.train!(loss, params(m), data, Flux.ADAM())
+    Flux.train!(loss, Flux.params(m), data, Flux.ADAM())
 end
 
 @testset "SpectralConv" begin
@@ -165,7 +165,7 @@ end
     graph = grid([10, 10])
     𝐱 = rand(Float32, channel, N, batch_size)
     l = WithGraph(FeaturedGraph(graph), GraphKernel(κ, channel))
-    @test repr(l.layer) == "GraphKernel(Dense(64, 32, relu), channel=32)"
+    @test repr(l.layer) == "GraphKernel(Dense(64 => 32, relu), channel=32)"
     @test size(l(𝐱)) == (channel, N, batch_size)
 
     g = Zygote.gradient(() -> sum(l(𝐱)), Flux.params(l))