try to unify the scatter API

Author: yuehhua

lib/NNlibCUDA/src/scatter.jl

@@ -1,7 +1,7 @@
 ATM_OPS = Dict((+) => CUDA.atomic_add!, (-) => CUDA.atomic_sub!, (max) => CUDA.atomic_max!, (min) => CUDA.atomic_min!,
                (*) => CUDA.atomic_mul!, (/) => CUDA.atomic_div!, (&) => CUDA.atomic_and!, (|) => CUDA.atomic_or!)
 
-function scatter!(op, dst::CuArray, src::CuArray, idx::CuArray{IntOrIntTuple})
+function scatter!(op, dst::CuArray{Tdst}, src::CuArray{Tsrc}, idx::CuArray{<:IntOrIntTuple}, dims::Val{N}) where {Tdst,Tsrc,N}
     function kernel!(atm_op, dst, src, idx)
         li = threadIdx().y + (blockIdx().y - 1) * blockDim().y
         i = threadIdx().x + (blockIdx().x - 1) * blockDim().x

lib/NNlibCUDA/test/runtests.jl

@@ -3,6 +3,7 @@ using NNlib
 using Zygote
 using NNlibCUDA
 using ForwardDiff: Dual
+using Statistics: mean
 using CUDA
 CUDA.allowscalar(false)