test: display layers

Author: avik-pal

test/deeponet_tests.jl

@@ -7,7 +7,7 @@
         out_size = (10, 5)
 
         deeponet = DeepONet(; branch=(64, 32, 32, 16), trunk=(1, 8, 8, 16))
-
+        display(deeponet)
         ps, st = Lux.setup(rng, deeponet) |> dev
 
         @inferred deeponet((u, y), ps, st)
@@ -18,7 +18,7 @@
 
         deeponet = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 16)),
             Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
-
+        display(deeponet)
         ps, st = Lux.setup(rng, deeponet) |> dev
 
         @inferred deeponet((u, y), ps, st)
@@ -29,13 +29,15 @@
 
         deeponet = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 20)),
             Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
+        display(deeponet)
         ps, st = Lux.setup(rng, deeponet) |> dev
         @test_throws ArgumentError deeponet((u, y), ps, st)
 
         @testset "higher-dim input #7" begin
             u = ones(Float32, 10, 10, 10) |> aType
             v = ones(Float32, 1, 10, 10) |> aType
             deeponet = DeepONet(; branch=(10, 10, 10), trunk=(1, 10, 10))
+            display(deeponet)
             ps, st = Lux.setup(rng, deeponet) |> dev
 
             y, st_ = deeponet((u, v), ps, st)

test/fno_tests.jl

@@ -10,12 +10,12 @@
 
         @testset "$(length(setup.modes))D: permuted = $(setup.permuted)" for setup in setups
             fno = FourierNeuralOperator(; setup.chs, setup.modes, setup.permuted)
+            display(fno)
+            ps, st = Lux.setup(rng, fno) |> dev
 
             x = rand(rng, Float32, setup.x_size...) |> aType
             y = rand(rng, Float32, setup.y_size...) |> aType
 
-            ps, st = Lux.setup(rng, fno) |> dev
-
             @inferred fno(x, ps, st)
             @jet fno(x, ps, st)
 

test/layers_tests.jl

@@ -22,6 +22,7 @@
             l1 = p ? Conv(ntuple(_ -> 1, length(setup.m)), in_chs => first(ch)) :
                  Dense(in_chs => first(ch))
             m = Chain(l1, op(ch, setup.m; setup.permuted))
+            display(m)
             ps, st = Lux.setup(rng, m) |> dev
 
             x = rand(rng, Float32, setup.x_size...) |> aType