Format code

EnzymeAD · May 28, 2024 · e108c7a · e108c7a
1 parent a4f685c
commit e108c7a
Show file tree

Hide file tree

Showing 26 changed files with 2,550 additions and 1,654 deletions.
diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl
@@ -7,22 +7,25 @@ function __init__()
     for (jlop, hloop) in ((:(NNlib.tanh), :tanh), (:(NNlib.tanh_fast), :tanh))
         @eval begin
             if $jlop != Base.tanh && $jlop != Base.FastMath.tanh_fast
-                function Reactant.elem_apply(::typeof($jlop),
-                                             lhs::Reactant.TracedRArray{ElType,Shape,N}) where {ElType,
-                                                                                                Shape,
-                                                                                                N}
-                    return Reactant.TracedRArray{ElType,Shape,N}((),
-                                                                 Reactant.MLIR.IR.result(Reactant.MLIR.Dialects.stablehlo.$hloop(lhs.mlir_data),
-                                                                                         1))
+                function Reactant.elem_apply(
+                    ::typeof($jlop), lhs::Reactant.TracedRArray{ElType,Shape,N}
+                ) where {ElType,Shape,N}
+                    return Reactant.TracedRArray{ElType,Shape,N}(
+                        (),
+                        Reactant.MLIR.IR.result(
+                            Reactant.MLIR.Dialects.stablehlo.$hloop(lhs.mlir_data), 1
+                        ),
+                    )
                 end
             end
         end
     end
 end
 
 # TODO handle non finite cases
-function NNlib.softmax!(out::Reactant.TracedRArray{T,Shape,N}, x::AbstractArray;
-                        dims=1) where {T,Shape,N}
+function NNlib.softmax!(
+    out::Reactant.TracedRArray{T,Shape,N}, x::AbstractArray; dims=1
+) where {T,Shape,N}
     max_ = NNlib.fast_maximum(x; dims)
     #if all(isfinite, max_)
     @fastmath out .= exp.(x .- max_)