Unify binarize function signature

Author: dferre97

GNNGraphs/ext/GNNGraphsCUDAExt.jl

@@ -21,13 +21,9 @@ GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz
 
 GNNGraphs.iscuarray(x::AnyCuArray) = true
 
-function GNNGraphs.binarize(Mat::CUSPARSE.CuSparseMatrixCSC)
-    bin_vals = fill!(similar(nonzeros(Mat), Bool), true)
-    return CUSPARSE.CuSparseMatrixCSC(Mat.colPtr, rowvals(Mat), bin_vals, size(Mat))
-end
-function GNNGraphs.binarize(Mat::CUSPARSE.CuSparseMatrixCSC, T::DataType)
+function GNNGraphs.binarize(Mat::CUSPARSE.CuSparseMatrixCSC, T::DataType = Bool)
+    @debug "Binarizing sparse matrix of type $(typeof(Mat)) to type $(T)"
     bin_vals = fill!(similar(nonzeros(Mat)), one(T))
-    # Binarize a CuSparseMatrixCSC by setting all nonzero values to one(T)
     return CUSPARSE.CuSparseMatrixCSC(Mat.colPtr, rowvals(Mat), bin_vals, size(Mat))
 end
 

GNNGraphs/src/utils.jl

@@ -295,9 +295,7 @@ function _rand_edges(rng, (n1, n2), m)
     return s, t, val
 end
 
-binarize(x) = map(>(0), x)
-# here just to allow CUDA extension to overload this function with correct type casting
-binarize(x, T::DataType) = binarize(x)
+binarize(x, T::DataType = Bool) = ifelse.(x .> 0, one(T), zero(T))
 
 CRC.@non_differentiable binarize(x...)
 CRC.@non_differentiable edge_encoding(x...)