added a testset for dims parameter.

FluxML · Dec 20, 2020 · 6d48ff8 · 6d48ff8
1 parent 39970e8
commit 6d48ff8
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/test/softmax.jl b/test/softmax.jl
@@ -5,6 +5,14 @@ using Statistics: mean
     @test softmax(Int[0, 0]) == [0.5, 0.5]
 end
 
+@testset "softmax dims" begin
+    xs = rand(fill(2,5)...)
+    out = similar(xs)
+    for (fn!, fn) in [(softmax!, softmax), (logsoftmax!, logsoftmax)], i = 1:ndims(xs)
+        @test fn!(out, xs; dims = i) == fn(xs; dims = i)
+    end
+end
+
 @testset "softmax" begin
     xs = rand(5, 5)
     @test all(sum(softmax(xs), dims = 1) .≈ 1)