outsource

aten/cpu/degree.cpp

+#ifndef DEGREE_CPP
+#define DEGREE_CPP
+
+#include <torch/torch.h>
+
+inline at::Tensor degree(at::Tensor index, int num_nodes,
+                         at::ScalarType scalar_type) {
+  auto zero = at::full(torch::CPU(scalar_type), {num_nodes}, 0);
+  auto one = at::full(zero.type(), {index.size(0)}, 1);
+  return zero.scatter_add_(0, index, one);
+}
+
+#endif // DEGREE_CPP

aten/cpu/graclus.cpp

 #include <torch/torch.h>
 
-inline std::tuple<at::Tensor, at::Tensor> remove_self_loops(at::Tensor row,
-                                                            at::Tensor col) {
-  auto mask = row != col;
-  return {row.masked_select(mask), col.masked_select(mask)};
-}
-
-inline std::tuple<at::Tensor, at::Tensor>
-randperm(at::Tensor row, at::Tensor col, int num_nodes) {
-  // Randomly reorder row and column indices.
-  auto perm = at::randperm(torch::CPU(at::kLong), row.size(0));
-  row = row.index_select(0, perm);
-  col = col.index_select(0, perm);
-
-  // Randomly swap row values.
-  auto node_rid = at::randperm(torch::CPU(at::kLong), num_nodes);
-  row = node_rid.index_select(0, row);
-
-  // Sort row and column indices row-wise.
-  std::tie(row, perm) = row.sort();
-  col = col.index_select(0, perm);
-
-  // Revert row value swaps.
-  row = std::get<1>(node_rid.sort()).index_select(0, row);
-
-  return {row, col};
-}
-
-inline at::Tensor degree(at::Tensor index, int num_nodes,
-                         at::ScalarType scalar_type) {
-  auto zero = at::full(torch::CPU(scalar_type), {num_nodes}, 0);
-  auto one = at::full(zero.type(), {index.size(0)}, 1);
-  return zero.scatter_add_(0, index, one);
-}
+#include "degree.cpp"
+#include "loop.cpp"
+#include "perm.cpp"
 
 at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes) {
   std::tie(row, col) = remove_self_loops(row, col);

aten/cpu/loop.cpp

+#ifndef LOOP_CPP
+#define LOOP_CPP
+
+#include <torch/torch.h>
+
+inline std::tuple<at::Tensor, at::Tensor> remove_self_loops(at::Tensor row,
+                                                            at::Tensor col) {
+  auto mask = row != col;
+  return {row.masked_select(mask), col.masked_select(mask)};
+}
+
+#endif // LOOP_CPP

aten/cpu/perm.cpp

+#ifndef PERM_CPP
+#define PERM_CPP
+
+#include <torch/torch.h>
+
+inline std::tuple<at::Tensor, at::Tensor>
+randperm(at::Tensor row, at::Tensor col, int num_nodes) {
+  // Randomly reorder row and column indices.
+  auto perm = at::randperm(torch::CPU(at::kLong), row.size(0));
+  row = row.index_select(0, perm);
+  col = col.index_select(0, perm);
+
+  // Randomly swap row values.
+  auto node_rid = at::randperm(torch::CPU(at::kLong), num_nodes);
+  row = node_rid.index_select(0, row);
+
+  // Sort row and column indices row-wise.
+  std::tie(row, perm) = row.sort();
+  col = col.index_select(0, perm);
+
+  // Revert row value swaps.
+  row = std::get<1>(node_rid.sort()).index_select(0, row);
+
+  return {row, col};
+}
+
+#endif // PERM_CPP