refactor: remove more elem_apply

EnzymeAD · Aug 4, 2024 · d0e1dbd · d0e1dbd
1 parent 3d93c66
commit d0e1dbd
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl
@@ -8,17 +8,13 @@ for (jlop, hloop) in (
     (:(NNlib.sigmoid_fast), :logistic),
     (:(NNlib.sigmoid), :logistic),
 )
-    @eval begin
-        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
+    @eval function $(jlop)(x::Reactant.TracedRArray{T,(),0}) where {T}
+        return Reactant.TracedRArray{T,(),0}(
+            (),
+            Reactant.MLIR.IR.result(
+                Reactant.MLIR.Dialects.stablehlo.$(hloop)(x.mlir_data), 1
+            ),
+        )
     end
 end