Use float union type alias and remove CRC for now

Author: ToucheSir

ext/NNlibCUDA/Project.toml

@@ -4,14 +4,12 @@ version = "0.1.11"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "1"
 CUDA = "3.3.1"
 NNlib = "0.7.31"
 julia = "1.6"

ext/NNlibCUDA/src/NNlibCUDA.jl

@@ -3,8 +3,6 @@ module NNlibCUDA
 using NNlib
 using CUDA
 using Random, Statistics
-using ChainRulesCore: NoTangent, ZeroTangent
-import ChainRulesCore: rrule
 
 const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}
 

ext/NNlibCUDA/src/cudnn/batchnorm.jl

@@ -44,7 +44,7 @@ function cudnnBNForward!(y::DenseCuArray{T}, g::DenseCuArray{T}, b::DenseCuArray
                         eps = T(1e-5),
                         training = true,
                         affine = true,
-                        track_stats = true) where T<:Union{Float32, Float64}
+                        track_stats = true) where T<:CUDNNFloat
   dims = _wsize(x)
   if eps < CUDNN_BN_MIN_EPSILON
     @warn "eps $eps is too small for CuDNN, setting to CUDNN_BN_MIN_EPSILON=$CUDNN_BN_MIN_EPSILON"
@@ -99,7 +99,7 @@ end
 
 function ∇batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T, 2}, dy::DenseCuArray{T, 2},
             running_mean, running_var, momentum;
-            kws...) where T<:Union{Float32, Float64}
+            kws...) where T<:CUDNNFloat
   dg, db, dx = ∇batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), reshape(dy, 1, 1, size(dy, 1),
                           size(dy, 2)), running_mean, running_var, momentum; kws...)
   (dg, db, dropdims(dx, dims = (1, 2)))
@@ -108,7 +108,7 @@ end
 
 function ∇batchnorm(g::DenseCuArray{T}, b::DenseCuArray{T}, x::DenseCuArray{T}, dy::DenseCuArray{T},
                     running_mean, running_var, momentum;
-                    affine=true, kws...) where T<:Union{Float32, Float64}
+                    affine=true, kws...) where T<:CUDNNFloat
   dg = similar(g)
   db = similar(b)
   dx = similar(x)
@@ -127,7 +127,7 @@ function cudnnBNBackward!(dg::DenseCuArray{T}, g::DenseCuArray{T}, db::DenseCuAr
                           momentum; cache = nothing, eps = T(1e-5),
                           alpha = T(1), beta = T(0),
                           dalpha = T(1), dbeta = T(0), training = true,
-                          track_stats = true) where T<:Union{Float32, Float64}
+                          track_stats = true) where T<:CUDNNFloat
   if !track_stats
     running_mean = CU_NULL
     running_var = CU_NULL