From 633db5b756b6bcaf088080c030024c8357c856bf Mon Sep 17 00:00:00 2001 From: rusty1s Date: Tue, 31 Oct 2023 08:59:09 +0000 Subject: [PATCH] update --- pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp b/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp index 0bc4cace9..f26ee13af 100644 --- a/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp +++ b/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp @@ -224,12 +224,14 @@ class NeighborSampler { // Case 2: Multinomial sampling: else { at::Tensor index; - if (replace) { // at::multinomial has good perfomance only when - // replace=true, e.g. - // https://github.com/pytorch/pytorch/issues/11931#top + 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 { // An Efficient Algorithm for Biased Sampling: - // https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf + } 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));