add tests softmax

biaslab · Aug 27, 2024 · 987b420 · 987b420
1 parent ef7decf
commit 987b420
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/test/layers/softmax_tests.jl b/test/layers/softmax_tests.jl
@@ -0,0 +1,16 @@
+@testitem "Softmax" begin
+
+    using UnboundedBNN: Softmax
+    using StatsFuns: softmax
+
+    A = randn(100, 200)
+    layer = Softmax(200)
+    @test all(exp.(layer(A)) .≈ softmax(A, dims = 1))
+
+    @test sum(exp, layer(randn(1000))) ≈ 1.0
+    @test sum(exp, layer(randn(1000, 100))) ≈ 100.0 
+
+    @test all(layer(A, zeros(size(A)))[1] .≈ layer(A))
+    @test layer(A, zeros(size(A)))[2] === nothing
+
+end