chore: fmt

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

test/nn/luxlib.jl

@@ -30,30 +30,26 @@ using LuxLib, Reactant, Enzyme, NNlib
         x_ra = Reactant.to_rarray(x)
         bias_ra = Reactant.to_rarray(bias)
 
-        f_compile = Reactant.with_config(;
+        y_compile = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile fused_dense_bias_activation(act, weight_ra, x_ra, bias_ra)
+            @jit fused_dense_bias_activation(act, weight_ra, x_ra, bias_ra)
         end
 
         y_res = fused_dense_bias_activation(act, weight, x, bias)
-        y_compile = f_compile(act, weight_ra, x_ra, bias_ra)
 
         @test y_res ≈ y_compile atol = 1e-5 rtol = 1e-2
 
         @testset "Enzyme: fused_dense_bias_activation" begin
             dw, dx, db = ∇fuseddense(act, weight, x, bias)
 
-            ∇fuseddense_compiled = Reactant.with_config(;
+            dw_compile, dx_compile, db_compile = Reactant.with_config(;
                 dot_general_precision=PrecisionConfig.HIGHEST,
                 convolution_precision=PrecisionConfig.HIGHEST,
             ) do
-                @compile ∇fuseddense(act, weight_ra, x_ra, bias_ra)
+                @jit ∇fuseddense(act, weight_ra, x_ra, bias_ra)
             end
-            dw_compile, dx_compile, db_compile = ∇fuseddense_compiled(
-                act, weight_ra, x_ra, bias_ra
-            )
 
             @test dw ≈ dw_compile atol = 1e-5 rtol = 1e-2
             @test dx ≈ dx_compile atol = 1e-5 rtol = 1e-2
@@ -103,51 +99,47 @@ end
         x_ra = Reactant.to_rarray(x)
         b_ra = Reactant.to_rarray(b)
 
-        f_compile = Reactant.with_config(;
+        y_compile = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile biasact(act, x_ra, b_ra)
+            @jit biasact(act, x_ra, b_ra)
         end
 
-        f_compile!! = Reactant.with_config(;
+        y_compile!! = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile biasact!!(act, x_ra, b_ra)
+            @jit biasact!!(act, x_ra, b_ra)
         end
 
         y_simple = biasact(act, x, b)
         y_simple!! = biasact!!(act, x, b)
-        y_compile = f_compile(act, x_ra, b_ra)
-        y_compile!! = f_compile!!(act, x_ra, b_ra)
 
         @test y_simple ≈ y_compile atol = 1e-5 rtol = 1e-2
         @test y_simple!! ≈ y_compile!! atol = 1e-5 rtol = 1e-2
 
         @testset "Enzyme: bias_activation" begin
             ∂x_enz, ∂b_enz = ∇biasact(act, x, b)
-            ∇biasact_compiled = Reactant.with_config(;
+            ∂x_compile, ∂b_compile = Reactant.with_config(;
                 dot_general_precision=PrecisionConfig.HIGHEST,
                 convolution_precision=PrecisionConfig.HIGHEST,
             ) do
-                @compile ∇biasact(act, x_ra, b_ra)
+                @jit ∇biasact(act, x_ra, b_ra)
             end
-            ∂x_compile, ∂b_compile = ∇biasact_compiled(act, x_ra, b_ra)
 
             @test ∂x_enz ≈ ∂x_compile atol = 1e-5 rtol = 1e-2
             @test ∂b_enz ≈ ∂b_compile atol = 1e-5 rtol = 1e-2
         end
 
         @testset "Enzyme: bias_activation!!" begin
             ∂x_enz!!, ∂b_enz!! = ∇biasact!!(act, x, b)
-            ∇biasact!!_compiled = Reactant.with_config(;
+            ∂x_compile!!, ∂b_compile!! = Reactant.with_config(;
                 dot_general_precision=PrecisionConfig.HIGHEST,
                 convolution_precision=PrecisionConfig.HIGHEST,
             ) do
-                @compile ∇biasact!!(act, x_ra, b_ra)
+                @jit ∇biasact!!(act, x_ra, b_ra)
             end
-            ∂x_compile!!, ∂b_compile!! = ∇biasact!!_compiled(act, x_ra, b_ra)
 
             @test ∂x_enz!! ≈ ∂x_compile!! atol = 1e-5 rtol = 1e-2
             @test ∂b_enz!! ≈ ∂b_compile!! atol = 1e-5 rtol = 1e-2
@@ -178,24 +170,21 @@ end
     @testset "Activation: $act" for act in (
         identity, relu, sigmoid, tanh, tanh_fast, sigmoid_fast, gelu, abs2
     )
-        f_compile = Reactant.with_config(;
+        y_simple = sumabs2(act, x_act)
+        y_simple!! = sumabs2!!(act, x_act)
+        y_compile = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile sumabs2(act, x_act_ca)
+            @jit sumabs2(act, x_act_ca)
         end
-        f_compile!! = Reactant.with_config(;
+        y_compile!! = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile sumabs2!!(act, x_act_ca)
+            @jit sumabs2!!(act, x_act_ca)
         end
 
-        y_simple = sumabs2(act, x_act)
-        y_simple!! = sumabs2!!(act, x_act)
-        y_compile = f_compile(act, x_act_ca)
-        y_compile!! = f_compile!!(act, x_act_ca)
-
         @test y_simple ≈ y_compile atol = 1e-5 rtol = 1e-2
         @test y_simple!! ≈ y_compile!! atol = 1e-5 rtol = 1e-2
 
@@ -205,21 +194,19 @@ end
         ∂x_enz!! = Enzyme.make_zero(x_act)
         Enzyme.autodiff(Reverse, sumabs2!!, Active, Const(act), Duplicated(x_act, ∂x_enz!!))
 
-        ∇sumabs2 = Reactant.with_config(;
+        ∂x_compile = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile ∇sumabs2(act, x_act_ca)
+            @jit ∇sumabs2(act, x_act_ca)
         end
-        ∂x_compile = ∇sumabs2(act, x_act_ca)
 
-        ∇sumabs2!! = Reactant.with_config(;
+        ∂x_compile!! = Reactant.with_config(;
             dot_general_precision=PrecisionConfig.HIGHEST,
             convolution_precision=PrecisionConfig.HIGHEST,
         ) do
-            @compile ∇sumabs2!!(act, x_act_ca)
+            @jit ∇sumabs2!!(act, x_act_ca)
         end
-        ∂x_compile!! = ∇sumabs2!!(act, x_act_ca)
 
         @test ∂x_enz ≈ ∂x_compile atol = 1e-5 rtol = 1e-2
         @test ∂x_enz!! ≈ ∂x_compile!! atol = 1e-5 rtol = 1e-2
@@ -242,16 +229,15 @@ end
 
             conv_dims = DenseConvDims(x, weight; stride, padding, dilation, groups)
 
-            fused_conv_compiled = Reactant.with_config(;
+            reactant_res = Reactant.with_config(;
                 dot_general_precision=PrecisionConfig.HIGHEST,
                 convolution_precision=PrecisionConfig.HIGHEST,
             ) do
-                @compile fused_conv_bias_activation(act, weight_reactant, x_reactant, bias_reactant, conv_dims)
+                @jit fused_conv_bias_activation(
+                    act, weight_reactant, x_reactant, bias_reactant, conv_dims
+                )
             end
 
-            reactant_res = fused_conv_compiled(
-                act, weight_reactant, x_reactant, bias_reactant, conv_dims
-            )
             luxlib_res = fused_conv_bias_activation(act, weight, x, bias, conv_dims)
 
             @test reactant_res ≈ luxlib_res atol = 1e-5 rtol = 1e-2