Add dims keyword argument to softmax and softmax! (#28)

* Add `dims` keyword argument to `softmax` and `softmax!`

* Bump version

* Apply suggestions from code review

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

* Fix test error

* Update documentation

* Fix

* Extend docstring

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/Manifest.toml

@@ -1,23 +1,58 @@
 # This file is machine-generated - editing it directly is not advised
 
+[[ANSIColoredPrinters]]
+git-tree-sha1 = "574baf8110975760d391c710b6341da1afa48d8c"
+uuid = "a4c015fc-c6ff-483c-b24f-f7ea428134e9"
+version = "0.0.1"
+
+[[ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+
+[[Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
+[[ChainRulesCore]]
+deps = ["Compat", "LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "30ee06de5ff870b45c78f529a6b093b3323256a3"
+uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+version = "1.3.1"
+
+[[Compat]]
+deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
+git-tree-sha1 = "4866e381721b30fac8dda4c8cb1d9db45c8d2994"
+uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
+version = "3.37.0"
+
 [[Dates]]
 deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
+[[DelimitedFiles]]
+deps = ["Mmap"]
+uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+
+[[Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
 [[DocStringExtensions]]
 deps = ["LibGit2"]
 git-tree-sha1 = "a32185f5428d3986f47c2ab78b1f216d5e6cc96f"
 uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 version = "0.8.5"
 
 [[Documenter]]
-deps = ["Base64", "Dates", "DocStringExtensions", "IOCapture", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"]
-git-tree-sha1 = "47f13b6305ab195edb73c86815962d84e31b0f48"
+deps = ["ANSIColoredPrinters", "Base64", "Dates", "DocStringExtensions", "IOCapture", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"]
+git-tree-sha1 = "fe0bc46b27cd3413df55859152fd70e50744025f"
 uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-version = "0.27.3"
+version = "0.27.6"
+
+[[Downloads]]
+deps = ["ArgTools", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 
 [[IOCapture]]
 deps = ["Logging", "Random"]
@@ -36,14 +71,26 @@ version = "0.1.0"
 
 [[JSON]]
 deps = ["Dates", "Mmap", "Parsers", "Unicode"]
-git-tree-sha1 = "81690084b6198a2e1da36fcfda16eeca9f9f24e4"
+git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
 uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-version = "0.21.1"
+version = "0.21.2"
+
+[[LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+
+[[LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
 
 [[LibGit2]]
 deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
+[[LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+
 [[Libdl]]
 uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 
@@ -52,10 +99,10 @@ deps = ["Libdl"]
 uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[LogExpFunctions]]
-deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
+deps = ["ChainRulesCore", "DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
 path = ".."
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
-version = "0.3.0"
+version = "0.3.3"
 
 [[Logging]]
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
@@ -64,17 +111,28 @@ uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 deps = ["Base64"]
 uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
+[[MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+
 [[Mmap]]
 uuid = "a63ad114-7e13-5084-954f-fe012c677804"
 
+[[MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+
 [[NetworkOptions]]
 uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
 
 [[Parsers]]
 deps = ["Dates"]
-git-tree-sha1 = "c8abc88faa3f7a3950832ac5d6e690881590d6dc"
+git-tree-sha1 = "438d35d2d95ae2c5e8780b330592b6de8494e779"
 uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
-version = "1.1.0"
+version = "2.0.3"
+
+[[Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [[Printf]]
 deps = ["Unicode"]
@@ -94,12 +152,48 @@ uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
+[[SharedArrays]]
+deps = ["Distributed", "Mmap", "Random", "Serialization"]
+uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
+
 [[Sockets]]
 uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
 
+[[SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+
+[[Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+
 [[Test]]
 deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[[UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
 [[Unicode]]
 uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+
+[[nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+
+[[p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"

src/basicfuns.jl

@@ -235,15 +235,35 @@ function logsubexp(x::Real, y::Real)
 end
 
 """
-$(SIGNATURES)
+    softmax!(r::AbstractArray{<:Real}, x::AbstractArray{<:Real}=r; dims=:)
+
+Overwrite `r` with the
+[softmax transformation](https://en.wikipedia.org/wiki/Softmax_function) of `x` over
+dimension `dims`.
 
-Overwrite `r` with the `softmax` (or _normalized exponential_) transformation of `x`
+That is, `r` is overwritten with `exp.(x)`, normalized to sum to 1 over the given
+dimensions.
 
-That is, `r` is overwritten with `exp.(x)`, normalized to sum to 1.
+See also: [`softmax`](@ref)
+"""
+softmax!(r::AbstractArray{<:Real}, x::AbstractArray{<:Real}=r; dims=:) =
+    _softmax!(r, x, dims)
 
-See the [Wikipedia entry](https://en.wikipedia.org/wiki/Softmax_function)
 """
-function softmax!(r::AbstractArray{<:Real}, x::AbstractArray{<:Real})
+    softmax(x::AbstractArray{<:Real}; dims=:)
+
+Return the
+[softmax transformation](https://en.wikipedia.org/wiki/Softmax_function) of `x` over
+dimension `dims`.
+
+That is, return `exp.(x)`, normalized to sum to 1 over the given dimensions.
+
+See also: [`softmax!`](@ref)
+"""
+softmax(x::AbstractArray{<:Real}; dims=:) =
+    softmax!(similar(x, float(eltype(x))), x; dims=dims)
+
+function _softmax!(r, x, ::Colon)
     length(r) == length(x) || throw(DimensionMismatch("inconsistent array lengths"))
     u = maximum(x)
     map!(r, x) do xi
@@ -253,18 +273,16 @@ function softmax!(r::AbstractArray{<:Real}, x::AbstractArray{<:Real})
     return r
 end
 
-"""
-$(SIGNATURES)
-
-Return the [`softmax transformation`](https://en.wikipedia.org/wiki/Softmax_function)
-applied to `x` *in place*.
-"""
-softmax!(x::AbstractArray{<:AbstractFloat}) = softmax!(x, x)
-
-"""
-$(SIGNATURES)
-
-Return the [`softmax transformation`](https://en.wikipedia.org/wiki/Softmax_function)
-applied to `x`.
-"""
-softmax(x::AbstractArray{<:Real}) = softmax!(similar(x, float(eltype(x))), x)
+function _softmax!(r, x, dims)
+    size(r) == size(x) || throw(DimensionMismatch("inconsistent array sizes"))
+    u = maximum(x; dims=dims)
+    r .= exp.(x .- u)
+    if u isa Array{eltype(r)}
+        # array can be reused
+        sum!(u, r)
+        r ./= u
+    else
+        r ./= sum(r; dims=dims)
+    end
+    return r
+end

src/chainrules.jl

@@ -34,27 +34,36 @@ function ChainRulesCore.rrule(::typeof(logsumexp), x::AbstractArray{<:Real}; dim
     return Ω, logsumexp_pullback
 end
 
-function ChainRulesCore.frule(
-    (_, _, Δx), ::typeof(softmax!), r::AbstractArray{<:Real}, x::AbstractArray{<:Real},
-)
-    softmax!(r, x)
-    _Δx = reshape(Δx, size(r))
-    Δr = r .* (_Δx .- LinearAlgebra.dot(r, _Δx))
-    return r, Δr
+# no rules for mutating functions currently:
+# https://juliadiff.org/ChainRulesCore.jl/stable/writing_good_rules.html#Which-functions-need-rules?
+function ChainRulesCore.frule((_, Δx), ::typeof(softmax), x::AbstractArray{<:Real}; dims=:)
+    Ω = softmax(x; dims=dims)
+    ΔΩ = if dims === (:)
+        Ω .* (Δx .- LinearAlgebra.dot(Ω, Δx))
+    else
+        ΩΔx = Ω .* Δx
+        ΩΔx .- Ω .* sum(ΩΔx; dims=dims)
+    end
+    return Ω, ΔΩ
 end
-function ChainRulesCore.rrule(
-    ::typeof(softmax!), r::AbstractArray{<:Real}, x::AbstractArray{<:Real},
-)
-    softmax!(r, x)
+function ChainRulesCore.rrule(::typeof(softmax), x::AbstractArray{<:Real}; dims=:)
+    Ω = softmax(x; dims=dims)
+    Ωcopy = copy(Ω)
     project_x = ChainRulesCore.ProjectTo(x)
-    rcopy = copy(reshape(r, size(x)))
-    function softmax!_pullback(r̄)
-        _r̄ = reshape(r̄, size(rcopy))
-        x̄ = ChainRulesCore.InplaceableThunk(
-            Δ -> Δ .+= rcopy .* (_r̄ .- LinearAlgebra.dot(rcopy, _r̄)),
-            ChainRulesCore.@thunk(project_x(rcopy .* (_r̄ .- LinearAlgebra.dot(rcopy, _r̄)))),
-        )
-        return ChainRulesCore.NoTangent(), ChainRulesCore.ZeroTangent(), x̄
+    function softmax_pullback(Ω̄)
+        x̄ = if dims === (:)
+            ChainRulesCore.InplaceableThunk(
+                Δ -> Δ .+= Ωcopy .* (Ω̄ .- LinearAlgebra.dot(Ωcopy, Ω̄)),
+                ChainRulesCore.@thunk(project_x(Ωcopy .* (Ω̄ .- LinearAlgebra.dot(Ωcopy, Ω̄)))),
+            )
+        else
+            ΩcopyΩ̄  = Ωcopy .* Ω̄
+            ChainRulesCore.InplaceableThunk(
+                Δ -> Δ .+= ΩcopyΩ̄  .- Ωcopy .* sum(ΩcopyΩ̄; dims=dims),
+                ChainRulesCore.@thunk(project_x(ΩcopyΩ̄  .- Ωcopy .* sum(ΩcopyΩ̄; dims=dims))),
+            )
+        end
+        return ChainRulesCore.NoTangent(), x̄
     end
-    return r, softmax!_pullback
+    return Ω, softmax_pullback
 end

test/basicfuns.jl

@@ -207,9 +207,23 @@ end
         softmax!(s, x)
         @test s ≈ r
 
+        fill!(s, zero(T))
+        softmax!(s, x; dims=1)
+        @test s ≈ r
+
         s = Matrix{T}(undef, 1, 3)
         softmax!(s, x)
         @test s ≈ permutedims(r)
+
+        @test_throws DimensionMismatch softmax!(s, x; dims=1)
+
+        fill!(s, zero(T))
+        softmax!(s, permutedims(x); dims=2)
+        @test s ≈ permutedims(r)
+
+        fill!(s, zero(T))
+        softmax!(s, permutedims(x); dims=1:2)
+        @test s ≈ permutedims(r)
     end
     softmax!(x)
     @test x ≈ r
@@ -219,6 +233,21 @@ end
         s = softmax(x)
         @test s ≈ r
         @test eltype(s) === T
+
+        x = repeat(S[1, 2, 3], 1, 3)
+        s = softmax(x; dims=1)
+        @test s ≈ repeat(r, 1, 3)
+        @test eltype(s) === T
+
+        x = repeat(S[1 2 3], 3, 1)
+        s = softmax(x; dims=2)
+        @test s ≈ repeat(permutedims(r), 3, 1)
+        @test eltype(s) === T
+
+        x = S[1 2 3]
+        s = softmax(x; dims=1:2)
+        @test s ≈ permutedims(r)
+        @test eltype(s) === T
     end
 
     x = [1//2, 2//3, 3//4]

test/chainrules.jl

@@ -76,9 +76,13 @@
     end
 
     for x in (randn(10), randn(10, 8))
-        for r in (similar(x), similar(x, 1, size(x)...))
-            test_frule(softmax!, r, x)
-            test_rrule(softmax!, r, x)
+        test_frule(softmax, x)
+        test_rrule(softmax, x)
+
+        for dims in (1, 1:2, 2)
+            all(d <= ndims(x) for d in dims) || continue
+            test_frule(softmax, x; fkwargs=(dims=dims,))
+            test_rrule(softmax, x; fkwargs=(dims=dims,))
         end
     end
 end