Remove MKL codepath in replace=False sampling (pyg-team#275)

DamianSzwichtenberg · Nov 9, 2023 · 664acb4 · 664acb4
1 parent 1fda0d7
commit 664acb4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 - Added support for `bfloat16` data type in `segment_matmul` and `grouped_matmul` (CPU only) ([#272](https://github.com/pyg-team/pyg-lib/pull/272))
 ### Changed
+- Dropped the MKL code path when sampling neighbors with `replace=False` since it does not correctly prevent duplicates ([#275](https://github.com/pyg-team/pyg-lib/pull/275))
 - 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

diff --git a/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp b/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp
@@ -174,7 +174,9 @@ class NeighborSampler {
     // Case 3: Sample without replacement:
     else {
       auto index_tracker = IndexTracker<scalar_t>(population);
-      if (population < (1 << 16)) {
+      if (false) {
+        // This logic is currently flawed since `arr` is not guaranteed to
+        // only contain unique values.
         const auto arr =
             std::move(generator.generate_range_of_ints(0, population, count));
         for (auto i = 0; i < arr.size(); ++i) {