random_walk implementation

cmake/test.cmake

@@ -8,10 +8,12 @@ set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
 FetchContent_MakeAvailable(googletest)
 
 enable_testing()
-FILE(GLOB cpp_tests test/csrc/test_*.cpp)
-
 include(GoogleTest)
-foreach(t ${cpp_tests})
+
+set(CTEST test/csrc)
+file(GLOB_RECURSE ALL_TESTS ${CTEST}/*.cpp)
+
+foreach(t ${ALL_TESTS})
     get_filename_component(name ${t} NAME_WE)
     add_executable(${name} ${t})
     target_link_libraries(${name} ${PROJECT_NAME} gtest_main)

pyg_lib/csrc/sampler/random_walk.cpp

@@ -0,0 +1,28 @@
+#include "random_walk.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/library.h>
+
+namespace pyg {
+namespace sampler {
+
+at::Tensor random_walk(const at::Tensor& rowptr,
+                       const at::Tensor& col,
+                       const at::Tensor& seed,
+                       int64_t walk_length,
+                       double p,
+                       double q) {
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("pyg::random_walk", "")
+                       .typed<decltype(random_walk)>();
+  return op.call(rowptr, col, seed, walk_length, p, q);
+}
+
+TORCH_LIBRARY_FRAGMENT(pyg, m) {
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "pyg::random_walk(Tensor rowptr, Tensor col, Tensor seed, int "
+      "walk_length, float p, float q) -> Tensor"));
+}
+
+}  // namespace sampler
+}  // namespace pyg

pyg_lib/csrc/sampler/random_walk.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include "../macros.h"
+
+namespace pyg {
+namespace sampler {
+
+PYG_API at::Tensor random_walk(const at::Tensor& rowptr,
+                               const at::Tensor& col,
+                               const at::Tensor& seed,
+                               int64_t walk_length,
+                               double p,
+                               double q);
+
+}  // namespace sampler
+}  // namespace pyg

test/csrc/sampler/test_random_walk.cpp

@@ -0,0 +1,9 @@
+#include <gtest/gtest.h>
+
+#include "../../../pyg_lib/csrc/sampler/random_walk.h"
+
+#include <iostream>
+
+TEST(RandomWalkTest, BasicAssertions) {
+  std::cout << "random walk test" << std::endl;
+}