Fix edge weight normalization

Author: pbielak

csrc/cpu/rw_cpu.cpp

@@ -148,10 +148,18 @@ void compute_cdf(const int64_t *rowptr, const float_t *edge_weight,
   at::parallel_for(0, numel - 1, at::internal::GRAIN_SIZE, [&](int64_t begin, int64_t end) {
     for(int64_t i = begin; i < end; i++) {
       int64_t row_start = rowptr[i], row_end = rowptr[i + 1];
+
+      // Compute sum to normalize weights
+      float_t sum = 0.0;
+
+      for(int64_t j = row_start; j < row_end; j++) {
+        sum += edge_weight[j];
+      }
+
       float_t acc = 0.0;
 
       for(int64_t j = row_start; j < row_end; j++) {
-        acc += edge_weight[j];
+        acc += edge_weight[j] / sum;
         edge_weight_cdf[j] = acc;
       }
     }

csrc/cuda/rw_cuda.cu

@@ -159,10 +159,16 @@ __global__ void cdf_kernel(const int64_t *rowptr, const float_t *edge_weight,
   if (thread_idx < numel - 1) {
     int64_t row_start = rowptr[thread_idx], row_end = rowptr[thread_idx + 1];
 
+    float_t sum = 0.0;
+
+    for(int64_t i = row_start; i < row_end; i++) {
+      sum += edge_weight[i];
+    }
+
     float_t acc = 0.0;
 
     for(int64_t i = row_start; i < row_end; i++) {
-      acc += edge_weight[i];
+      acc += edge_weight[i] / sum;
       edge_weight_cdf[i] = acc;
     }
   }

torch_cluster/rw.py

@@ -62,9 +62,6 @@ def random_walk(
             rowptr, col, start, walk_length, p, q,
         )
     else:
-        # Normalize edge weights by node degrees
-        edge_weight = edge_weight / deg[row]
-
         node_seq, edge_seq = torch.ops.torch_cluster.random_walk_weighted(
             rowptr, col, edge_weight, start, walk_length, p, q,
         )