Allow general indices in softmax! (#17)

Co-authored-by: Milan Bouchet-Valat <[email protected]>

Project.toml

@@ -1,17 +1,19 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
-julia = "1"
 DocStringExtensions = "0.8"
+julia = "1"
 
 [extras]
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["OffsetArrays", "Test"]

src/LogExpFunctions.jl

@@ -2,6 +2,7 @@ module LogExpFunctions
 
 using DocStringExtensions: SIGNATURES
 using Base: Math.@horner, @irrational
+import LinearAlgebra
 
 export loghalf, logtwo, logπ, log2π, log4π
 export xlogx, xlogy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,

src/basicfuns.jl

@@ -233,17 +233,12 @@ That is, `r` is overwritten with `exp.(x)`, normalized to sum to 1.
 See the [Wikipedia entry](https://en.wikipedia.org/wiki/Softmax_function)
 """
 function softmax!(r::AbstractArray{<:Real}, x::AbstractArray{<:Real})
-    n = length(x)
-    length(r) == n || throw(DimensionMismatch("Inconsistent array lengths."))
+    length(r) == length(x) || throw(DimensionMismatch("inconsistent array lengths"))
     u = maximum(x)
-    s = zero(eltype(r))
-    @inbounds for i = 1:n
-        s += (r[i] = exp(x[i] - u))
-    end
-    invs = inv(s)
-    @inbounds for i = 1:n
-        r[i] *= invs
+    map!(r, x) do xi
+        return exp(xi - u)
     end
+    LinearAlgebra.lmul!(inv(sum(r)), r)
     return r
 end
 

test/basicfuns.jl

@@ -149,22 +149,37 @@ end
 @testset "softmax" begin
     x = [1.0, 2.0, 3.0]
     r = exp.(x) ./ sum(exp.(x))
-    @test softmax(x) ≈ r
+
+    # in-place versions
+    for T in (Float32, Float64)
+        s = Vector{T}(undef, 3)
+        softmax!(s, x)
+        @test s ≈ r
+
+        s = Matrix{T}(undef, 1, 3)
+        softmax!(s, x)
+        @test s ≈ permutedims(r)
+    end
     softmax!(x)
     @test x ≈ r
-
-    x = [1, 2, 3]
-    r = exp.(x) ./ sum(exp.(x))
-    @test softmax(x) ≈ r
-    @test eltype(softmax(x)) == Float64
-
+
+    for (S, T) in ((Int, Float64), (Float64, Float64), (Float32, Float32))
+        x = S[1, 2, 3]
+        s = softmax(x)
+        @test s ≈ r
+        @test eltype(s) === T
+    end
+
     x = [1//2, 2//3, 3//4]
     r = exp.(x) ./ sum(exp.(x))
-    @test softmax(x) ≈ r
-    @test eltype(softmax(x)) == Float64
+    s = softmax(x)
+    @test s ≈ r
+    @test eltype(s) === Float64
 
-    x = Float32[1, 2, 3]
-    r = exp.(x) ./ sum(exp.(x))
-    @test softmax(x) ≈ r
-    @test eltype(softmax(x)) == Float32
+    # non-standard indices: #12
+    x = OffsetArray(1:3, -2:0)
+    s = softmax(x)
+    @test s isa OffsetArray{Float64}
+    @test axes(s, 1) == OffsetArrays.IdOffsetRange(-2:0)
+    @test collect(s) ≈ softmax(1:3)
 end

test/runtests.jl

@@ -1,3 +1,5 @@
-using LogExpFunctions, Test
+using LogExpFunctions
+using OffsetArrays
+using Test
 
 include("basicfuns.jl")