[GraphBolt] Slight `gb.expand_indptr` improvement on CPU. (#7189)

8c28f982 · Muhammed Fatih BALIN · GitHub · 290dadb9 · 8c28f982
Unverified Commit 8c28f982 authored Mar 03, 2024 by Muhammed Fatih BALIN Committed by GitHub Mar 03, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

graphbolt/src/expand_indptr.cc graphbolt/src/expand_indptr.cc +1 -1

No files found.
--- a/graphbolt/src/expand_indptr.cc
+++ b/graphbolt/src/expand_indptr.cc
@@ -23,7 +23,7 @@ torch::Tensor ExpandIndptr(
    });
  }
  if (!node_ids.has_value()) {
-    node_ids = torch::arange(indptr.size(0) - 1, indptr.options().dtype(dtype));
+    return torch::repeat_interleave(indptr.diff(), output_size).to(dtype);
  }
  return node_ids.value().to(dtype).repeat_interleave(
      indptr.diff(), 0, output_size);