Merge pull request #48 from DomCRose/activation_eltype_restriction

Restrict element type of activation overrides to CUDNN datatypes

Author: ToucheSir

ext/NNlibCUDA/src/cudnn/activations.jl

@@ -18,13 +18,13 @@ for (f, op) in [
     @eval begin
         # in-place
         function Base.materialize!(dst::DenseCuArray{<:CUDNNFloat},
-                                   bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray}})
+                                   bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray{<:CUDNNFloat}}})
             $op(bc.args[1], dst)
             return dst
         end
 
         # out of place
-        function Base.materialize(bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray}})
+        function Base.materialize(bc::Broadcast.Broadcasted{<:Any,<:Any,typeof($f),<:Tuple{DenseCuArray{<:CUDNNFloat}}})
             ElType = Broadcast.combine_eltypes(bc.f, bc.args)
             dst = similar(bc, ElType)
             $op(bc.args[1], dst)

ext/NNlibCUDA/test/activations.jl

@@ -7,9 +7,20 @@
 end
 
 @testset "forward diff" begin
-    f(x) = logσ.(x)
-    ds = Dual.(rand(5),1)
-    @test f(ds) ≈ collect(f(CuArray(ds)))
+    for f in NNlib.ACTIVATIONS
+        if f ∉ [:rrelu]
+            @eval gputest(x -> $f.(x), Dual.(rand(5), 1))
+        end
+    end
+end
+
+# Broadcasting over complex CuArray works without NNlibCUDA, this test checks that
+# NNlibCUDA does not cause such operations to take a fast path which does not support
+# complex numbers (e.g. CUDNN)
+@testset "complex" begin
+    f(x) = tanh.(x)
+    cs = rand(ComplexF64, 5)
+    @test f(cs) ≈ collect(f(CuArray(cs)))
 end
 
 @testset "softplus" begin