Fix edge weight normalization

rusty1s · Sep 9, 2022 · bcda1c7 · bcda1c7
1 parent c4154c8
commit bcda1c7
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/torch_cluster/rw.py b/torch_cluster/rw.py
@@ -62,8 +62,17 @@ def random_walk(
             rowptr, col, start, walk_length, p, q,
         )
     else:
-        # Normalize edge weights by node degrees
-        edge_weight = edge_weight / deg[row]
+        # Normalize edge weights
+        from torch_sparse import SparseTensor
+
+        adj = SparseTensor(
+            row=row,
+            col=col,
+            value=edge_weight,
+            sparse_sizes=(num_nodes, num_nodes),
+        )
+
+        edge_weight = edge_weight / adj.sum(dim=1).repeat_interleave(deg)
 
         node_seq, edge_seq = torch.ops.torch_cluster.random_walk_weighted(
             rowptr, col, edge_weight, start, walk_length, p, q,