update softmax tests to match NNlib 0.8.3 (#44)

* update tests for NNlib 0.8.3

* lower precision

Author: mcabbott

ext/NNlibCUDA/Project.toml

@@ -1,6 +1,6 @@
 name = "NNlibCUDA"
 uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -11,7 +11,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 CUDA = "3.3.1"
-NNlib = "0.8"
+NNlib = "0.8.3"
 julia = "1.6"
 
 [extras]

ext/NNlibCUDA/test/softmax.jl

@@ -1,12 +1,20 @@
 @testset "softmax" begin
     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)
+
+        y = softmax(x, dims=dims)
         gputest(softmax, x, dims=dims)
-        gputest(∇softmax, dy, x, y, dims=dims, checkgrad=false)
-        y = logsoftmax(x, dims=dims)
+        gputest(NNlib.∇softmax_data, dy, y; dims=dims)
+
+        y2 = logsoftmax(x, dims=dims)
         gputest(logsoftmax, x, dims=dims)
-        gputest(∇logsoftmax, dy, x, y, dims=dims, checkgrad=false) 
+        gputest(NNlib.∇logsoftmax_data, dy, y2; dims=dims)
+
+        # From NNlib 0.8.3, ∇softmax! is not used in the gradient.
+        # But NNlibCUDA still knows how to call CUDNN routines, let's test they agree:
+        @test NNlib.∇softmax_data(dy, y; dims=dims) ≈ collect(∇softmax!(similar(cu(x)), cu(dy), cu(x), cu(y); dims=dims)) atol=1e-4
+        @test NNlib.∇logsoftmax_data(dy, y2; dims=dims) ≈ collect(∇logsoftmax!(similar(cu(x)), cu(dy), cu(x), cu(y2); dims=dims)) atol=1e-4
+        # (Note that ∇softmax! does not depend on x, it's just there to disambiguate from an even older signature.)
     end
 end