test: enable Enzyme testing partly

Author: avik-pal

.JuliaFormatter.toml

@@ -7,4 +7,3 @@ format_docstrings = true
 separate_kwargs_with_semicolon = true
 always_for_in = true
 annotate_untyped_fields_with_any = false
-join_lines_based_on_source = false

.gitignore

@@ -1,4 +1,5 @@
 Manifest.toml
+Manifest-v*.toml
 .vscode
 wip
 examples

test/deeponet_tests.jl

@@ -10,7 +10,8 @@
             (u_size=(64, 3, 5), y_size=(4, 10, 5), out_size=(3, 10, 5),
                 branch=(64, 32, 32, 16), trunk=(4, 8, 8, 16), name="Vector"),
             (u_size=(64, 4, 3, 3, 5), y_size=(4, 10, 5), out_size=(4, 3, 3, 10, 5),
-                branch=(64, 32, 32, 16), trunk=(4, 8, 8, 16), name="Tensor")]
+                branch=(64, 32, 32, 16), trunk=(4, 8, 8, 16), name="Tensor")
+        ]
 
         @testset "$(setup.name)" for setup in setups
             u = rand(Float32, setup.u_size...) |> aType
@@ -34,7 +35,8 @@
                 additional=Dense(16 => 4), name="Scalar II"),
             (u_size=(64, 3, 5), y_size=(8, 10, 5), out_size=(4, 3, 10, 5),
                 branch=(64, 32, 32, 16), trunk=(8, 8, 8, 16),
-                additional=Dense(16 => 4), name="Vector")]
+                additional=Dense(16 => 4), name="Vector")
+        ]
 
         @testset "Additional layer: $(setup.name)" for setup in setups
             u = rand(Float32, setup.u_size...) |> aType
@@ -50,16 +52,17 @@
             @test setup.out_size == size(pred)
 
             __f = (u, y, ps) -> sum(abs2, first(deeponet((u, y), ps, st)))
-            test_gradients(
-                __f, u, y, ps; atol=1.0f-3, rtol=1.0f-3, skip_backends=[AutoEnzyme()])
+            @test_gradients(__f, u, y, ps; atol=1.0f-3, rtol=1.0f-3)
         end
 
         @testset "Embedding layer mismatch" begin
             u = rand(Float32, 64, 5) |> aType
             y = rand(Float32, 1, 10, 5) |> aType
 
-            deeponet = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 20)),
-                Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
+            deeponet = DeepONet(
+                Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 20)),
+                Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16))
+            )
 
             ps, st = Lux.setup(rng, deeponet) |> dev
             @test_throws ArgumentError deeponet((u, y), ps, st)

test/layers_tests.jl

@@ -9,7 +9,8 @@
             (; m=(10, 10), permuted=Val(false),
                 x_size=(1, 22, 22, 5), y_size=(64, 22, 22, 5)),
             (; m=(10, 10), permuted=Val(true),
-                x_size=(22, 22, 1, 5), y_size=(22, 22, 64, 5))]
+                x_size=(22, 22, 1, 5), y_size=(22, 22, 64, 5))
+        ]
 
         @testset "$(op) $(length(setup.m))D: permuted = $(setup.permuted)" for setup in setups,
             op in opconv
@@ -37,8 +38,8 @@
             end
 
             __f = (x, ps) -> sum(abs2, first(m(x, ps, st)))
-            test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3,
-                skip_backends=[AutoEnzyme(), AutoTracker(), AutoReverseDiff()])
+            @test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3,
+                skip_backends=[AutoTracker(), AutoEnzyme(), AutoReverseDiff()])
         end
     end
 end

test/nomad_tests.jl

@@ -6,7 +6,8 @@
             (u_size=(1, 5), y_size=(1, 5), out_size=(1, 5),
                 approximator=(1, 16, 16, 15), decoder=(16, 8, 4, 1), name="Scalar"),
             (u_size=(8, 5), y_size=(2, 5), out_size=(8, 5),
-                approximator=(8, 32, 32, 16), decoder=(18, 16, 8, 8), name="Vector")]
+                approximator=(8, 32, 32, 16), decoder=(18, 16, 8, 8), name="Vector")
+        ]
 
         @testset "$(setup.name)" for setup in setups
             u = rand(Float32, setup.u_size...) |> aType
@@ -21,8 +22,7 @@
             @test setup.out_size == size(pred)
 
             __f = (u, y, ps) -> sum(abs2, first(nomad((u, y), ps, st)))
-            test_gradients(
-                __f, u, y, ps; atol=1.0f-3, rtol=1.0f-3, skip_backends=[AutoEnzyme()])
+            @test_gradients(__f, u, y, ps; atol=1.0f-3, rtol=1.0f-3)
         end
     end
 end

test/utils_tests.jl

@@ -21,7 +21,8 @@
                 additional=Dense(16 => 4), name="additional : Vector"),
             (b_size=(16, 4, 3, 3, 5), t_size=(16, 10, 5), out_size=(3, 4, 3, 4, 10, 5),
                 additional=Chain(Dense(16 => 4), ReshapeLayer((3, 4, 3, 4, 10))),
-                name="additional : Tensor")]
+                name="additional : Tensor")
+        ]
 
         @testset "project : $(setup.name)" for setup in setups
             b = rand(Float32, setup.b_size...) |> aType