use broadcasting

Author: jumerckx

src/softmax.jl

@@ -19,15 +19,9 @@ independent.
 """
 softmax(xs; dims=1) = softmax!(similar(xs), xs, dims)
 
-function softmax!(out::AbstractVecOrMat{T}, xs::AbstractVecOrMat{T}, dims) where {T}
-    # First, store column-wise maximum in the last element of `out`
-    maxdims = ntuple(d -> (d in dims) ? (1:1) : (:), ndims(xs))
-    out[maxdims...] = maximum!(out[maxdims...], xs)
-
-    # Subtract the column-wise maximums to normalize, take exp()
-    out .= exp.(xs.-out[maxdims...])
-
-    # Normalize by sum of the entire thing
+function softmax!(out::AbstractArray{T}, xs::AbstractArray{T}, dims) where {T}
+    max = maximum(xs, dims)
+    out .= exp.(xs.-max)
     out ./= sum(out, dims=dims)
     return out
 end