Merge pull request #17 from ToucheSir/bc/test_softmax_dims

Test `softmax` with varying dimensions

test/softmax.jl

@@ -1,12 +1,12 @@
 @testset "softmax" begin
-    for dims in [(5,5), (5,)]
-        x = randn(Float64, dims)
-        y = softmax(x)
-        dy = randn(Float64, dims)
-        gputest(softmax, x)
-        gputest(∇softmax, dy, x, y, checkgrad=false)
-        y = logsoftmax(x)
-        gputest(logsoftmax, x)
-        gputest(∇logsoftmax, dy, x, y, checkgrad=false) 
+    for (sz, dims) in [((5,), :), ((5,), 1), ((5,5), :), ((5,5), 1), ((5,5), 2), ((5,5,5,5), (2,3)), ((5,5,5,5), (2,4))]
+        x = randn(Float64, sz)
+        y = softmax(x, dims=dims)
+        dy = randn(Float64, sz)
+        gputest(softmax, x, dims=dims)
+        gputest(∇softmax, dy, x, y, dims=dims, checkgrad=false)
+        y = logsoftmax(x, dims=dims)
+        gputest(logsoftmax, x, dims=dims)
+        gputest(∇logsoftmax, dy, x, y, dims=dims, checkgrad=false) 
     end
 end

test/test_utils.jl

@@ -1,9 +1,9 @@
-function gputest(f, xs...; checkgrad=true, atol=1e-10)
+function gputest(f, xs...; checkgrad=true, atol=1e-10, kws...)
     cpu_in = xs
     gpu_in = CuArray.(xs)
 
-    cpu_out = f(cpu_in...)
-    gpu_out = f(gpu_in...)
+    cpu_out = f(cpu_in...; kws...)
+    gpu_out = f(gpu_in...; kws...)
     @test collect(cpu_out) ≈ collect(gpu_out)
 
     if checkgrad