Optimization of biased sampling (#270)

Implemented biased sampling using this algorithm:
https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf

![image](https://github.com/pyg-team/pyg-lib/assets/57872493/48ca46ab-cb50-4286-9f09-1443863bb08f)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: rusty1s <[email protected]>

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 ### Changed
 - Added `--biased` parameter to run benchmarks for biased sampling ([#267](https://github.com/pyg-team/pyg-lib/pull/267))
+- Improved speed of biased sampling ([#270](https://github.com/pyg-team/pyg-lib/pull/270))
 ### Removed
 
 ## [0.3.0] - 2023-10-11

pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp

@@ -223,7 +223,19 @@ class NeighborSampler {
 
     // Case 2: Multinomial sampling:
     else {
-      const auto index = at::multinomial(weight, count, replace);
+      at::Tensor index;
+      if (replace) {
+        // at::multinomial only has good perfomance for `replace=true`, see:
+        // https://github.com/pytorch/pytorch/issues/11931
+        index = at::multinomial(weight, count, replace);
+      } else {
+        // For `replace=false`, we make use of the implementation of the
+        // "Weighted Random Sampling" paper:
+        // https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf
+        const auto rand = at::empty_like(weight).uniform_();
+        const auto key = (rand.log() / weight);
+        index = std::get<1>(key.topk(count));
+      }
       const auto index_data = index.data_ptr<int64_t>();
       for (size_t i = 0; i < index.numel(); ++i) {
         add(row_start + index_data[i], global_src_node, local_src_node,