chore: apply formatting suggestion

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

Author: avik-pal

ext/ReactantNNlibExt.jl

@@ -27,7 +27,6 @@ function NNlib.gelu(x::Reactant.TracedRArray{T,0}) where {T}
     return x * sigmoid(λλ * x * muladd(x^2, α, one(T)))
 end
 
-
 # TODO handle non finite cases
 function NNlib.softmax!(
     out::Reactant.TracedRArray{T,N}, x::AbstractArray; dims=1

test/nn/lux.jl

@@ -1,5 +1,34 @@
 using Reactant, Lux, Random, Statistics, Enzyme, Functors, OneHotArrays
 
+# Lux.Exprimental.TrainState is very specialized for Lux models, so we write out the
+# training loop manually:
+function crossentropy(ŷ, y)
+    logŷ = log.(ŷ)
+    result = y .* logŷ
+    return -sum(result)
+end
+
+function loss_function(model, x, y, ps, st)
+    y_hat, _ = model(x, ps, st)
+    # return CrossEntropyLoss()(y_hat, y)
+    return crossentropy(y_hat, y)
+end
+
+function gradient_loss_function(model, x, y, ps, st)
+    dps = Enzyme.make_zero(ps)
+    _, res = Enzyme.autodiff(
+        ReverseWithPrimal,
+        loss_function,
+        Active,
+        Const(model),
+        Const(x),
+        Const(y),
+        Duplicated(ps, dps),
+        Const(st),
+    )
+    return res, dps
+end
+
 @testset "Lux.jl Integration" begin
     # Generate some data for the XOR problem: vectors of length 2, as columns of a matrix:
     noisy = rand(Float32, 2, 1000)                                        # 2×1000 Matrix{Float32}
@@ -33,36 +62,6 @@ using Reactant, Lux, Random, Statistics, Enzyme, Functors, OneHotArrays
     ctarget = Reactant.ConcreteRArray(Array{Float32}(target))
     # ctarget = Reactant.to_rarray(target)
 
-    # Lux.Exprimental.TrainState is very specialized for Lux models, so we write out the
-    # training loop manually:
-    function crossentropy(ŷ, y)
-        logŷ = log.(ŷ)
-        result = y .* logŷ
-        # result = ifelse.(y .== 0.0f0, zero.(result), result)
-        return -sum(result)
-    end
-
-    function loss_function(model, x, y, ps, st)
-        y_hat, _ = model(x, ps, st)
-        # return CrossEntropyLoss()(y_hat, y)
-        return crossentropy(y_hat, y)
-    end
-
-    function gradient_loss_function(model, x, y, ps, st)
-        dps = Enzyme.make_zero(ps)
-        _, res = Enzyme.autodiff(
-            ReverseWithPrimal,
-            loss_function,
-            Active,
-            Const(model),
-            Const(x),
-            Const(y),
-            Duplicated(ps, dps),
-            Const(st),
-        )
-        return res, dps
-    end
-
     res, dps = gradient_loss_function(model, noisy, target, ps, st)
 
     compiled_gradient = Reactant.compile(