in-place softmax

FluxML · Aug 22, 2017 · 6379100 · 6379100
1 parent 86e0e30
commit 6379100
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/activation.jl b/src/activation.jl
@@ -2,6 +2,9 @@
 
 relu(x) = max(0, x)
 
-softmax(xs::AbstractVector) = exp.(xs) ./ sum(exp.(xs))
+softmax!(out::AbstractVecOrMat, xs::AbstractVecOrMat) =
+  out .= exp.(xs) ./ sum(exp, xs, 1)
 
-softmax(xs::AbstractMatrix) = exp.(xs) ./ sum(exp.(xs), 1)
+softmax!(xs) = softmax!(xs, xs)
+
+softmax(xs) = softmax!(similar(xs), xs)