add leakyrelu def

JuliaSIMD · May 3, 2024 · dbc1ded · dbc1ded
1 parent 51ee029
commit dbc1ded
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/ext/ForwardDiffExt.jl b/ext/ForwardDiffExt.jl
@@ -377,3 +377,15 @@ for f in (:vmapt, :vmapnt, :vmapntt)
   end
 end
 end
+
+if Base.ifelse !== IfElse.ifelse
+  @inline function NNlib.leakyrelu(x::AbstractSIMD)
+    fx = float(x)
+    NNlib.leakyrelu(fx, convert(typeof(fx), NNlib.leakyrelu_a))
+  end
+  @inline function NNlib.leakyrelu(x::AbstractSIMD, a)
+    fx = float(x)
+    ax = convert(typeof(fx), a * x)
+    ifelse(x > 0, fx, ax)  # max(a*x, x) is 3x slower
+  end
+end