chore: import upstream snapshot with attribution

tests/cpp/common.h

@@ -0,0 +1,56 @@
+#ifndef TEST_COMMON_H_
+#define TEST_COMMON_H_
+
+#include <dgl/runtime/ndarray.h>
+
+static constexpr DGLContext CTX = DGLContext{kDGLCPU, 0};
+static constexpr DGLContext CPU = DGLContext{kDGLCPU, 0};
+#ifdef DGL_USE_CUDA
+static constexpr DGLContext GPU = DGLContext{kDGLCUDA, 0};
+#endif
+
+template <typename T>
+inline T* Ptr(dgl::runtime::NDArray nd) {
+  return static_cast<T*>(nd->data);
+}
+
+inline int64_t* PI64(dgl::runtime::NDArray nd) {
+  return static_cast<int64_t*>(nd->data);
+}
+
+inline int32_t* PI32(dgl::runtime::NDArray nd) {
+  return static_cast<int32_t*>(nd->data);
+}
+
+inline int64_t Len(dgl::runtime::NDArray nd) { return nd->shape[0]; }
+
+template <typename T>
+inline bool ArrayEQ(dgl::runtime::NDArray a1, dgl::runtime::NDArray a2) {
+  if (a1->ndim != a2->ndim) return false;
+  if (a1->dtype != a2->dtype) return false;
+  if (a1->ctx != a2->ctx) return false;
+  if (a1.NumElements() != a2.NumElements()) return false;
+  if (a1.NumElements() == 0) return true;
+  int64_t num = 1;
+  for (int i = 0; i < a1->ndim; ++i) {
+    if (a1->shape[i] != a2->shape[i]) return false;
+    num *= a1->shape[i];
+  }
+  a1 = a1.CopyTo(CPU);
+  a2 = a2.CopyTo(CPU);
+  for (int64_t i = 0; i < num; ++i)
+    if (static_cast<T*>(a1->data)[i] != static_cast<T*>(a2->data)[i])
+      return false;
+  return true;
+}
+
+template <typename T>
+inline bool IsInArray(dgl::runtime::NDArray a, T x) {
+  if (!a.defined() || a->shape[0] == 0) return false;
+  for (int64_t i = 0; i < a->shape[0]; ++i) {
+    if (x == static_cast<T*>(a->data)[i]) return true;
+  }
+  return false;
+}
+
+#endif  // TEST_COMMON_H_

tests/cpp/graph_index_test.cc

@@ -0,0 +1,13 @@
+/**
+ *  Copyright (c) 2019 by Contributors
+ * @file graph_index_test.cc
+ * @brief Test GraphIndex
+ */
+#include <dgl/graph.h>
+#include <gtest/gtest.h>
+
+TEST(GraphTest, TestNumVertices) {
+  dgl::Graph g;
+  g.AddVertices(10);
+  ASSERT_EQ(g.NumVertices(), 10);
+};

tests/cpp/message_queue_test.cc

@@ -0,0 +1,105 @@
+/**
+ *  Copyright (c) 2019 by Contributors
+ * @file msg_queue.cc
+ * @brief Message queue for DGL distributed training.
+ */
+#include <gtest/gtest.h>
+
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "../src/rpc/network/msg_queue.h"
+
+using dgl::network::Message;
+using dgl::network::MessageQueue;
+using std::string;
+
+TEST(MessageQueueTest, AddRemove) {
+  MessageQueue queue(5, 1);  // size:5, num_of_producer:1
+  // msg 1
+  std::string str_1("111");
+  Message msg_1 = {const_cast<char*>(str_1.data()), 3};
+  EXPECT_EQ(queue.Add(msg_1), ADD_SUCCESS);
+  // msg 2
+  std::string str_2("22");
+  Message msg_2 = {const_cast<char*>(str_2.data()), 2};
+  EXPECT_EQ(queue.Add(msg_2), ADD_SUCCESS);
+  // msg 3
+  std::string str_3("xxxx");
+  Message msg_3 = {const_cast<char*>(str_3.data()), 4};
+  EXPECT_EQ(queue.Add(msg_3, false), QUEUE_FULL);
+  // msg 4
+  Message msg_4;
+  EXPECT_EQ(queue.Remove(&msg_4), REMOVE_SUCCESS);
+  EXPECT_EQ(string(msg_4.data, msg_4.size), string("111"));
+  // msg 5
+  Message msg_5;
+  EXPECT_EQ(queue.Remove(&msg_5), REMOVE_SUCCESS);
+  EXPECT_EQ(string(msg_5.data, msg_5.size), string("22"));
+  // msg 6
+  std::string str_6("33333");
+  Message msg_6 = {const_cast<char*>(str_6.data()), 5};
+  EXPECT_EQ(queue.Add(msg_6), ADD_SUCCESS);
+  // msg 7
+  Message msg_7;
+  EXPECT_EQ(queue.Remove(&msg_7), REMOVE_SUCCESS);
+  EXPECT_EQ(string(msg_7.data, msg_7.size), string("33333"));
+  // msg 8
+  Message msg_8;
+  EXPECT_EQ(queue.Remove(&msg_8, false), QUEUE_EMPTY);  // non-blocking remove
+  // msg 9
+  std::string str_9("666666");
+  Message msg_9 = {const_cast<char*>(str_9.data()), 6};
+  EXPECT_EQ(queue.Add(msg_9), MSG_GT_SIZE);  // exceed queue size
+  // msg 10
+  std::string str_10("55555");
+  Message msg_10 = {const_cast<char*>(str_10.data()), 5};
+  EXPECT_EQ(queue.Add(msg_10), ADD_SUCCESS);
+  // msg 11
+  Message msg_11;
+  EXPECT_EQ(queue.Remove(&msg_11), REMOVE_SUCCESS);
+}
+
+TEST(MessageQueueTest, EmptyAndNoMoreAdd) {
+  MessageQueue queue(5, 2);  // size:5, num_of_producer:2
+  EXPECT_EQ(queue.EmptyAndNoMoreAdd(), false);
+  EXPECT_EQ(queue.Empty(), true);
+  queue.SignalFinished(1);
+  queue.SignalFinished(1);
+  EXPECT_EQ(queue.EmptyAndNoMoreAdd(), false);
+  queue.SignalFinished(2);
+  EXPECT_EQ(queue.EmptyAndNoMoreAdd(), true);
+}
+
+const int kNumOfProducer = 100;
+const int kNumOfMessage = 100;
+
+std::string str_apple("apple");
+
+void start_add(MessageQueue* queue, int id) {
+  for (int i = 0; i < kNumOfMessage; ++i) {
+    Message msg = {const_cast<char*>(str_apple.data()), 5};
+    EXPECT_EQ(queue->Add(msg), ADD_SUCCESS);
+  }
+  queue->SignalFinished(id);
+}
+
+TEST(MessageQueueTest, MultiThread) {
+  MessageQueue queue(100000, kNumOfProducer);
+  EXPECT_EQ(queue.EmptyAndNoMoreAdd(), false);
+  EXPECT_EQ(queue.Empty(), true);
+  std::vector<std::thread> thread_pool(kNumOfProducer);
+  for (int i = 0; i < kNumOfProducer; ++i) {
+    thread_pool[i] = std::thread(start_add, &queue, i);
+  }
+  for (int i = 0; i < kNumOfProducer * kNumOfMessage; ++i) {
+    Message msg;
+    EXPECT_EQ(queue.Remove(&msg), REMOVE_SUCCESS);
+    EXPECT_EQ(string(msg.data, msg.size), string("apple"));
+  }
+  for (int i = 0; i < kNumOfProducer; ++i) {
+    thread_pool[i].join();
+  }
+  EXPECT_EQ(queue.EmptyAndNoMoreAdd(), true);
+}

tests/cpp/socket_communicator_test.cc

@@ -0,0 +1,246 @@
+/**
+ *  Copyright (c) 2019 by Contributors
+ * @file socket_communicator_test.cc
+ * @brief Test SocketCommunicator
+ */
+#include "../src/rpc/network/socket_communicator.h"
+
+#include <gtest/gtest.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+
+#include <chrono>
+#include <fstream>
+#include <streambuf>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "../src/rpc/network/msg_queue.h"
+
+using std::string;
+
+using dgl::network::DefaultMessageDeleter;
+using dgl::network::Message;
+using dgl::network::SocketReceiver;
+using dgl::network::SocketSender;
+
+const int64_t kQueueSize = 500 * 1024;
+const int kThreadNum = 2;
+const int kMaxTryTimes = 1024;
+
+#ifndef WIN32
+
+const int kNumSender = 3;
+const int kNumReceiver = 3;
+const int kNumMessage = 10;
+
+const char* ip_addr[] = {
+    "tcp://127.0.0.1:50091", "tcp://127.0.0.1:50092", "tcp://127.0.0.1:50093"};
+
+static void start_client();
+static void start_server(int id);
+
+TEST(SocketCommunicatorTest, SendAndRecv) {
+  // start 10 client
+  std::vector<std::thread> client_thread(kNumSender);
+  for (int i = 0; i < kNumSender; ++i) {
+    client_thread[i] = std::thread(start_client);
+  }
+  // start 10 server
+  std::vector<std::thread> server_thread(kNumReceiver);
+  for (int i = 0; i < kNumReceiver; ++i) {
+    server_thread[i] = std::thread(start_server, i);
+  }
+  for (int i = 0; i < kNumSender; ++i) {
+    client_thread[i].join();
+  }
+  for (int i = 0; i < kNumReceiver; ++i) {
+    server_thread[i].join();
+  }
+}
+
+TEST(SocketCommunicatorTest, SendAndRecvTimeout) {
+  std::atomic_bool stop{false};
+  // start 1 client, connect to 1 server, send 2 messsage
+  auto client = std::thread([&stop]() {
+    SocketSender sender(kQueueSize, kThreadNum);
+    sender.ConnectReceiver(ip_addr[0], 0);
+    sender.ConnectReceiverFinalize(kMaxTryTimes);
+    for (int i = 0; i < 2; ++i) {
+      char* str_data = new char[9];
+      memcpy(str_data, "123456789", 9);
+      Message msg = {str_data, 9};
+      msg.deallocator = DefaultMessageDeleter;
+      EXPECT_EQ(sender.Send(msg, 0), ADD_SUCCESS);
+    }
+    while (!stop) {
+    }
+    sender.Finalize();
+  });
+  // start 1 server, accept 1 client, receive 2 message
+  auto server = std::thread([&stop]() {
+    SocketReceiver receiver(kQueueSize, kThreadNum);
+    receiver.Wait(ip_addr[0], 1);
+    Message msg;
+    int recv_id;
+    // receive 1st message
+    EXPECT_EQ(receiver.RecvFrom(&msg, 0, 0), REMOVE_SUCCESS);
+    EXPECT_EQ(string(msg.data, msg.size), string("123456789"));
+    msg.deallocator(&msg);
+    // receive 2nd message
+    EXPECT_EQ(receiver.Recv(&msg, &recv_id, 0), REMOVE_SUCCESS);
+    EXPECT_EQ(string(msg.data, msg.size), string("123456789"));
+    msg.deallocator(&msg);
+    // timed out
+    EXPECT_EQ(receiver.RecvFrom(&msg, 0, 1000), QUEUE_EMPTY);
+    EXPECT_EQ(receiver.Recv(&msg, &recv_id, 1000), QUEUE_EMPTY);
+    stop = true;
+    receiver.Finalize();
+  });
+  // join
+  client.join();
+  server.join();
+}
+
+void start_client() {
+  SocketSender sender(kQueueSize, kThreadNum);
+  for (int i = 0; i < kNumReceiver; ++i) {
+    sender.ConnectReceiver(ip_addr[i], i);
+  }
+  sender.ConnectReceiverFinalize(kMaxTryTimes);
+  for (int i = 0; i < kNumMessage; ++i) {
+    for (int n = 0; n < kNumReceiver; ++n) {
+      char* str_data = new char[9];
+      memcpy(str_data, "123456789", 9);
+      Message msg = {str_data, 9};
+      msg.deallocator = DefaultMessageDeleter;
+      EXPECT_EQ(sender.Send(msg, n), ADD_SUCCESS);
+    }
+  }
+  for (int i = 0; i < kNumMessage; ++i) {
+    for (int n = 0; n < kNumReceiver; ++n) {
+      char* str_data = new char[9];
+      memcpy(str_data, "123456789", 9);
+      Message msg = {str_data, 9};
+      msg.deallocator = DefaultMessageDeleter;
+      EXPECT_EQ(sender.Send(msg, n), ADD_SUCCESS);
+    }
+  }
+  sender.Finalize();
+}
+
+void start_server(int id) {
+  sleep(5);
+  SocketReceiver receiver(kQueueSize, kThreadNum);
+  receiver.Wait(ip_addr[id], kNumSender);
+  for (int i = 0; i < kNumMessage; ++i) {
+    for (int n = 0; n < kNumSender; ++n) {
+      Message msg;
+      EXPECT_EQ(receiver.RecvFrom(&msg, n), REMOVE_SUCCESS);
+      EXPECT_EQ(string(msg.data, msg.size), string("123456789"));
+      msg.deallocator(&msg);
+    }
+  }
+  for (int n = 0; n < kNumSender * kNumMessage; ++n) {
+    Message msg;
+    int recv_id;
+    EXPECT_EQ(receiver.Recv(&msg, &recv_id), REMOVE_SUCCESS);
+    EXPECT_EQ(string(msg.data, msg.size), string("123456789"));
+    msg.deallocator(&msg);
+  }
+  receiver.Finalize();
+}
+
+TEST(SocketCommunicatorTest, TCPSocketBind) {
+  dgl::network::TCPSocket socket;
+  testing::internal::CaptureStderr();
+  EXPECT_EQ(socket.Bind("127.0.0", 50001), false);
+  const std::string stderr = testing::internal::GetCapturedStderr();
+  EXPECT_NE(stderr.find("Invalid IP: 127.0.0"), std::string::npos);
+}
+
+#else
+
+#include <windows.h>
+#include <winsock2.h>
+
+#pragma comment(lib, "ws2_32.lib")
+
+void sleep(int seconds) { Sleep(seconds * 1000); }
+
+static void start_client();
+static bool start_server();
+
+DWORD WINAPI _ClientThreadFunc(LPVOID param) {
+  start_client();
+  return 0;
+}
+
+DWORD WINAPI _ServerThreadFunc(LPVOID param) { return start_server() ? 1 : 0; }
+
+TEST(SocketCommunicatorTest, SendAndRecv) {
+  HANDLE hThreads[2];
+  WSADATA wsaData;
+  DWORD retcode, exitcode;
+
+  srand((unsigned)time(NULL));
+  int port = (rand() % (5000 - 3000 + 1)) + 3000;
+  std::string ip_addr = "tcp://127.0.0.1:" + std::to_string(port);
+  std::ofstream out("addr.txt");
+  out << ip_addr;
+  out.close();
+
+  ASSERT_EQ(::WSAStartup(MAKEWORD(2, 2), &wsaData), 0);
+
+  hThreads[0] =
+      ::CreateThread(NULL, 0, _ClientThreadFunc, NULL, 0, NULL);  // client
+  ASSERT_TRUE(hThreads[0] != NULL);
+  hThreads[1] =
+      ::CreateThread(NULL, 0, _ServerThreadFunc, NULL, 0, NULL);  // server
+  ASSERT_TRUE(hThreads[1] != NULL);
+
+  retcode = ::WaitForMultipleObjects(2, hThreads, TRUE, INFINITE);
+  EXPECT_TRUE((retcode <= WAIT_OBJECT_0 + 1) && (retcode >= WAIT_OBJECT_0));
+
+  EXPECT_EQ(::GetExitCodeThread(hThreads[1], &exitcode), TRUE);
+  EXPECT_EQ(exitcode, 1);
+
+  EXPECT_EQ(::CloseHandle(hThreads[0]), TRUE);
+  EXPECT_EQ(::CloseHandle(hThreads[1]), TRUE);
+
+  ::WSACleanup();
+}
+
+static void start_client() {
+  std::ifstream t("addr.txt");
+  std::string ip_addr(
+      (std::istreambuf_iterator<char>(t)), std::istreambuf_iterator<char>());
+  t.close();
+  SocketSender sender(kQueueSize, kThreadNum);
+  sender.ConnectReceiver(ip_addr.c_str(), 0);
+  sender.ConnectReceiverFinalize(kMaxTryTimes);
+  char* str_data = new char[9];
+  memcpy(str_data, "123456789", 9);
+  Message msg = {str_data, 9};
+  msg.deallocator = DefaultMessageDeleter;
+  sender.Send(msg, 0);
+  sender.Finalize();
+}
+
+static bool start_server() {
+  sleep(5);
+  std::ifstream t("addr.txt");
+  std::string ip_addr(
+      (std::istreambuf_iterator<char>(t)), std::istreambuf_iterator<char>());
+  t.close();
+  SocketReceiver receiver(kQueueSize, kThreadNum);
+  receiver.Wait(ip_addr.c_str(), 1);
+  Message msg;
+  EXPECT_EQ(receiver.RecvFrom(&msg, 0), REMOVE_SUCCESS);
+  receiver.Finalize();
+  return string("123456789") == string(msg.data, msg.size);
+}
+
+#endif

tests/cpp/string_test.cc

@@ -0,0 +1,52 @@
+/**
+ *  Copyright (c) 2019 by Contributors
+ * @file string_test.cc
+ * @brief Test String Common
+ */
+#include <gtest/gtest.h>
+
+#include <string>
+#include <vector>
+
+#include "../src/rpc/network/common.h"
+
+using dgl::network::SplitStringUsing;
+using dgl::network::SStringPrintf;
+using dgl::network::StringAppendF;
+using dgl::network::StringPrintf;
+
+TEST(SplitStringTest, SplitStringUsingCompoundDelim) {
+  std::string full(" apple \torange ");
+  std::vector<std::string> subs;
+  SplitStringUsing(full, " \t", &subs);
+  EXPECT_EQ(subs.size(), 2);
+  EXPECT_EQ(subs[0], std::string("apple"));
+  EXPECT_EQ(subs[1], std::string("orange"));
+}
+
+TEST(SplitStringTest, testSplitStringUsingSingleDelim) {
+  std::string full(" apple orange ");
+  std::vector<std::string> subs;
+  SplitStringUsing(full, " ", &subs);
+  EXPECT_EQ(subs.size(), 2);
+  EXPECT_EQ(subs[0], std::string("apple"));
+  EXPECT_EQ(subs[1], std::string("orange"));
+}
+
+TEST(SplitStringTest, testSplitingNoDelimString) {
+  std::string full("apple");
+  std::vector<std::string> subs;
+  SplitStringUsing(full, " ", &subs);
+  EXPECT_EQ(subs.size(), 1);
+  EXPECT_EQ(subs[0], std::string("apple"));
+}
+
+TEST(StringPrintf, normal) {
+  using std::string;
+  EXPECT_EQ(StringPrintf("%d", 1), string("1"));
+  string target;
+  SStringPrintf(&target, "%d", 1);
+  EXPECT_EQ(target, string("1"));
+  StringAppendF(&target, "%d", 2);
+  EXPECT_EQ(target, string("12"));
+}

tests/cpp/test_concurrent_id_hash_map.cc

@@ -0,0 +1,75 @@
+#include <dgl/array.h>
+#include <dgl/runtime/parallel_for.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <set>
+
+#include "../../src/array/cpu/concurrent_id_hash_map.h"
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+using namespace dgl::aten;
+
+namespace {
+
+template <typename IdType>
+size_t ConstructRandomSet(
+    size_t size, IdType range, std::vector<IdType>& id_vec) {
+  id_vec.resize(size);
+  std::srand(std::time(nullptr));
+  for (size_t i = 0; i < size; i++) {
+    id_vec[i] = static_cast<IdType>(std::rand() % range);
+  }
+
+  size_t num_seeds = size / 5 + 1;
+  std::sort(id_vec.begin(), id_vec.begin() + num_seeds);
+  return std::unique(id_vec.begin(), id_vec.begin() + num_seeds) -
+         id_vec.begin();
+}
+
+template <typename IdType, size_t size, IdType range>
+void _TestIdMap() {
+  std::vector<IdType> id_vec;
+  auto num_seeds = ConstructRandomSet(size, range, id_vec);
+  std::set<IdType> id_set(id_vec.begin(), id_vec.end());
+  IdArray ids = VecToIdArray(id_vec, sizeof(IdType) * 8, CTX);
+  ConcurrentIdHashMap<IdType> id_map;
+  IdArray unique_ids = id_map.Init(ids, num_seeds);
+  auto unique_num = static_cast<size_t>(unique_ids->shape[0]);
+  IdType* unique_id_data = unique_ids.Ptr<IdType>();
+  EXPECT_EQ(id_set.size(), unique_num);
+
+  parallel_for(0, num_seeds, 64, [&](int64_t s, int64_t e) {
+    for (int64_t i = s; i < e; i++) {
+      EXPECT_EQ(id_vec[i], unique_id_data[i]);
+    }
+  });
+
+  parallel_for(num_seeds, unique_num, 128, [&](int64_t s, int64_t e) {
+    for (int64_t i = s; i < e; i++) {
+      EXPECT_TRUE(id_set.find(unique_id_data[i]) != id_set.end());
+    }
+  });
+
+  IdArray new_ids = id_map.MapIds(unique_ids);
+  EXPECT_TRUE(new_ids.IsContiguous());
+  ids->shape[0] = num_seeds;
+  IdArray new_seed_ids = id_map.MapIds(ids);
+  EXPECT_TRUE(new_seed_ids.IsContiguous());
+  EXPECT_EQ(new_seed_ids.Ptr<IdType>()[0], static_cast<IdType>(0));
+}
+
+TEST(ConcurrentIdHashMapTest, TestConcurrentIdHashMap) {
+  _TestIdMap<int32_t, 1, 10>();
+  _TestIdMap<int64_t, 1, 10>();
+  _TestIdMap<int32_t, 1000, 500000>();
+  _TestIdMap<int64_t, 1000, 500000>();
+  _TestIdMap<int32_t, 50000, 1000000>();
+  _TestIdMap<int64_t, 50000, 1000000>();
+  _TestIdMap<int32_t, 100000, 40000000>();
+  _TestIdMap<int64_t, 100000, 40000000>();
+}
+
+};  // namespace

tests/cpp/test_csrmm.cc

@@ -0,0 +1,215 @@
+#include <dgl/array.h>
+#include <dgl/kernel.h>
+#include <gtest/gtest.h>
+
+#include "../../src/array/cpu/array_utils.h"  // PairHash
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+
+namespace {
+
+// Unit tests:
+// CSRMM(A, B) == A_mm_B
+// CSRSum({A, C}) == A_plus_C
+// CSRMask(A, C) = A_mask_C
+
+template <typename IdType, typename DType>
+std::unordered_map<std::pair<IdType, IdType>, DType, aten::PairHash> COOToMap(
+    aten::COOMatrix coo, NDArray weights) {
+  std::unordered_map<std::pair<IdType, IdType>, DType, aten::PairHash> map;
+
+  for (int64_t i = 0; i < coo.row->shape[0]; ++i) {
+    IdType irow = aten::IndexSelect<IdType>(coo.row, i);
+    IdType icol = aten::IndexSelect<IdType>(coo.col, i);
+    IdType ieid =
+        aten::COOHasData(coo) ? aten::IndexSelect<IdType>(coo.data, i) : i;
+    DType idata = aten::IndexSelect<DType>(weights, ieid);
+    map.insert({{irow, icol}, idata});
+  }
+  return map;
+}
+
+template <typename IdType, typename DType>
+bool CSRIsClose(
+    aten::CSRMatrix A, aten::CSRMatrix B, NDArray A_weights, NDArray B_weights,
+    DType rtol, DType atol) {
+  auto Amap = COOToMap<IdType, DType>(CSRToCOO(A, false), A_weights);
+  auto Bmap = COOToMap<IdType, DType>(CSRToCOO(B, false), B_weights);
+
+  if (Amap.size() != Bmap.size()) return false;
+
+  for (auto itA : Amap) {
+    auto itB = Bmap.find(itA.first);
+    if (itB == Bmap.end()) return false;
+    if (fabs(itA.second - itB->second) >= rtol * fabs(itA.second) + atol)
+      return false;
+  }
+
+  return true;
+}
+
+template <typename IdType, typename DType>
+std::pair<aten::CSRMatrix, NDArray> CSR_A(DGLContext ctx = CTX) {
+  // matrix([[0. , 0. , 1. , 0.7, 0. ],
+  //         [0. , 0. , 0.5, 0.+, 0. ],
+  //         [0.4, 0.7, 0. , 0.2, 0. ],
+  //         [0. , 0. , 0. , 0. , 0.2]])
+  // (0.+ indicates that the entry exists but the value is 0.)
+  auto csr = aten::CSRMatrix(
+      4, 5, NDArray::FromVector(std::vector<IdType>({0, 2, 4, 7, 8}), ctx),
+      NDArray::FromVector(std::vector<IdType>({2, 3, 2, 3, 0, 1, 3, 4}), ctx),
+      NDArray::FromVector(std::vector<IdType>({1, 0, 2, 3, 4, 5, 6, 7}), ctx));
+  auto weights = NDArray::FromVector(
+      std::vector<DType>({0.7, 1.0, 0.5, 0.0, 0.4, 0.7, 0.2, 0.2}), ctx);
+  return {csr, weights};
+}
+
+template <typename IdType, typename DType>
+std::pair<aten::CSRMatrix, NDArray> CSR_B(DGLContext ctx = CTX) {
+  // matrix([[0. , 0.9, 0. , 0.6, 0. , 0.3],
+  //         [0. , 0. , 0. , 0. , 0. , 0.4],
+  //         [0.+, 0. , 0. , 0. , 0. , 0.9],
+  //         [0.8, 0.2, 0.3, 0.2, 0. , 0. ],
+  //         [0.2, 0.4, 0. , 0. , 0. , 0. ]])
+  // (0.+ indicates that the entry exists but the value is 0.)
+  auto csr = aten::CSRMatrix(
+      5, 6, NDArray::FromVector(std::vector<IdType>({0, 3, 4, 6, 10, 12}), ctx),
+      NDArray::FromVector(
+          std::vector<IdType>({1, 3, 5, 5, 0, 5, 0, 1, 2, 3, 0, 1}), ctx));
+  auto weights = NDArray::FromVector(
+      std::vector<DType>(
+          {0.9, 0.6, 0.3, 0.4, 0.0, 0.9, 0.8, 0.2, 0.3, 0.2, 0.2, 0.4}),
+      ctx);
+  return {csr, weights};
+}
+
+template <typename IdType, typename DType>
+std::pair<aten::CSRMatrix, NDArray> CSR_C(DGLContext ctx = CTX) {
+  // matrix([[0. , 0. , 0. , 0.2, 0. ],
+  //         [0. , 0. , 0. , 0.5, 0.4],
+  //         [0. , 0.2, 0. , 0.9, 0.2],
+  //         [0. , 1. , 0. , 0.7, 0. ]])
+  auto csr = aten::CSRMatrix(
+      4, 5, NDArray::FromVector(std::vector<IdType>({0, 1, 3, 6, 8}), ctx),
+      NDArray::FromVector(std::vector<IdType>({3, 3, 4, 1, 3, 4, 1, 3}), ctx));
+  auto weights = NDArray::FromVector(
+      std::vector<DType>({0.2, 0.5, 0.4, 0.2, 0.9, 0.2, 1., 0.7}), ctx);
+  return {csr, weights};
+}
+
+template <typename IdType, typename DType>
+std::pair<aten::CSRMatrix, NDArray> CSR_A_mm_B(DGLContext ctx = CTX) {
+  // matrix([[0.56, 0.14, 0.21, 0.14, 0.  , 0.9 ],
+  //         [0.+ , 0.+ , 0.+ , 0.+ , 0.  , 0.45],
+  //         [0.16, 0.4 , 0.06, 0.28, 0.  , 0.4 ],
+  //         [0.04, 0.08, 0.  , 0.  , 0.  , 0.  ]])
+  // (0.+ indicates that the entry exists but the value is 0.)
+  auto csr = aten::CSRMatrix(
+      4, 6, NDArray::FromVector(std::vector<IdType>({0, 5, 10, 15, 17}), ctx),
+      NDArray::FromVector(
+          std::vector<IdType>(
+              {0, 1, 2, 3, 5, 0, 1, 2, 3, 5, 0, 1, 2, 3, 5, 0, 1}),
+          ctx));
+  auto weights = NDArray::FromVector(
+      std::vector<DType>(
+          {0.56, 0.14, 0.21, 0.14, 0.9, 0., 0., 0., 0., 0.45, 0.16, 0.4, 0.06,
+           0.28, 0.4, 0.04, 0.08}),
+      ctx);
+  return {csr, weights};
+}
+
+template <typename IdType, typename DType>
+std::pair<aten::CSRMatrix, NDArray> CSR_A_plus_C(DGLContext ctx = CTX) {
+  auto csr = aten::CSRMatrix(
+      4, 5, NDArray::FromVector(std::vector<IdType>({0, 2, 5, 9, 12}), ctx),
+      NDArray::FromVector(
+          std::vector<IdType>({2, 3, 2, 3, 4, 0, 1, 3, 4, 1, 3, 4}), ctx));
+  auto weights = NDArray::FromVector(
+      std::vector<DType>(
+          {1., 0.9, 0.5, 0.5, 0.4, 0.4, 0.9, 1.1, 0.2, 1., 0.7, 0.2}),
+      ctx);
+  return {csr, weights};
+}
+
+template <typename DType>
+NDArray CSR_A_mask_C(DGLContext ctx = CTX) {
+  return NDArray::FromVector(
+      std::vector<DType>({0.7, 0.0, 0.0, 0.7, 0.2, 0.0, 0.0, 0.0}), ctx);
+}
+
+template <typename IdType, typename DType>
+void _TestCsrmm(DGLContext ctx = CTX) {
+  auto A = CSR_A<IdType, DType>(ctx);
+  auto B = CSR_B<IdType, DType>(ctx);
+  auto A_mm_B = aten::CSRMM(A.first, A.second, B.first, B.second);
+  auto A_mm_B2 = CSR_A_mm_B<IdType, DType>(ctx);
+  bool result = CSRIsClose<IdType, DType>(
+      A_mm_B.first, A_mm_B2.first, A_mm_B.second, A_mm_B2.second, 1e-4, 1e-4);
+  ASSERT_TRUE(result);
+}
+
+template <typename IdType, typename DType>
+void _TestCsrsum(DGLContext ctx = CTX) {
+  auto A = CSR_A<IdType, DType>(ctx);
+  auto C = CSR_C<IdType, DType>(ctx);
+  auto A_plus_C = aten::CSRSum({A.first, C.first}, {A.second, C.second});
+  auto A_plus_C2 = CSR_A_plus_C<IdType, DType>(ctx);
+  bool result = CSRIsClose<IdType, DType>(
+      A_plus_C.first, A_plus_C2.first, A_plus_C.second, A_plus_C2.second, 1e-4,
+      1e-4);
+  ASSERT_TRUE(result);
+}
+
+template <typename IdType, typename DType>
+void _TestCsrmask(DGLContext ctx = CTX) {
+  auto A = CSR_A<IdType, DType>(ctx);
+  auto C = CSR_C<IdType, DType>(ctx);
+  auto C_coo = CSRToCOO(C.first, false);
+  auto A_mask_C =
+      aten::CSRGetData<DType>(A.first, C_coo.row, C_coo.col, A.second, 0);
+  auto A_mask_C2 = CSR_A_mask_C<DType>(ctx);
+  ASSERT_TRUE(ArrayEQ<DType>(A_mask_C, A_mask_C2));
+}
+
+TEST(CsrmmTest, TestCsrmm) {
+  _TestCsrmm<int32_t, float>(CPU);
+  _TestCsrmm<int32_t, double>(CPU);
+  _TestCsrmm<int64_t, float>(CPU);
+  _TestCsrmm<int64_t, double>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCsrmm<int32_t, float>(GPU);
+  _TestCsrmm<int32_t, double>(GPU);
+  _TestCsrmm<int64_t, float>(GPU);
+  _TestCsrmm<int64_t, double>(GPU);
+#endif
+}
+
+TEST(CsrmmTest, TestCsrsum) {
+  _TestCsrsum<int32_t, float>(CPU);
+  _TestCsrsum<int32_t, double>(CPU);
+  _TestCsrsum<int64_t, float>(CPU);
+  _TestCsrsum<int64_t, double>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCsrsum<int32_t, float>(GPU);
+  _TestCsrsum<int32_t, double>(GPU);
+  _TestCsrsum<int64_t, float>(GPU);
+  _TestCsrsum<int64_t, double>(GPU);
+#endif
+}
+
+TEST(CsrmmTest, TestCsrmask) {
+  _TestCsrmask<int32_t, float>(CPU);
+  _TestCsrmask<int32_t, double>(CPU);
+  _TestCsrmask<int64_t, float>(CPU);
+  _TestCsrmask<int64_t, double>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCsrmask<int32_t, float>(GPU);
+  _TestCsrmask<int32_t, double>(GPU);
+  _TestCsrmask<int64_t, float>(GPU);
+  _TestCsrmask<int64_t, double>(GPU);
+#endif
+}
+
+};  // namespace

tests/cpp/test_partition.cc

@@ -0,0 +1,196 @@
+#include <gtest/gtest.h>
+
+#include "../../src/partition/ndarray_partition.h"
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::partition;
+
+template <DGLDeviceType XPU, typename IdType>
+void _TestRemainder_GeneratePermutation() {
+  const int64_t size = 160000;
+  const int num_parts = 7;
+  NDArrayPartitionRef part = CreatePartitionRemainderBased(size, num_parts);
+
+  IdArray idxs =
+      aten::Range(0, size / 10, sizeof(IdType) * 8, DGLContext{XPU, 0});
+
+  std::pair<IdArray, IdArray> result = part->GeneratePermutation(idxs);
+
+  // first part of result should be the permutation
+  IdArray perm = result.first.CopyTo(DGLContext{kDGLCPU, 0});
+  ASSERT_TRUE(perm.Ptr<IdType>() != nullptr);
+  ASSERT_EQ(perm->shape[0], idxs->shape[0]);
+  const IdType* const perm_cpu = static_cast<const IdType*>(perm->data);
+
+  // second part of result should be the counts
+  IdArray counts = result.second.CopyTo(DGLContext{kDGLCPU, 0});
+  ASSERT_TRUE(counts.Ptr<int64_t>() != nullptr);
+  ASSERT_EQ(counts->shape[0], num_parts);
+  const int64_t* const counts_cpu = static_cast<const int64_t*>(counts->data);
+
+  std::vector<int64_t> prefix(num_parts + 1, 0);
+  for (int p = 0; p < num_parts; ++p) {
+    prefix[p + 1] = prefix[p] + counts_cpu[p];
+  }
+  ASSERT_EQ(prefix.back(), idxs->shape[0]);
+
+  // copy original indexes to cpu
+  idxs = idxs.CopyTo(DGLContext{kDGLCPU, 0});
+  const IdType* const idxs_cpu = static_cast<const IdType*>(idxs->data);
+
+  for (int p = 0; p < num_parts; ++p) {
+    for (int64_t i = prefix[p]; i < prefix[p + 1]; ++i) {
+      EXPECT_EQ(idxs_cpu[perm_cpu[i]] % num_parts, p);
+    }
+  }
+}
+
+template <DGLDeviceType XPU, typename IdType>
+void _TestRemainder_MapToX() {
+  const int64_t size = 160000;
+  const int num_parts = 7;
+  NDArrayPartitionRef part = CreatePartitionRemainderBased(size, num_parts);
+
+  for (int part_id = 0; part_id < num_parts; ++part_id) {
+    IdArray local = aten::Range(
+        0, part->PartSize(part_id), sizeof(IdType) * 8, DGLContext{XPU, 0});
+    IdArray global = part->MapToGlobal(local, part_id);
+    IdArray act_local = part->MapToLocal(global).CopyTo(CPU);
+
+    // every global index should have the same remainder as the part id
+    ASSERT_EQ(global->shape[0], local->shape[0]);
+    global = global.CopyTo(CPU);
+    for (int64_t i = 0; i < global->shape[0]; ++i) {
+      EXPECT_EQ(Ptr<IdType>(global)[i] % num_parts, part_id)
+          << "i=" << i << ", num_parts=" << num_parts
+          << ", part_id=" << part_id;
+    }
+
+    // the remapped local indices to should match the original
+    local = local.CopyTo(CPU);
+    ASSERT_EQ(local->shape[0], act_local->shape[0]);
+    for (int64_t i = 0; i < act_local->shape[0]; ++i) {
+      EXPECT_EQ(Ptr<IdType>(local)[i], Ptr<IdType>(act_local)[i]);
+    }
+  }
+}
+
+TEST(PartitionTest, TestRemainderPartition) {
+#ifdef DGL_USE_CUDA
+  _TestRemainder_GeneratePermutation<kDGLCUDA, int32_t>();
+  _TestRemainder_GeneratePermutation<kDGLCUDA, int64_t>();
+
+  _TestRemainder_MapToX<kDGLCUDA, int32_t>();
+  _TestRemainder_MapToX<kDGLCUDA, int64_t>();
+#endif
+  // CPU is not implemented
+}
+
+template <typename INDEX, typename RANGE>
+int _FindPart(const INDEX idx, const RANGE* const range, const int num_parts) {
+  for (int i = 0; i < num_parts; ++i) {
+    if (range[i + 1] > idx) {
+      return i;
+    }
+  }
+
+  return -1;
+}
+
+template <DGLDeviceType XPU, typename IdType>
+void _TestRange_GeneratePermutation() {
+  const int64_t size = 160000;
+  const int num_parts = 7;
+  IdArray range = aten::NewIdArray(
+      num_parts + 1, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
+  for (int i = 0; i < num_parts; ++i) {
+    range.Ptr<IdType>()[i] = (size / num_parts) * i;
+  }
+  range.Ptr<IdType>()[num_parts] = size;
+  NDArrayPartitionRef part = CreatePartitionRangeBased(
+      size, num_parts, range.CopyTo(DGLContext{XPU, 0}));
+
+  IdArray idxs =
+      aten::Range(0, size / 10, sizeof(IdType) * 8, DGLContext{XPU, 0});
+
+  std::pair<IdArray, IdArray> result = part->GeneratePermutation(idxs);
+
+  // first part of result should be the permutation
+  IdArray perm = result.first.CopyTo(DGLContext{kDGLCPU, 0});
+  ASSERT_TRUE(perm.Ptr<IdType>() != nullptr);
+  ASSERT_EQ(perm->shape[0], idxs->shape[0]);
+  const IdType* const perm_cpu = static_cast<const IdType*>(perm->data);
+
+  // second part of result should be the counts
+  IdArray counts = result.second.CopyTo(DGLContext{kDGLCPU, 0});
+  ASSERT_TRUE(counts.Ptr<int64_t>() != nullptr);
+  ASSERT_EQ(counts->shape[0], num_parts);
+  const int64_t* const counts_cpu = static_cast<const int64_t*>(counts->data);
+
+  std::vector<int64_t> prefix(num_parts + 1, 0);
+  for (int p = 0; p < num_parts; ++p) {
+    prefix[p + 1] = prefix[p] + counts_cpu[p];
+  }
+  ASSERT_EQ(prefix.back(), idxs->shape[0]);
+
+  // copy original indexes to cpu
+  idxs = idxs.CopyTo(DGLContext{kDGLCPU, 0});
+  const IdType* const idxs_cpu = static_cast<const IdType*>(idxs->data);
+
+  for (int p = 0; p < num_parts; ++p) {
+    for (int64_t i = prefix[p]; i < prefix[p + 1]; ++i) {
+      EXPECT_EQ(
+          _FindPart(idxs_cpu[perm_cpu[i]], range.Ptr<IdType>(), num_parts), p);
+    }
+  }
+}
+
+template <DGLDeviceType XPU, typename IdType>
+void _TestRange_MapToX() {
+  const int64_t size = 160000;
+  const int num_parts = 7;
+  IdArray range = aten::NewIdArray(
+      num_parts + 1, DGLContext{kDGLCPU, 0}, sizeof(IdType) * 8);
+  for (int i = 0; i < num_parts; ++i) {
+    Ptr<IdType>(range)[i] = (size / num_parts) * i;
+  }
+  range.Ptr<IdType>()[num_parts] = size;
+  NDArrayPartitionRef part = CreatePartitionRangeBased(
+      size, num_parts, range.CopyTo(DGLContext{XPU, 0}));
+
+  for (int part_id = 0; part_id < num_parts; ++part_id) {
+    IdArray local = aten::Range(
+        0, part->PartSize(part_id), sizeof(IdType) * 8, DGLContext{XPU, 0});
+    IdArray global = part->MapToGlobal(local, part_id);
+    IdArray act_local = part->MapToLocal(global).CopyTo(CPU);
+
+    ASSERT_EQ(global->shape[0], local->shape[0]);
+    global = global.CopyTo(CPU);
+    for (int64_t i = 0; i < global->shape[0]; ++i) {
+      EXPECT_EQ(
+          _FindPart(Ptr<IdType>(global)[i], Ptr<IdType>(range), num_parts),
+          part_id)
+          << "i=" << i << ", num_parts=" << num_parts << ", part_id=" << part_id
+          << ", shape=" << global->shape[0];
+    }
+
+    // the remapped local indices to should match the original
+    local = local.CopyTo(CPU);
+    ASSERT_EQ(local->shape[0], act_local->shape[0]);
+    for (int64_t i = 0; i < act_local->shape[0]; ++i) {
+      EXPECT_EQ(Ptr<IdType>(local)[i], Ptr<IdType>(act_local)[i]);
+    }
+  }
+}
+
+TEST(PartitionTest, TestRangePartition) {
+#ifdef DGL_USE_CUDA
+  _TestRange_GeneratePermutation<kDGLCUDA, int32_t>();
+  _TestRange_GeneratePermutation<kDGLCUDA, int64_t>();
+
+  _TestRange_MapToX<kDGLCUDA, int32_t>();
+  _TestRange_MapToX<kDGLCUDA, int64_t>();
+#endif
+  // CPU is not implemented
+}

tests/cpp/test_rowwise.cc

@@ -0,0 +1,905 @@
+#include <dgl/array.h>
+#include <gtest/gtest.h>
+
+#include <set>
+#include <tuple>
+
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+using namespace dgl::aten;
+
+template <typename Idx>
+using ETuple = std::tuple<Idx, Idx, Idx>;
+
+template <typename Idx>
+std::set<ETuple<Idx>> AllEdgeSet(bool has_data) {
+  if (has_data) {
+    std::set<ETuple<Idx>> eset;
+    eset.insert(ETuple<Idx>{0, 0, 2});
+    eset.insert(ETuple<Idx>{0, 1, 3});
+    eset.insert(ETuple<Idx>{1, 1, 0});
+    eset.insert(ETuple<Idx>{3, 2, 1});
+    eset.insert(ETuple<Idx>{3, 3, 4});
+    return eset;
+  } else {
+    std::set<ETuple<Idx>> eset;
+    eset.insert(ETuple<Idx>{0, 0, 0});
+    eset.insert(ETuple<Idx>{0, 1, 1});
+    eset.insert(ETuple<Idx>{1, 1, 2});
+    eset.insert(ETuple<Idx>{3, 2, 3});
+    eset.insert(ETuple<Idx>{3, 3, 4});
+    return eset;
+  }
+}
+
+template <typename Idx>
+std::set<ETuple<Idx>> AllEdgePerEtypeSet(bool has_data) {
+  if (has_data) {
+    std::set<ETuple<Idx>> eset;
+    eset.insert(ETuple<Idx>{0, 0, 0});
+    eset.insert(ETuple<Idx>{0, 1, 1});
+    eset.insert(ETuple<Idx>{0, 2, 4});
+    eset.insert(ETuple<Idx>{0, 3, 6});
+    eset.insert(ETuple<Idx>{3, 2, 5});
+    eset.insert(ETuple<Idx>{3, 3, 3});
+    return eset;
+  } else {
+    std::set<ETuple<Idx>> eset;
+    eset.insert(ETuple<Idx>{0, 0, 0});
+    eset.insert(ETuple<Idx>{0, 1, 1});
+    eset.insert(ETuple<Idx>{0, 2, 2});
+    eset.insert(ETuple<Idx>{0, 3, 3});
+    eset.insert(ETuple<Idx>{3, 3, 5});
+    eset.insert(ETuple<Idx>{3, 2, 6});
+    return eset;
+  }
+}
+
+template <typename Idx>
+std::set<ETuple<Idx>> ToEdgeSet(COOMatrix mat) {
+  std::set<ETuple<Idx>> eset;
+  Idx* row = static_cast<Idx*>(mat.row->data);
+  Idx* col = static_cast<Idx*>(mat.col->data);
+  Idx* data = static_cast<Idx*>(mat.data->data);
+  for (int64_t i = 0; i < mat.row->shape[0]; ++i) {
+    // std::cout << row[i] << " " << col[i] <<  " " << data[i] << std::endl;
+    eset.emplace(row[i], col[i], data[i]);
+  }
+  return eset;
+}
+
+template <typename Idx>
+void CheckSampledResult(COOMatrix mat, IdArray rows, bool has_data) {
+  ASSERT_EQ(mat.num_rows, 4);
+  ASSERT_EQ(mat.num_cols, 4);
+  Idx* row = static_cast<Idx*>(mat.row->data);
+  Idx* col = static_cast<Idx*>(mat.col->data);
+  Idx* data = static_cast<Idx*>(mat.data->data);
+  const auto& gt = AllEdgeSet<Idx>(has_data);
+  for (int64_t i = 0; i < mat.row->shape[0]; ++i) {
+    ASSERT_TRUE(gt.count(std::make_tuple(row[i], col[i], data[i])));
+    ASSERT_TRUE(IsInArray(rows, row[i]));
+  }
+}
+
+template <typename Idx>
+void CheckSampledPerEtypeResult(COOMatrix mat, IdArray rows, bool has_data) {
+  ASSERT_EQ(mat.num_rows, 4);
+  ASSERT_EQ(mat.num_cols, 4);
+  Idx* row = static_cast<Idx*>(mat.row->data);
+  Idx* col = static_cast<Idx*>(mat.col->data);
+  Idx* data = static_cast<Idx*>(mat.data->data);
+  const auto& gt = AllEdgePerEtypeSet<Idx>(has_data);
+  for (int64_t i = 0; i < mat.row->shape[0]; ++i) {
+    int64_t count = gt.count(std::make_tuple(row[i], col[i], data[i]));
+    ASSERT_TRUE(count);
+    ASSERT_TRUE(IsInArray(rows, row[i]));
+  }
+}
+
+template <typename Idx>
+CSRMatrix CSR(bool has_data) {
+  IdArray indptr = NDArray::FromVector(std::vector<Idx>({0, 2, 3, 3, 5}));
+  IdArray indices = NDArray::FromVector(std::vector<Idx>({0, 1, 1, 2, 3}));
+  IdArray data = NDArray::FromVector(std::vector<Idx>({2, 3, 0, 1, 4}));
+  if (has_data)
+    return CSRMatrix(4, 4, indptr, indices, data);
+  else
+    return CSRMatrix(4, 4, indptr, indices);
+}
+
+template <typename Idx>
+COOMatrix COO(bool has_data) {
+  IdArray row = NDArray::FromVector(std::vector<Idx>({0, 0, 1, 3, 3}));
+  IdArray col = NDArray::FromVector(std::vector<Idx>({0, 1, 1, 2, 3}));
+  IdArray data = NDArray::FromVector(std::vector<Idx>({2, 3, 0, 1, 4}));
+  if (has_data)
+    return COOMatrix(4, 4, row, col, data);
+  else
+    return COOMatrix(4, 4, row, col);
+}
+
+template <typename Idx>
+std::pair<CSRMatrix, std::vector<int64_t>> CSREtypes(bool has_data) {
+  IdArray indptr = NDArray::FromVector(std::vector<Idx>({0, 4, 5, 5, 7}));
+  IdArray indices =
+      NDArray::FromVector(std::vector<Idx>({0, 1, 2, 3, 1, 3, 2}));
+  IdArray data = NDArray::FromVector(std::vector<Idx>({0, 1, 4, 6, 2, 3, 5}));
+  auto eid2etype_offsets = std::vector<int64_t>({0, 4, 5, 6, 7});
+  if (has_data)
+    return {CSRMatrix(4, 4, indptr, indices, data), eid2etype_offsets};
+  else
+    return {CSRMatrix(4, 4, indptr, indices), eid2etype_offsets};
+}
+
+template <typename Idx>
+std::pair<COOMatrix, std::vector<int64_t>> COOEtypes(bool has_data) {
+  IdArray row = NDArray::FromVector(std::vector<Idx>({0, 0, 0, 0, 1, 3, 3}));
+  IdArray col = NDArray::FromVector(std::vector<Idx>({0, 1, 2, 3, 1, 3, 2}));
+  IdArray data = NDArray::FromVector(std::vector<Idx>({0, 1, 4, 6, 2, 3, 5}));
+  auto eid2etype_offsets = std::vector<int64_t>({0, 4, 5, 6, 7});
+  if (has_data)
+    return {COOMatrix(4, 4, row, col, data), eid2etype_offsets};
+  else
+    return {COOMatrix(4, 4, row, col), eid2etype_offsets};
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRSampling(bool has_data) {
+  auto mat = CSR<Idx>(has_data);
+  FloatArray prob =
+      NDArray::FromVector(std::vector<FloatType>({.5, .5, .5, .5, .5}));
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSampling(mat, rows, 2, prob, false);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 4);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+  prob = NDArray::FromVector(std::vector<FloatType>({.0, .5, .5, .0, .5}));
+  for (int k = 0; k < 100; ++k) {
+    auto rst = CSRRowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 1, 3)));
+    } else {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_FALSE(eset.count(std::make_tuple(3, 2, 3)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCSRSampling) {
+  _TestCSRSampling<int32_t, float>(true);
+  _TestCSRSampling<int64_t, float>(true);
+  _TestCSRSampling<int32_t, double>(true);
+  _TestCSRSampling<int64_t, double>(true);
+  _TestCSRSampling<int32_t, float>(false);
+  _TestCSRSampling<int64_t, float>(false);
+  _TestCSRSampling<int32_t, double>(false);
+  _TestCSRSampling<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRSamplingUniform(bool has_data) {
+  auto mat = CSR<Idx>(has_data);
+  FloatArray prob = aten::NullArray();
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSampling(mat, rows, 2, prob, false);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCSRSamplingUniform) {
+  _TestCSRSamplingUniform<int32_t, float>(true);
+  _TestCSRSamplingUniform<int64_t, float>(true);
+  _TestCSRSamplingUniform<int32_t, double>(true);
+  _TestCSRSamplingUniform<int64_t, double>(true);
+  _TestCSRSamplingUniform<int32_t, float>(false);
+  _TestCSRSamplingUniform<int64_t, float>(false);
+  _TestCSRSamplingUniform<int32_t, double>(false);
+  _TestCSRSamplingUniform<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRPerEtypeSampling(bool has_data) {
+  auto pair = CSREtypes<Idx>(has_data);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      NDArray::FromVector(std::vector<FloatType>({.5, .5, .5, .5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 2, 4));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 3, 6));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 2));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 3));
+      ASSERT_EQ(counts, 1);
+    } else {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      counts += eset.count(std::make_tuple(0, 2, 2));
+      counts += eset.count(std::make_tuple(0, 3, 3));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 4));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 6));
+      ASSERT_EQ(counts, 1);
+    }
+  }
+
+  prob = {
+      NDArray::FromVector(std::vector<FloatType>({.0, .5, .0, .0})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+    } else {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 2, 2)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 3, 3)));
+    }
+  }
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRPerEtypeSamplingSorted() {
+  auto pair = CSREtypes<Idx>(true);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      NDArray::FromVector(std::vector<FloatType>({.5, .5, .5, .5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, true);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, true);
+    auto eset = ToEdgeSet<Idx>(rst);
+    int counts = 0;
+    counts += eset.count(std::make_tuple(0, 0, 0));
+    counts += eset.count(std::make_tuple(0, 1, 1));
+    ASSERT_EQ(counts, 2);
+    counts = 0;
+    counts += eset.count(std::make_tuple(0, 2, 4));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(0, 3, 6));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(1, 1, 2));
+    ASSERT_EQ(counts, 0);
+    counts = 0;
+    counts += eset.count(std::make_tuple(3, 2, 5));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(3, 3, 3));
+    ASSERT_EQ(counts, 1);
+  }
+
+  prob = {
+      NDArray::FromVector(std::vector<FloatType>({.0, .5, .0, .0})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, true);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+  }
+}
+
+TEST(RowwiseTest, TestCSRPerEtypeSampling) {
+  _TestCSRPerEtypeSampling<int32_t, float>(true);
+  _TestCSRPerEtypeSampling<int64_t, float>(true);
+  _TestCSRPerEtypeSampling<int32_t, double>(true);
+  _TestCSRPerEtypeSampling<int64_t, double>(true);
+  _TestCSRPerEtypeSampling<int32_t, float>(false);
+  _TestCSRPerEtypeSampling<int64_t, float>(false);
+  _TestCSRPerEtypeSampling<int32_t, double>(false);
+  _TestCSRPerEtypeSampling<int64_t, double>(false);
+  _TestCSRPerEtypeSamplingSorted<int32_t, float>();
+  _TestCSRPerEtypeSamplingSorted<int64_t, float>();
+  _TestCSRPerEtypeSamplingSorted<int32_t, double>();
+  _TestCSRPerEtypeSamplingSorted<int64_t, double>();
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRPerEtypeSamplingUniform(bool has_data) {
+  auto pair = CSREtypes<Idx>(has_data);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      aten::NullArray(), aten::NullArray(), aten::NullArray(),
+      aten::NullArray()};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 2, 4));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 3, 6));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 2));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 3));
+      ASSERT_EQ(counts, 1);
+    } else {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      counts += eset.count(std::make_tuple(0, 2, 2));
+      counts += eset.count(std::make_tuple(0, 3, 3));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 4));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 6));
+      ASSERT_EQ(counts, 1);
+    }
+  }
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRPerEtypeSamplingUniformSorted() {
+  auto pair = CSREtypes<Idx>(true);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      aten::NullArray(), aten::NullArray(), aten::NullArray(),
+      aten::NullArray()};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, true);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, true);
+    auto eset = ToEdgeSet<Idx>(rst);
+    int counts = 0;
+    counts += eset.count(std::make_tuple(0, 0, 0));
+    counts += eset.count(std::make_tuple(0, 1, 1));
+    ASSERT_EQ(counts, 2);
+    counts = 0;
+    counts += eset.count(std::make_tuple(0, 2, 4));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(0, 3, 6));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(1, 1, 2));
+    ASSERT_EQ(counts, 0);
+    counts = 0;
+    counts += eset.count(std::make_tuple(3, 2, 5));
+    ASSERT_EQ(counts, 1);
+    counts = 0;
+    counts += eset.count(std::make_tuple(3, 3, 3));
+    ASSERT_EQ(counts, 1);
+  }
+}
+
+TEST(RowwiseTest, TestCSRPerEtypeSamplingUniform) {
+  _TestCSRPerEtypeSamplingUniform<int32_t, float>(true);
+  _TestCSRPerEtypeSamplingUniform<int64_t, float>(true);
+  _TestCSRPerEtypeSamplingUniform<int32_t, double>(true);
+  _TestCSRPerEtypeSamplingUniform<int64_t, double>(true);
+  _TestCSRPerEtypeSamplingUniform<int32_t, float>(false);
+  _TestCSRPerEtypeSamplingUniform<int64_t, float>(false);
+  _TestCSRPerEtypeSamplingUniform<int32_t, double>(false);
+  _TestCSRPerEtypeSamplingUniform<int64_t, double>(false);
+  _TestCSRPerEtypeSamplingUniformSorted<int32_t, float>();
+  _TestCSRPerEtypeSamplingUniformSorted<int64_t, float>();
+  _TestCSRPerEtypeSamplingUniformSorted<int32_t, double>();
+  _TestCSRPerEtypeSamplingUniformSorted<int64_t, double>();
+}
+
+template <typename Idx, typename FloatType>
+void _TestCOOSampling(bool has_data) {
+  auto mat = COO<Idx>(has_data);
+  FloatArray prob =
+      NDArray::FromVector(std::vector<FloatType>({.5, .5, .5, .5, .5}));
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWiseSampling(mat, rows, 2, prob, false);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 4);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+  prob = NDArray::FromVector(std::vector<FloatType>({.0, .5, .5, .0, .5}));
+  for (int k = 0; k < 100; ++k) {
+    auto rst = COORowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 1, 3)));
+    } else {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_FALSE(eset.count(std::make_tuple(3, 2, 3)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCOOSampling) {
+  _TestCOOSampling<int32_t, float>(true);
+  _TestCOOSampling<int64_t, float>(true);
+  _TestCOOSampling<int32_t, double>(true);
+  _TestCOOSampling<int64_t, double>(true);
+  _TestCOOSampling<int32_t, float>(false);
+  _TestCOOSampling<int64_t, float>(false);
+  _TestCOOSampling<int32_t, double>(false);
+  _TestCOOSampling<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCOOSamplingUniform(bool has_data) {
+  auto mat = COO<Idx>(has_data);
+  FloatArray prob = aten::NullArray();
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWiseSampling(mat, rows, 2, prob, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWiseSampling(mat, rows, 2, prob, false);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCOOSamplingUniform) {
+  _TestCOOSamplingUniform<int32_t, float>(true);
+  _TestCOOSamplingUniform<int64_t, float>(true);
+  _TestCOOSamplingUniform<int32_t, double>(true);
+  _TestCOOSamplingUniform<int64_t, double>(true);
+  _TestCOOSamplingUniform<int32_t, float>(false);
+  _TestCOOSamplingUniform<int64_t, float>(false);
+  _TestCOOSamplingUniform<int32_t, double>(false);
+  _TestCOOSamplingUniform<int64_t, double>(false);
+}
+
+// COOPerEtypeSampling with rowwise_etype_sorted == true is not meaningful as
+// it's never used in practice.
+
+template <typename Idx, typename FloatType>
+void _TestCOOPerEtypeSampling(bool has_data) {
+  auto pair = COOEtypes<Idx>(has_data);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      NDArray::FromVector(std::vector<FloatType>({.5, .5, .5, .5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 2, 4));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 3, 6));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 2));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 3));
+      ASSERT_EQ(counts, 1);
+    } else {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      counts += eset.count(std::make_tuple(0, 2, 2));
+      counts += eset.count(std::make_tuple(0, 3, 3));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 4));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 6));
+      ASSERT_EQ(counts, 1);
+    }
+  }
+
+  prob = {
+      NDArray::FromVector(std::vector<FloatType>({.0, .5, .0, .0})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5})),
+      NDArray::FromVector(std::vector<FloatType>({.5}))};
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+    } else {
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 2, 2)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 3, 3)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCOOPerEtypeSampling) {
+  _TestCOOPerEtypeSampling<int32_t, float>(true);
+  _TestCOOPerEtypeSampling<int64_t, float>(true);
+  _TestCOOPerEtypeSampling<int32_t, double>(true);
+  _TestCOOPerEtypeSampling<int64_t, double>(true);
+  _TestCOOPerEtypeSampling<int32_t, float>(false);
+  _TestCOOPerEtypeSampling<int64_t, float>(false);
+  _TestCOOPerEtypeSampling<int32_t, double>(false);
+  _TestCOOPerEtypeSampling<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCOOPerEtypeSamplingUniform(bool has_data) {
+  auto pair = COOEtypes<Idx>(has_data);
+  auto mat = pair.first;
+  auto eid2etype_offset = pair.second;
+  std::vector<FloatArray> prob = {
+      aten::NullArray(), aten::NullArray(), aten::NullArray(),
+      aten::NullArray()};
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, true);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = COORowWisePerEtypeSampling(
+        mat, rows, eid2etype_offset, {2, 2, 2, 2}, prob, false);
+    CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 2, 4));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(0, 3, 6));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 2));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 3));
+      ASSERT_EQ(counts, 1);
+    } else {
+      int counts = 0;
+      counts += eset.count(std::make_tuple(0, 0, 0));
+      counts += eset.count(std::make_tuple(0, 1, 1));
+      counts += eset.count(std::make_tuple(0, 2, 2));
+      counts += eset.count(std::make_tuple(0, 3, 3));
+      ASSERT_EQ(counts, 2);
+      counts = 0;
+      counts += eset.count(std::make_tuple(1, 1, 4));
+      ASSERT_EQ(counts, 0);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 3, 5));
+      ASSERT_EQ(counts, 1);
+      counts = 0;
+      counts += eset.count(std::make_tuple(3, 2, 6));
+      ASSERT_EQ(counts, 1);
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCOOPerEtypeSamplingUniform) {
+  _TestCOOPerEtypeSamplingUniform<int32_t, float>(true);
+  _TestCOOPerEtypeSamplingUniform<int64_t, float>(true);
+  _TestCOOPerEtypeSamplingUniform<int32_t, double>(true);
+  _TestCOOPerEtypeSamplingUniform<int64_t, double>(true);
+  _TestCOOPerEtypeSamplingUniform<int32_t, float>(false);
+  _TestCOOPerEtypeSamplingUniform<int64_t, float>(false);
+  _TestCOOPerEtypeSamplingUniform<int32_t, double>(false);
+  _TestCOOPerEtypeSamplingUniform<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRTopk(bool has_data) {
+  auto mat = CSR<Idx>(has_data);
+  FloatArray weight =
+      NDArray::FromVector(std::vector<FloatType>({.1f, .0f, -.1f, .2f, .5f}));
+  // -.1, .2, .1, .0, .5
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+
+  {
+    auto rst = CSRRowWiseTopk(mat, rows, 1, weight, true);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 2);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+    }
+  }
+
+  {
+    auto rst = CSRRowWiseTopk(mat, rows, 1, weight, false);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 2);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCSRTopk) {
+  _TestCSRTopk<int32_t, float>(true);
+  _TestCSRTopk<int64_t, float>(true);
+  _TestCSRTopk<int32_t, double>(true);
+  _TestCSRTopk<int64_t, double>(true);
+  _TestCSRTopk<int32_t, float>(false);
+  _TestCSRTopk<int64_t, float>(false);
+  _TestCSRTopk<int32_t, double>(false);
+  _TestCSRTopk<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCOOTopk(bool has_data) {
+  auto mat = COO<Idx>(has_data);
+  FloatArray weight =
+      NDArray::FromVector(std::vector<FloatType>({.1f, .0f, -.1f, .2f, .5f}));
+  // -.1, .2, .1, .0, .5
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 3}));
+
+  {
+    auto rst = COORowWiseTopk(mat, rows, 1, weight, true);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 2);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 1)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 2, 3)));
+    }
+  }
+
+  {
+    auto rst = COORowWiseTopk(mat, rows, 1, weight, false);
+    auto eset = ToEdgeSet<Idx>(rst);
+    ASSERT_EQ(eset.size(), 2);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCOOTopk) {
+  _TestCOOTopk<int32_t, float>(true);
+  _TestCOOTopk<int64_t, float>(true);
+  _TestCOOTopk<int32_t, double>(true);
+  _TestCOOTopk<int64_t, double>(true);
+  _TestCOOTopk<int32_t, float>(false);
+  _TestCOOTopk<int64_t, float>(false);
+  _TestCOOTopk<int32_t, double>(false);
+  _TestCOOTopk<int64_t, double>(false);
+}
+
+template <typename Idx, typename FloatType>
+void _TestCSRSamplingBiased(bool has_data) {
+  auto mat = CSR<Idx>(has_data);
+  // 0 - 0,1
+  // 1 - 1
+  // 3 - 2,3
+  NDArray tag_offset = NDArray::FromVector(
+      std::vector<Idx>({0, 1, 2, 0, 0, 1, 0, 0, 0, 0, 1, 2}));
+  tag_offset = tag_offset.CreateView({4, 3}, tag_offset->dtype);
+  IdArray rows = NDArray::FromVector(std::vector<Idx>({0, 1, 3}));
+  FloatArray bias = NDArray::FromVector(std::vector<FloatType>({0, 0.5}));
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSamplingBiased(mat, rows, 1, tag_offset, bias, false);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(1, 1, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(1, 1, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+    }
+  }
+  for (int k = 0; k < 10; ++k) {
+    auto rst = CSRRowWiseSamplingBiased(mat, rows, 3, tag_offset, bias, true);
+    CheckSampledResult<Idx>(rst, rows, has_data);
+    auto eset = ToEdgeSet<Idx>(rst);
+    if (has_data) {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 3)));
+      ASSERT_TRUE(eset.count(std::make_tuple(1, 1, 0)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 2)));
+      ASSERT_FALSE(eset.count(std::make_tuple(3, 2, 1)));
+    } else {
+      ASSERT_TRUE(eset.count(std::make_tuple(0, 1, 1)));
+      ASSERT_TRUE(eset.count(std::make_tuple(1, 1, 2)));
+      ASSERT_TRUE(eset.count(std::make_tuple(3, 3, 4)));
+      ASSERT_FALSE(eset.count(std::make_tuple(0, 0, 0)));
+      ASSERT_FALSE(eset.count(std::make_tuple(3, 2, 3)));
+    }
+  }
+}
+
+TEST(RowwiseTest, TestCSRSamplingBiased) {
+  _TestCSRSamplingBiased<int32_t, float>(true);
+  _TestCSRSamplingBiased<int32_t, float>(false);
+  _TestCSRSamplingBiased<int64_t, float>(true);
+  _TestCSRSamplingBiased<int64_t, float>(false);
+  _TestCSRSamplingBiased<int32_t, double>(true);
+  _TestCSRSamplingBiased<int32_t, double>(false);
+  _TestCSRSamplingBiased<int64_t, double>(true);
+  _TestCSRSamplingBiased<int64_t, double>(false);
+}

tests/cpp/test_sampler.cc

@@ -0,0 +1,280 @@
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <iostream>
+#include <vector>
+
+#include "../../src/random/cpu/sample_utils.h"
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::aten;
+
+// TODO: adapt this to Random::Choice
+
+template <typename Idx, typename DType>
+void _TestWithReplacement(RandomEngine* re) {
+  Idx n_categories = 100;
+  Idx n_rolls = 1000000;
+  std::vector<DType> _prob;
+  DType accum = 0.;
+  for (Idx i = 0; i < n_categories; ++i) {
+    _prob.push_back(re->Uniform<DType>());
+    accum += _prob.back();
+  }
+  for (Idx i = 0; i < n_categories; ++i) _prob[i] /= accum;
+  FloatArray prob = NDArray::FromVector(_prob);
+
+  auto _check_given_sampler = [n_categories, n_rolls,
+                               &_prob](utils::BaseSampler<Idx>* s) {
+    std::vector<Idx> counter(n_categories, 0);
+    for (Idx i = 0; i < n_rolls; ++i) {
+      Idx dice = s->Draw();
+      counter[dice]++;
+    }
+    for (Idx i = 0; i < n_categories; ++i)
+      ASSERT_NEAR(static_cast<DType>(counter[i]) / n_rolls, _prob[i], 1e-2);
+  };
+
+  auto _check_random_choice = [n_categories, n_rolls, &_prob, prob]() {
+    std::vector<int64_t> counter(n_categories, 0);
+    for (Idx i = 0; i < n_rolls; ++i) {
+      Idx dice = RandomEngine::ThreadLocal()->Choice<int64_t>(prob);
+      counter[dice]++;
+    }
+    for (Idx i = 0; i < n_categories; ++i)
+      ASSERT_NEAR(static_cast<DType>(counter[i]) / n_rolls, _prob[i], 1e-2);
+  };
+
+  utils::AliasSampler<Idx, DType, true> as(re, prob);
+  utils::CDFSampler<Idx, DType, true> cs(re, prob);
+  utils::TreeSampler<Idx, DType, true> ts(re, prob);
+  _check_given_sampler(&as);
+  _check_given_sampler(&cs);
+  _check_given_sampler(&ts);
+  _check_random_choice();
+}
+
+TEST(SampleUtilsTest, TestWithReplacement) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  re->SetSeed(42);
+  _TestWithReplacement<int32_t, float>(re);
+  re->SetSeed(42);
+  _TestWithReplacement<int32_t, double>(re);
+  re->SetSeed(42);
+  _TestWithReplacement<int64_t, float>(re);
+  re->SetSeed(42);
+  _TestWithReplacement<int64_t, double>(re);
+};
+
+template <typename Idx, typename DType>
+void _TestWithoutReplacementOrder(RandomEngine* re) {
+  // TODO(BarclayII): is there a reliable way to do this test?
+  std::vector<DType> _prob = {1e6f, 1e-6f, 1e-2f, 1e2f};
+  FloatArray prob = NDArray::FromVector(_prob);
+  std::vector<Idx> ground_truth = {0, 3, 2, 1};
+
+  auto _check_given_sampler = [&ground_truth](utils::BaseSampler<Idx>* s) {
+    for (size_t i = 0; i < ground_truth.size(); ++i) {
+      Idx dice = s->Draw();
+      ASSERT_EQ(dice, ground_truth[i]);
+    }
+  };
+
+  utils::AliasSampler<Idx, DType, false> as(re, prob);
+  utils::CDFSampler<Idx, DType, false> cs(re, prob);
+  utils::TreeSampler<Idx, DType, false> ts(re, prob);
+  _check_given_sampler(&as);
+  _check_given_sampler(&cs);
+  _check_given_sampler(&ts);
+}
+
+TEST(SampleUtilsTest, TestWithoutReplacementOrder) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  re->SetSeed(42);
+  _TestWithoutReplacementOrder<int32_t, float>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementOrder<int32_t, double>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementOrder<int64_t, float>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementOrder<int64_t, double>(re);
+};
+
+template <typename Idx, typename DType>
+void _TestWithoutReplacementUnique(RandomEngine* re) {
+  Idx N = 1000000;
+  std::vector<DType> _likelihood;
+  for (Idx i = 0; i < N; ++i) _likelihood.push_back(re->Uniform<DType>());
+  FloatArray likelihood = NDArray::FromVector(_likelihood);
+
+  auto _check_given_sampler = [N](utils::BaseSampler<Idx>* s) {
+    std::vector<int> cnt(N, 0);
+    for (Idx i = 0; i < N; ++i) {
+      Idx dice = s->Draw();
+      cnt[dice]++;
+    }
+    for (Idx i = 0; i < N; ++i) ASSERT_EQ(cnt[i], 1);
+  };
+
+  utils::AliasSampler<Idx, DType, false> as(re, likelihood);
+  utils::CDFSampler<Idx, DType, false> cs(re, likelihood);
+  utils::TreeSampler<Idx, DType, false> ts(re, likelihood);
+  _check_given_sampler(&as);
+  _check_given_sampler(&cs);
+  _check_given_sampler(&ts);
+}
+
+TEST(SampleUtilsTest, TestWithoutReplacementUnique) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  re->SetSeed(42);
+  _TestWithoutReplacementUnique<int32_t, float>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementUnique<int32_t, double>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementUnique<int64_t, float>(re);
+  re->SetSeed(42);
+  _TestWithoutReplacementUnique<int64_t, double>(re);
+};
+
+template <typename Idx, typename DType>
+void _TestChoice(RandomEngine* re) {
+  re->SetSeed(42);
+  std::vector<DType> prob_vec = {1., 0., 0., 0., 2., 2., 0., 0.};
+  FloatArray prob = FloatArray::FromVector(prob_vec);
+  {
+    for (int k = 0; k < 1000; ++k) {
+      Idx x = re->Choice<Idx>(prob);
+      ASSERT_TRUE(x == 0 || x == 4 || x == 5);
+    }
+  }
+  // num = 0
+  {
+    IdArray rst = re->Choice<Idx, DType>(0, prob, true);
+    ASSERT_EQ(rst->shape[0], 0);
+  }
+  // w/ replacement
+  {
+    IdArray rst = re->Choice<Idx, DType>(1000, prob, true);
+    ASSERT_EQ(rst->shape[0], 1000);
+    for (int64_t i = 0; i < 1000; ++i) {
+      Idx x = static_cast<Idx*>(rst->data)[i];
+      ASSERT_TRUE(x == 0 || x == 4 || x == 5);
+    }
+  }
+  // w/o replacement
+  {
+    IdArray rst = re->Choice<Idx, DType>(3, prob, false);
+    ASSERT_EQ(rst->shape[0], 3);
+    std::set<Idx> idxset;
+    for (int64_t i = 0; i < 3; ++i) {
+      Idx x = static_cast<Idx*>(rst->data)[i];
+      idxset.insert(x);
+    }
+    ASSERT_EQ(idxset.size(), 3);
+    ASSERT_EQ(idxset.count(0), 1);
+    ASSERT_EQ(idxset.count(4), 1);
+    ASSERT_EQ(idxset.count(5), 1);
+  }
+}
+
+TEST(RandomTest, TestChoice) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  _TestChoice<int32_t, float>(re);
+  _TestChoice<int64_t, float>(re);
+  _TestChoice<int32_t, double>(re);
+  _TestChoice<int64_t, double>(re);
+}
+
+template <typename Idx>
+void _TestUniformChoice(RandomEngine* re) {
+  re->SetSeed(42);
+  // num == 0
+  {
+    IdArray rst = re->UniformChoice<Idx>(0, 100, true);
+    ASSERT_EQ(rst->shape[0], 0);
+  }
+  // w/ replacement
+  {
+    IdArray rst = re->UniformChoice<Idx>(1000, 100, true);
+    ASSERT_EQ(rst->shape[0], 1000);
+    for (int64_t i = 0; i < 1000; ++i) {
+      Idx x = static_cast<Idx*>(rst->data)[i];
+      ASSERT_TRUE(x >= 0 && x < 100);
+    }
+  }
+  // w/o replacement
+  {
+    IdArray rst = re->UniformChoice<Idx>(99, 100, false);
+    ASSERT_EQ(rst->shape[0], 99);
+    std::set<Idx> idxset;
+    for (int64_t i = 0; i < 99; ++i) {
+      Idx x = static_cast<Idx*>(rst->data)[i];
+      ASSERT_TRUE(x >= 0 && x < 100);
+      idxset.insert(x);
+    }
+    ASSERT_EQ(idxset.size(), 99);
+  }
+}
+
+TEST(RandomTest, TestUniformChoice) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  _TestUniformChoice<int32_t>(re);
+  _TestUniformChoice<int64_t>(re);
+  _TestUniformChoice<int32_t>(re);
+  _TestUniformChoice<int64_t>(re);
+}
+
+template <typename Idx, typename FloatType>
+void _TestBiasedChoice(RandomEngine* re) {
+  re->SetSeed(42);
+  // num == 0
+  {
+    Idx split[] = {0, 1, 2};
+    FloatArray bias = NDArray::FromVector(std::vector<FloatType>({1, 3}));
+    IdArray rst = re->BiasedChoice<Idx, FloatType>(0, split, bias, true);
+    ASSERT_EQ(rst->shape[0], 0);
+  }
+  // basic test
+  {
+    Idx sample_num = 100000;
+    Idx population = 1000000;
+    Idx split[] = {0, population / 2, population};
+    FloatArray bias = NDArray::FromVector(std::vector<FloatType>({1, 3}));
+
+    IdArray rst =
+        re->BiasedChoice<Idx, FloatType>(sample_num, split, bias, true);
+    auto rst_data = static_cast<Idx*>(rst->data);
+    Idx larger = 0;
+    for (Idx i = 0; i < sample_num; ++i)
+      if (rst_data[i] >= population / 2) larger++;
+    ASSERT_LE(fabs((double)larger / sample_num - 0.75), 1e-2);
+  }
+  // without replacement
+  {
+    Idx sample_num = 500;
+    Idx population = 1000;
+    Idx split[] = {0, sample_num, population};
+    FloatArray bias = NDArray::FromVector(std::vector<FloatType>({1, 0}));
+
+    IdArray rst =
+        re->BiasedChoice<Idx, FloatType>(sample_num, split, bias, false);
+    auto rst_data = static_cast<Idx*>(rst->data);
+
+    std::set<Idx> idxset;
+    for (int64_t i = 0; i < sample_num; ++i) {
+      Idx x = rst_data[i];
+      ASSERT_LT(x, sample_num);
+      idxset.insert(x);
+    }
+    ASSERT_EQ(idxset.size(), sample_num);
+  }
+}
+
+TEST(RandomTest, TestBiasedChoice) {
+  RandomEngine* re = RandomEngine::ThreadLocal();
+  _TestBiasedChoice<int32_t, float>(re);
+  _TestBiasedChoice<int64_t, float>(re);
+  _TestBiasedChoice<int32_t, double>(re);
+  _TestBiasedChoice<int64_t, double>(re);
+}

tests/cpp/test_serialize.cc

@@ -0,0 +1,103 @@
+#include <dgl/graph_serializer.h>
+#include <dgl/immutable_graph.h>
+#include <dmlc/memory_io.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <iostream>
+#include <memory>
+#include <vector>
+
+#include "../../src/graph/heterograph.h"
+#include "../../src/graph/unit_graph.h"
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::aten;
+using namespace dmlc;
+
+TEST(Serialize, UnitGraph_COO) {
+  aten::CSRMatrix csr_matrix;
+  auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto mg = std::dynamic_pointer_cast<UnitGraph>(
+      dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, COO_CODE));
+
+  std::string blob;
+  dmlc::MemoryStringStream ifs(&blob);
+
+  static_cast<dmlc::Stream *>(&ifs)->Write(mg);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  auto ug2 = Serializer::make_shared<UnitGraph>();
+  static_cast<dmlc::Stream *>(&ofs)->Read(&ug2);
+  EXPECT_EQ(ug2->NumVertices(0), 9);
+  EXPECT_EQ(ug2->NumVertices(1), 8);
+  EXPECT_EQ(ug2->NumEdges(0), 4);
+  EXPECT_EQ(ug2->FindEdge(0, 1).first, 2);
+  EXPECT_EQ(ug2->FindEdge(0, 1).second, 6);
+}
+
+TEST(Serialize, UnitGraph_CSR) {
+  aten::CSRMatrix csr_matrix;
+  auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto coo_g = std::dynamic_pointer_cast<UnitGraph>(
+      dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst));
+  auto csr_g =
+      std::dynamic_pointer_cast<UnitGraph>(coo_g->GetGraphInFormat(CSR_CODE));
+
+  std::string blob;
+  dmlc::MemoryStringStream ifs(&blob);
+
+  static_cast<dmlc::Stream *>(&ifs)->Write(csr_g);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  auto ug2 = Serializer::make_shared<UnitGraph>();
+  static_cast<dmlc::Stream *>(&ofs)->Read(&ug2);
+  // Query operation is not supported on CSR, how to check it?
+}
+
+TEST(Serialize, ImmutableGraph) {
+  auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto gptr = ImmutableGraph::CreateFromCOO(10, src, dst);
+  std::string blob;
+  dmlc::MemoryStringStream ifs(&blob);
+
+  static_cast<dmlc::Stream *>(&ifs)->Write(gptr);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  auto rptr_read = dgl::Serializer::make_shared<ImmutableGraph>();
+  static_cast<dmlc::Stream *>(&ofs)->Read(&rptr_read);
+  EXPECT_EQ(rptr_read->NumEdges(), 4);
+  EXPECT_EQ(rptr_read->NumVertices(), 10);
+  EXPECT_EQ(rptr_read->FindEdge(2).first, 5);
+  EXPECT_EQ(rptr_read->FindEdge(2).second, 2);
+}
+
+TEST(Serialize, HeteroGraph) {
+  auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto mg1 = dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst);
+  src = VecToIdArray<int64_t>({6, 2, 5, 1, 8});
+  dst = VecToIdArray<int64_t>({5, 2, 4, 8, 0});
+  auto mg2 = dgl::UnitGraph::CreateFromCOO(1, 9, 9, src, dst);
+  std::vector<HeteroGraphPtr> relgraphs;
+  relgraphs.push_back(mg1);
+  relgraphs.push_back(mg2);
+  src = VecToIdArray<int64_t>({0, 0});
+  dst = VecToIdArray<int64_t>({1, 0});
+  auto meta_gptr = ImmutableGraph::CreateFromCOO(3, src, dst);
+  auto hrptr = std::make_shared<HeteroGraph>(meta_gptr, relgraphs);
+
+  std::string blob;
+  dmlc::MemoryStringStream ifs(&blob);
+  static_cast<dmlc::Stream *>(&ifs)->Write(hrptr);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  auto gptr = dgl::Serializer::make_shared<HeteroGraph>();
+  static_cast<dmlc::Stream *>(&ofs)->Read(&gptr);
+  EXPECT_EQ(gptr->NumVertices(0), 9);
+  EXPECT_EQ(gptr->NumVertices(1), 8);
+}

tests/cpp/test_smart_ptr_serialize.cc

@@ -0,0 +1,100 @@
+#include <dgl/runtime/serializer.h>
+#include <dgl/runtime/smart_ptr_serializer.h>
+#include <dmlc/io.h>
+#include <dmlc/logging.h>
+#include <dmlc/memory_io.h>
+#include <dmlc/parameter.h>
+#include <gtest/gtest.h>
+
+#include <cstring>
+#include <iostream>
+#include <sstream>
+#include <unordered_map>
+
+using namespace std;
+
+class MyClass {
+ public:
+  MyClass() {}
+  MyClass(std::string data) : data_(data) {}
+  inline void Save(dmlc::Stream *strm) const { strm->Write(this->data_); }
+  inline bool Load(dmlc::Stream *strm) { return strm->Read(&data_); }
+  inline bool operator==(const MyClass &other) const {
+    return data_ == other.data_;
+  }
+
+ public:
+  std::string data_;
+};
+// need to declare the traits property of my class to dmlc
+namespace dmlc {
+DMLC_DECLARE_TRAITS(has_saveload, MyClass, true);
+}
+
+template <typename T>
+class SmartPtrTest : public ::testing::Test {
+ public:
+  typedef T SmartPtr;
+};
+
+using SmartPtrTypes =
+    ::testing::Types<std::shared_ptr<MyClass>, std::unique_ptr<MyClass>>;
+TYPED_TEST_SUITE(SmartPtrTest, SmartPtrTypes);
+
+TYPED_TEST(SmartPtrTest, Obj_Test) {
+  std::string blob;
+  dmlc::MemoryStringStream fs(&blob);
+  using SmartPtr = typename TestFixture::SmartPtr;
+  auto myc = SmartPtr(new MyClass("1111"));
+  { static_cast<dmlc::Stream *>(&fs)->Write(myc); }
+  fs.Seek(0);
+  auto copy_data = SmartPtr(new MyClass());
+  CHECK(static_cast<dmlc::Stream *>(&fs)->Read(&copy_data));
+
+  EXPECT_EQ(myc->data_, copy_data->data_);
+}
+
+TYPED_TEST(SmartPtrTest, Vector_Test1) {
+  std::string blob;
+  dmlc::MemoryStringStream fs(&blob);
+  using SmartPtr = typename TestFixture::SmartPtr;
+  typedef std::pair<std::string, SmartPtr> Pair;
+
+  std::vector<Pair> myclasses;
+  myclasses.emplace_back("a", SmartPtr(new MyClass("@A@B")));
+  myclasses.emplace_back("b", SmartPtr(new MyClass("2222")));
+  static_cast<dmlc::Stream *>(&fs)->Write<std::vector<Pair>>(myclasses);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  std::vector<Pair> copy_myclasses;
+  static_cast<dmlc::Stream *>(&ofs)->Read<std::vector<Pair>>(&copy_myclasses);
+
+  EXPECT_TRUE(std::equal(
+      myclasses.begin(), myclasses.end(), copy_myclasses.begin(),
+      [](const Pair &left, const Pair &right) {
+        return (left.second->data_ == right.second->data_) &&
+               (left.first == right.first);
+      }));
+}
+
+TYPED_TEST(SmartPtrTest, Vector_Test2) {
+  std::string blob;
+  dmlc::MemoryStringStream fs(&blob);
+  using SmartPtr = typename TestFixture::SmartPtr;
+
+  std::vector<SmartPtr> myclasses;
+  myclasses.emplace_back(new MyClass("@A@"));
+  myclasses.emplace_back(new MyClass("2222"));
+  static_cast<dmlc::Stream *>(&fs)->Write<std::vector<SmartPtr>>(myclasses);
+
+  dmlc::MemoryStringStream ofs(&blob);
+  std::vector<SmartPtr> copy_myclasses;
+  static_cast<dmlc::Stream *>(&ofs)->Read<std::vector<SmartPtr>>(
+      &copy_myclasses);
+
+  EXPECT_TRUE(std::equal(
+      myclasses.begin(), myclasses.end(), copy_myclasses.begin(),
+      [](const SmartPtr &left, const SmartPtr &right) {
+        return left->data_ == right->data_;
+      }));
+}

tests/cpp/test_spmat_coo.cc

@@ -0,0 +1,576 @@
+#include <dgl/array.h>
+#include <dmlc/omp.h>
+#include <gtest/gtest.h>
+#include <omp.h>
+
+#include <random>
+
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+
+namespace {
+
+template <typename IDX>
+aten::CSRMatrix CSR1(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 3, 1, 4]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 3, 5, 5}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 3, 4, 1}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::CSRMatrix CSR2(DGLContext ctx = CTX) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 5, 3, 1, 4]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::COOMatrix COO1(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 3, 1, 4]
+  // row : [0, 2, 0, 1, 2]
+  // col : [1, 2, 2, 0, 3]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 1, 2, 4}), sizeof(IDX) * 8, ctx));
+}
+
+template <typename IDX>
+aten::COOMatrix COO2(DGLContext ctx = CTX) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 5, 3, 1, 4]
+  // row : [0, 2, 0, 1, 2, 0]
+  // col : [1, 2, 2, 0, 3, 2]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 3, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 1, 2, 3, 4, 5}), sizeof(IDX) * 8, ctx));
+}
+
+template <typename IDX>
+aten::CSRMatrix SR_CSR3(DGLContext ctx) {
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::CSRMatrix SRC_CSR3(DGLContext ctx) {
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 0, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::COOMatrix COO3(DGLContext ctx) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // row : [0, 2, 0, 1, 2, 0]
+  // col : [2, 2, 1, 0, 3, 2]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 2, 1, 0, 3, 2}), sizeof(IDX) * 8, ctx));
+}
+
+template <typename IDX>
+aten::COOMatrix COORandomized(IDX rows_and_cols, int64_t nnz, int seed) {
+  std::vector<IDX> vec_rows(nnz);
+  std::vector<IDX> vec_cols(nnz);
+  std::vector<IDX> vec_data(nnz);
+
+#pragma omp parallel
+  {
+    const int64_t num_threads = omp_get_num_threads();
+    const int64_t thread_id = omp_get_thread_num();
+    const int64_t chunk = nnz / num_threads;
+    const int64_t size = (thread_id == num_threads - 1)
+                             ? nnz - chunk * (num_threads - 1)
+                             : chunk;
+    auto rows = vec_rows.data() + thread_id * chunk;
+    auto cols = vec_cols.data() + thread_id * chunk;
+    auto data = vec_data.data() + thread_id * chunk;
+
+    std::mt19937_64 gen64(seed + thread_id);
+    std::mt19937 gen32(seed + thread_id);
+
+    for (int64_t i = 0; i < size; ++i) {
+      rows[i] = gen64() % rows_and_cols;
+      cols[i] = gen64() % rows_and_cols;
+      data[i] = gen32() % 90 + 1;
+    }
+  }
+
+  return aten::COOMatrix(
+      rows_and_cols, rows_and_cols,
+      aten::VecToIdArray(vec_rows, sizeof(IDX) * 8, CTX),
+      aten::VecToIdArray(vec_cols, sizeof(IDX) * 8, CTX),
+      aten::VecToIdArray(vec_data, sizeof(IDX) * 8, CTX), false, false);
+}
+
+struct SparseCOOCSR {
+  static constexpr uint64_t NUM_ROWS = 100;
+  static constexpr uint64_t NUM_COLS = 150;
+  static constexpr uint64_t NUM_NZ = 5;
+  template <typename IDX>
+  static aten::COOMatrix COOSparse(const DGLContext &ctx = CTX) {
+    return aten::COOMatrix(
+        NUM_ROWS, NUM_COLS,
+        aten::VecToIdArray(
+            std::vector<IDX>({0, 1, 2, 3, 4}), sizeof(IDX) * 8, ctx),
+        aten::VecToIdArray(
+            std::vector<IDX>({1, 2, 3, 4, 5}), sizeof(IDX) * 8, ctx));
+  }
+
+  template <typename IDX>
+  static aten::CSRMatrix CSRSparse(const DGLContext &ctx = CTX) {
+    auto &&indptr = std::vector<IDX>(NUM_ROWS + 1, NUM_NZ);
+    for (size_t i = 0; i < NUM_NZ; ++i) {
+      indptr[i + 1] = static_cast<IDX>(i + 1);
+    }
+    indptr[0] = 0;
+    return aten::CSRMatrix(
+        NUM_ROWS, NUM_COLS, aten::VecToIdArray(indptr, sizeof(IDX) * 8, ctx),
+        aten::VecToIdArray(
+            std::vector<IDX>({1, 2, 3, 4, 5}), sizeof(IDX) * 8, ctx),
+        aten::VecToIdArray(
+            std::vector<IDX>({1, 1, 1, 1, 1}), sizeof(IDX) * 8, ctx),
+        false);
+  }
+};
+
+template <typename IDX>
+aten::COOMatrix RowSorted_NullData_COO(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // row : [0, 0, 1, 2, 2]
+  // col : [1, 2, 0, 2, 3]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::NullArray(), true, false);
+}
+
+template <typename IDX>
+aten::CSRMatrix RowSorted_NullData_CSR(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 1, 2, 3, 4]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 3, 5, 5}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 1, 2, 3, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+}  // namespace
+
+template <typename IDX>
+void _TestCOOToCSR(DGLContext ctx) {
+  auto coo = COO1<IDX>(ctx);
+  auto csr = CSR1<IDX>(ctx);
+  auto tcsr = aten::COOToCSR(coo);
+  ASSERT_FALSE(coo.row_sorted);
+  ASSERT_EQ(csr.num_rows, tcsr.num_rows);
+  ASSERT_EQ(csr.num_cols, tcsr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(csr.indptr, tcsr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(csr.indices, tcsr.indices));
+
+  coo = COO2<IDX>(ctx);
+  csr = CSR2<IDX>(ctx);
+  tcsr = aten::COOToCSR(coo);
+  ASSERT_EQ(coo.num_rows, csr.num_rows);
+  ASSERT_EQ(coo.num_cols, csr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(csr.indptr, tcsr.indptr));
+
+  // Convert from row sorted coo
+  coo = COO1<IDX>(ctx);
+  auto rs_coo = aten::COOSort(coo, false);
+  auto rs_csr = CSR1<IDX>(ctx);
+  auto rs_tcsr = aten::COOToCSR(rs_coo);
+  ASSERT_TRUE(rs_coo.row_sorted);
+  ASSERT_EQ(coo.num_rows, rs_tcsr.num_rows);
+  ASSERT_EQ(coo.num_cols, rs_tcsr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_csr.indptr, rs_tcsr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_tcsr.indices, rs_coo.col));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_tcsr.data, rs_coo.data));
+
+  coo = COO3<IDX>(ctx);
+  rs_coo = aten::COOSort(coo, false);
+  rs_csr = SR_CSR3<IDX>(ctx);
+  rs_tcsr = aten::COOToCSR(rs_coo);
+  ASSERT_EQ(coo.num_rows, rs_tcsr.num_rows);
+  ASSERT_EQ(coo.num_cols, rs_tcsr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_csr.indptr, rs_tcsr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_tcsr.indices, rs_coo.col));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_tcsr.data, rs_coo.data));
+
+  rs_coo = RowSorted_NullData_COO<IDX>(ctx);
+  ASSERT_TRUE(rs_coo.row_sorted);
+  rs_csr = RowSorted_NullData_CSR<IDX>(ctx);
+  rs_tcsr = aten::COOToCSR(rs_coo);
+  ASSERT_EQ(coo.num_rows, rs_tcsr.num_rows);
+  ASSERT_EQ(rs_csr.num_rows, rs_tcsr.num_rows);
+  ASSERT_EQ(coo.num_cols, rs_tcsr.num_cols);
+  ASSERT_EQ(rs_csr.num_cols, rs_tcsr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_csr.indptr, rs_tcsr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_csr.indices, rs_tcsr.indices));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_csr.data, rs_tcsr.data));
+  ASSERT_TRUE(ArrayEQ<IDX>(rs_coo.col, rs_tcsr.indices));
+  ASSERT_FALSE(ArrayEQ<IDX>(rs_coo.data, rs_tcsr.data));
+
+  // Convert from col sorted coo
+  coo = COO1<IDX>(ctx);
+  auto src_coo = aten::COOSort(coo, true);
+  auto src_csr = CSR1<IDX>(ctx);
+  auto src_tcsr = aten::COOToCSR(src_coo);
+  ASSERT_EQ(coo.num_rows, src_tcsr.num_rows);
+  ASSERT_EQ(coo.num_cols, src_tcsr.num_cols);
+  ASSERT_TRUE(src_tcsr.sorted);
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.indptr, src_csr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.indices, src_coo.col));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.data, src_coo.data));
+
+  coo = COO3<IDX>(ctx);
+  src_coo = aten::COOSort(coo, true);
+  src_csr = SRC_CSR3<IDX>(ctx);
+  src_tcsr = aten::COOToCSR(src_coo);
+  ASSERT_EQ(coo.num_rows, src_tcsr.num_rows);
+  ASSERT_EQ(coo.num_cols, src_tcsr.num_cols);
+  ASSERT_TRUE(src_tcsr.sorted);
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.indptr, src_csr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.indices, src_coo.col));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_tcsr.data, src_coo.data));
+
+  coo = SparseCOOCSR::COOSparse<IDX>(ctx);
+  csr = SparseCOOCSR::CSRSparse<IDX>(ctx);
+  tcsr = aten::COOToCSR(coo);
+  ASSERT_FALSE(coo.row_sorted);
+  ASSERT_EQ(csr.num_rows, tcsr.num_rows);
+  ASSERT_EQ(csr.num_cols, tcsr.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(csr.indptr, tcsr.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(csr.indices, tcsr.indices));
+}
+
+TEST(SpmatTest, COOToCSR) {
+  _TestCOOToCSR<int32_t>(CPU);
+  _TestCOOToCSR<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCOOToCSR<int32_t>(GPU);
+  _TestCOOToCSR<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCOOHasDuplicate() {
+  auto coo = COO1<IDX>();
+  ASSERT_FALSE(aten::COOHasDuplicate(coo));
+  coo = COO2<IDX>();
+  ASSERT_TRUE(aten::COOHasDuplicate(coo));
+}
+
+TEST(SpmatTest, TestCOOHasDuplicate) {
+  _TestCOOHasDuplicate<int32_t>();
+  _TestCOOHasDuplicate<int64_t>();
+}
+
+template <typename IDX>
+void _TestCOOSort(DGLContext ctx) {
+  auto coo = COO3<IDX>(ctx);
+
+  auto sr_coo = COOSort(coo, false);
+  ASSERT_EQ(coo.num_rows, sr_coo.num_rows);
+  ASSERT_EQ(coo.num_cols, sr_coo.num_cols);
+  ASSERT_TRUE(sr_coo.row_sorted);
+  auto flags = COOIsSorted(sr_coo);
+  ASSERT_TRUE(flags.first);
+  flags = COOIsSorted(coo);  // original coo should stay the same
+  ASSERT_FALSE(flags.first);
+  ASSERT_FALSE(flags.second);
+
+  auto src_coo = COOSort(coo, true);
+  ASSERT_EQ(coo.num_rows, src_coo.num_rows);
+  ASSERT_EQ(coo.num_cols, src_coo.num_cols);
+  ASSERT_TRUE(src_coo.row_sorted);
+  ASSERT_TRUE(src_coo.col_sorted);
+  flags = COOIsSorted(src_coo);
+  ASSERT_TRUE(flags.first);
+  ASSERT_TRUE(flags.second);
+
+  // sort inplace
+  COOSort_(&coo);
+  ASSERT_TRUE(coo.row_sorted);
+  flags = COOIsSorted(coo);
+  ASSERT_TRUE(flags.first);
+  COOSort_(&coo, true);
+  ASSERT_TRUE(coo.row_sorted);
+  ASSERT_TRUE(coo.col_sorted);
+  flags = COOIsSorted(coo);
+  ASSERT_TRUE(flags.first);
+  ASSERT_TRUE(flags.second);
+
+  // COO3
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 1, 2, 3, 4, 5]
+  // row : [0, 2, 0, 1, 2, 0]
+  // col : [2, 2, 1, 0, 3, 2]
+  // Row Sorted
+  // data: [0, 2, 5, 3, 1, 4]
+  // row : [0, 0, 0, 1, 2, 2]
+  // col : [2, 1, 2, 0, 2, 3]
+  // Row Col Sorted
+  // data: [2, 0, 5, 3, 1, 4]
+  // row : [0, 0, 0, 1, 2, 2]
+  // col : [1, 2, 2, 0, 2, 3]
+  auto sort_row = aten::VecToIdArray(
+      std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
+  auto sort_col = aten::VecToIdArray(
+      std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx);
+  auto sort_col_data = aten::VecToIdArray(
+      std::vector<IDX>({2, 0, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx);
+
+  ASSERT_TRUE(ArrayEQ<IDX>(sr_coo.row, sort_row));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_coo.row, sort_row));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_coo.col, sort_col));
+  ASSERT_TRUE(ArrayEQ<IDX>(src_coo.data, sort_col_data));
+}
+
+TEST(SpmatTest, COOSort) {
+  _TestCOOSort<int32_t>(CPU);
+  _TestCOOSort<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCOOSort<int32_t>(GPU);
+  _TestCOOSort<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCOOReorder() {
+  auto coo = COO2<IDX>();
+  auto new_row =
+      aten::VecToIdArray(std::vector<IDX>({2, 0, 3, 1}), sizeof(IDX) * 8, CTX);
+  auto new_col = aten::VecToIdArray(
+      std::vector<IDX>({2, 0, 4, 3, 1}), sizeof(IDX) * 8, CTX);
+  auto new_coo = COOReorder(coo, new_row, new_col);
+  ASSERT_EQ(new_coo.num_rows, coo.num_rows);
+  ASSERT_EQ(new_coo.num_cols, coo.num_cols);
+}
+
+TEST(SpmatTest, TestCOOReorder) {
+  _TestCOOReorder<int32_t>();
+  _TestCOOReorder<int64_t>();
+}
+
+template <typename IDX>
+void _TestCOOGetData(DGLContext ctx) {
+  auto coo = COO2<IDX>(ctx);
+  // test get all data
+  auto x = aten::COOGetAllData(coo, 0, 0);
+  auto tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::COOGetAllData(coo, 0, 2);
+  tx = aten::VecToIdArray(std::vector<IDX>({2, 5}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data
+  auto r =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  auto c =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::COOGetData(coo, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data on sorted
+  coo = aten::COOSort(coo);
+  r = aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::COOGetData(coo, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data w/ broadcasting
+  r = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX) * 8, ctx);
+  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::COOGetData(coo, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+TEST(SpmatTest, COOGetData) {
+  _TestCOOGetData<int32_t>(CPU);
+  _TestCOOGetData<int64_t>(CPU);
+  // #ifdef DGL_USE_CUDA
+  //_TestCOOGetData<int32_t>(GPU);
+  //_TestCOOGetData<int64_t>(GPU);
+  // #endif
+}
+
+template <typename IDX>
+void _TestCOOGetDataAndIndices() {
+  auto coo = COO2<IDX>();
+  auto r =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, CTX);
+  auto c =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, CTX);
+  auto x = aten::COOGetDataAndIndices(coo, r, c);
+  auto tr =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, CTX);
+  auto tc =
+      aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX) * 8, CTX);
+  auto td =
+      aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, CTX);
+  ASSERT_TRUE(ArrayEQ<IDX>(x[0], tr));
+  ASSERT_TRUE(ArrayEQ<IDX>(x[1], tc));
+  ASSERT_TRUE(ArrayEQ<IDX>(x[2], td));
+}
+
+TEST(SpmatTest, COOGetDataAndIndices) {
+  _TestCOOGetDataAndIndices<int32_t>();
+  _TestCOOGetDataAndIndices<int64_t>();
+}
+
+template <typename IDX>
+void _TestCOOToCSRAlgs() {
+  // Compare results between different CPU COOToCSR implementations.
+  // NNZ is chosen to be bigger than the limit for the "small" matrix algorithm.
+  // N is set to lay on border between "sparse" and "dense" algorithm choice.
+
+  const int64_t num_threads = std::min(256, omp_get_max_threads());
+  const int64_t min_num_threads = 3;
+
+  if (num_threads < min_num_threads) {
+    std::cerr << "[          ] [ INFO ]"
+              << "This test requires at least 3 OMP threads to work properly"
+              << std::endl;
+    GTEST_SKIP();
+    return;
+  }
+
+  // Select N and NNZ for COO matrix in a way than depending on number of
+  // threads different algorithm will be used.
+  // See WhichCOOToCSR in src/array/cpu/spmat_op_impl_coo.cc for details
+  const int64_t type_scale = sizeof(IDX) >> 1;
+  const int64_t small = 50 * num_threads * type_scale * type_scale;
+  // NNZ should be bigger than limit for small matrix algorithm
+  const int64_t nnz = small + 1234;
+  // N is chosen to lay on sparse/dense border
+  const int64_t n = type_scale * nnz / num_threads;
+  const IDX rows_nad_cols = n + 1;  // should be bigger than sparse/dense border
+
+  // Note that it will be better to set the seed to a random value when gtest
+  // allows to use --gtest_random_seed without --gtest_shuffle and report this
+  // value for reproduction. This way we can find unforeseen situations and
+  // potential bugs.
+  const auto seed = 123321;
+  auto coo = COORandomized<IDX>(rows_nad_cols, nnz, seed);
+
+  omp_set_num_threads(1);
+  // UnSortedSmallCOOToCSR will be used
+  auto tcsr_small = aten::COOToCSR(coo);
+  ASSERT_EQ(coo.num_rows, tcsr_small.num_rows);
+  ASSERT_EQ(coo.num_cols, tcsr_small.num_cols);
+
+  omp_set_num_threads(num_threads - 1);
+  // UnSortedDenseCOOToCSR will be used
+  auto tcsr_dense = aten::COOToCSR(coo);
+  ASSERT_EQ(tcsr_small.num_rows, tcsr_dense.num_rows);
+  ASSERT_EQ(tcsr_small.num_cols, tcsr_dense.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.indptr, tcsr_dense.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.indices, tcsr_dense.indices));
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.data, tcsr_dense.data));
+
+  omp_set_num_threads(num_threads);
+  // UnSortedSparseCOOToCSR will be used
+  auto tcsr_sparse = aten::COOToCSR(coo);
+  ASSERT_EQ(tcsr_small.num_rows, tcsr_sparse.num_rows);
+  ASSERT_EQ(tcsr_small.num_cols, tcsr_sparse.num_cols);
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.indptr, tcsr_sparse.indptr));
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.indices, tcsr_sparse.indices));
+  ASSERT_TRUE(ArrayEQ<IDX>(tcsr_small.data, tcsr_sparse.data));
+  return;
+}
+
+TEST(SpmatTest, COOToCSRAlgs) {
+  _TestCOOToCSRAlgs<int32_t>();
+  _TestCOOToCSRAlgs<int64_t>();
+}

tests/cpp/test_spmat_csr.cc

@@ -0,0 +1,760 @@
+#include <dgl/array.h>
+#include <gtest/gtest.h>
+
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+
+namespace {
+
+template <typename IDX>
+aten::CSRMatrix CSR1(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 3, 1, 4]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 3, 5, 5}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 0, 3, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 3, 4, 1}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::CSRMatrix CSR2(DGLContext ctx = CTX) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 5, 3, 1, 4]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::CSRMatrix CSR3(DGLContext ctx = CTX) {
+  // has duplicate entries and the columns are not sorted
+  // [[0, 1, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0, 0],
+  //  [0, 0, 0, 0, 0, 0],
+  //  [1, 1, 1, 0, 0, 0],
+  //  [0, 0, 0, 1, 0, 0],
+  //  [0, 0, 0, 0, 0, 0],
+  //  [1, 2, 1, 1, 0, 0],
+  //  [0, 1, 0, 0, 0, 1]],
+  // data: [5, 2, 0, 3, 1, 4, 8, 7, 6, 9, 12, 13, 11, 10, 14, 15, 16]
+  return aten::CSRMatrix(
+      9, 6,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6, 9, 10, 10, 15, 17}), sizeof(IDX) * 8,
+          ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({3, 2, 1, 0, 2, 3, 1, 2, 0, 3, 1, 2, 1, 3, 0, 5, 1}),
+          sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>(
+              {0, 2, 5, 3, 1, 4, 6, 8, 7, 9, 13, 10, 11, 14, 12, 16, 15}),
+          sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::COOMatrix COO1(DGLContext ctx = CTX) {
+  // [[0, 1, 1, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 3, 1, 4]
+  // row : [0, 2, 0, 1, 2]
+  // col : [1, 2, 2, 0, 3]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 1, 2, 4}), sizeof(IDX) * 8, ctx));
+}
+
+template <typename IDX>
+aten::COOMatrix COO2(DGLContext ctx = CTX) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 5, 3, 1, 4]
+  // row : [0, 2, 0, 1, 2, 0]
+  // col : [1, 2, 2, 0, 3, 2]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 3, 2}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 1, 2, 3, 4, 5}), sizeof(IDX) * 8, ctx));
+}
+
+template <typename IDX>
+aten::CSRMatrix SR_CSR3(DGLContext ctx) {
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 1, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::CSRMatrix SRC_CSR3(DGLContext ctx) {
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  return aten::CSRMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 3, 4, 6, 6}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({1, 2, 2, 0, 2, 3}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 0, 5, 3, 1, 4}), sizeof(IDX) * 8, ctx),
+      false);
+}
+
+template <typename IDX>
+aten::COOMatrix COO3(DGLContext ctx) {
+  // has duplicate entries
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // row : [0, 2, 0, 1, 2, 0]
+  // col : [2, 2, 1, 0, 3, 2]
+  return aten::COOMatrix(
+      4, 5,
+      aten::VecToIdArray(
+          std::vector<IDX>({0, 2, 0, 1, 2, 0}), sizeof(IDX) * 8, ctx),
+      aten::VecToIdArray(
+          std::vector<IDX>({2, 2, 1, 0, 3, 2}), sizeof(IDX) * 8, ctx));
+}
+
+}  // namespace
+
+template <typename IDX>
+void _TestCSRIsNonZero1(DGLContext ctx) {
+  auto csr = CSR1<IDX>(ctx);
+  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
+  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
+  IdArray r =
+      aten::VecToIdArray(std::vector<IDX>({2, 2, 0, 0}), sizeof(IDX) * 8, ctx);
+  IdArray c =
+      aten::VecToIdArray(std::vector<IDX>({1, 1, 1, 3}), sizeof(IDX) * 8, ctx);
+  IdArray x = aten::CSRIsNonZero(csr, r, c);
+  IdArray tx =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 1, 0}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+template <typename IDX>
+void _TestCSRIsNonZero2(DGLContext ctx) {
+  auto csr = CSR3<IDX>(ctx);
+  ASSERT_TRUE(aten::CSRIsNonZero(csr, 0, 1));
+  ASSERT_FALSE(aten::CSRIsNonZero(csr, 0, 0));
+  IdArray r = aten::VecToIdArray(
+      std::vector<IDX>({
+          0,
+          0,
+          0,
+          0,
+          0,
+      }),
+      sizeof(IDX) * 8, ctx);
+  IdArray c = aten::VecToIdArray(
+      std::vector<IDX>({
+          0,
+          1,
+          2,
+          3,
+          4,
+      }),
+      sizeof(IDX) * 8, ctx);
+  IdArray x = aten::CSRIsNonZero(csr, r, c);
+  IdArray tx = aten::VecToIdArray(
+      std::vector<IDX>({0, 1, 1, 1, 0}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx)) << " x = " << x << ", tx = " << tx;
+}
+
+TEST(SpmatTest, TestCSRIsNonZero) {
+  _TestCSRIsNonZero1<int32_t>(CPU);
+  _TestCSRIsNonZero1<int64_t>(CPU);
+  _TestCSRIsNonZero2<int32_t>(CPU);
+  _TestCSRIsNonZero2<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRIsNonZero1<int32_t>(GPU);
+  _TestCSRIsNonZero1<int64_t>(GPU);
+  _TestCSRIsNonZero2<int32_t>(GPU);
+  _TestCSRIsNonZero2<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRGetRowNNZ(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 0), 3);
+  ASSERT_EQ(aten::CSRGetRowNNZ(csr, 3), 0);
+  IdArray r =
+      aten::VecToIdArray(std::vector<IDX>({0, 3}), sizeof(IDX) * 8, ctx);
+  IdArray x = aten::CSRGetRowNNZ(csr, r);
+  IdArray tx =
+      aten::VecToIdArray(std::vector<IDX>({3, 0}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+TEST(SpmatTest, TestCSRGetRowNNZ) {
+  _TestCSRGetRowNNZ<int32_t>(CPU);
+  _TestCSRGetRowNNZ<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRGetRowNNZ<int32_t>(GPU);
+  _TestCSRGetRowNNZ<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRGetRowColumnIndices(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  auto x = aten::CSRGetRowColumnIndices(csr, 0);
+  auto tx =
+      aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::CSRGetRowColumnIndices(csr, 1);
+  tx = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::CSRGetRowColumnIndices(csr, 3);
+  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+TEST(SpmatTest, TestCSRGetRowColumnIndices) {
+  _TestCSRGetRowColumnIndices<int32_t>(CPU);
+  _TestCSRGetRowColumnIndices<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRGetRowColumnIndices<int32_t>(GPU);
+  _TestCSRGetRowColumnIndices<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRGetRowData(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  auto x = aten::CSRGetRowData(csr, 0);
+  auto tx =
+      aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::CSRGetRowData(csr, 1);
+  tx = aten::VecToIdArray(std::vector<IDX>({3}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::CSRGetRowData(csr, 3);
+  tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+TEST(SpmatTest, TestCSRGetRowData) {
+  _TestCSRGetRowData<int32_t>(CPU);
+  _TestCSRGetRowData<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRGetRowData<int32_t>(GPU);
+  _TestCSRGetRowData<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRGetData(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  // test get all data
+  auto x = aten::CSRGetAllData(csr, 0, 0);
+  auto tx = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+  x = aten::CSRGetAllData(csr, 0, 2);
+  tx = aten::VecToIdArray(std::vector<IDX>({2, 5}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data
+  auto r =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  auto c =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRGetData(csr, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data on sorted
+  csr = aten::CSRSort(csr);
+  r = aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRGetData(csr, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+
+  // test get data w/ broadcasting
+  r = aten::VecToIdArray(std::vector<IDX>({0}), sizeof(IDX) * 8, ctx);
+  c = aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRGetData(csr, r, c);
+  tx = aten::VecToIdArray(std::vector<IDX>({-1, 0, 2}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x, tx));
+}
+
+TEST(SpmatTest, CSRGetData) {
+  _TestCSRGetData<int32_t>(CPU);
+  _TestCSRGetData<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRGetData<int32_t>(GPU);
+  _TestCSRGetData<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRGetDataAndIndices(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  auto r =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  auto c =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+  auto x = aten::CSRGetDataAndIndices(csr, r, c);
+  auto tr =
+      aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+  auto tc =
+      aten::VecToIdArray(std::vector<IDX>({1, 2, 2}), sizeof(IDX) * 8, ctx);
+  auto td =
+      aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x[0], tr));
+  ASSERT_TRUE(ArrayEQ<IDX>(x[1], tc));
+  ASSERT_TRUE(ArrayEQ<IDX>(x[2], td));
+}
+
+TEST(SpmatTest, CSRGetDataAndIndices) {
+  _TestCSRGetDataAndIndices<int32_t>(CPU);
+  _TestCSRGetDataAndIndices<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRGetDataAndIndices<int32_t>(GPU);
+  _TestCSRGetDataAndIndices<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRTranspose(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  auto csr_t = aten::CSRTranspose(csr);
+  // [[0, 1, 0, 0],
+  //  [1, 0, 0, 0],
+  //  [2, 0, 1, 0],
+  //  [0, 0, 1, 0],
+  //  [0, 0, 0, 0]]
+  // data: [3, 0, 2, 5, 1, 4]
+  ASSERT_EQ(csr_t.num_rows, 5);
+  ASSERT_EQ(csr_t.num_cols, 4);
+  auto tp = aten::VecToIdArray(
+      std::vector<IDX>({0, 1, 2, 5, 6, 6}), sizeof(IDX) * 8, ctx);
+  auto ti = aten::VecToIdArray(
+      std::vector<IDX>({1, 0, 0, 0, 2, 2}), sizeof(IDX) * 8, ctx);
+  auto td = aten::VecToIdArray(
+      std::vector<IDX>({3, 0, 2, 5, 1, 4}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(csr_t.data, td));
+}
+
+TEST(SpmatTest, CSRTranspose) {
+  _TestCSRTranspose<int32_t>(CPU);
+  _TestCSRTranspose<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRTranspose<int32_t>(GPU);
+  _TestCSRTranspose<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRToCOO(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  {
+    auto coo = CSRToCOO(csr, false);
+    ASSERT_EQ(coo.num_rows, 4);
+    ASSERT_EQ(coo.num_cols, 5);
+    ASSERT_TRUE(coo.row_sorted);
+    auto tr = aten::VecToIdArray(
+        std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tr));
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, csr.indices));
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.data, csr.data));
+
+    // convert from sorted csr
+    auto s_csr = CSRSort(csr);
+    coo = CSRToCOO(s_csr, false);
+    ASSERT_EQ(coo.num_rows, 4);
+    ASSERT_EQ(coo.num_cols, 5);
+    ASSERT_TRUE(coo.row_sorted);
+    ASSERT_TRUE(coo.col_sorted);
+    tr = aten::VecToIdArray(
+        std::vector<IDX>({0, 0, 0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tr));
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, s_csr.indices));
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.data, s_csr.data));
+  }
+  {
+    auto coo = CSRToCOO(csr, true);
+    ASSERT_EQ(coo.num_rows, 4);
+    ASSERT_EQ(coo.num_cols, 5);
+    auto tcoo = COO2<IDX>(ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.row, tcoo.row));
+    ASSERT_TRUE(ArrayEQ<IDX>(coo.col, tcoo.col));
+  }
+}
+
+TEST(SpmatTest, CSRToCOO) {
+  _TestCSRToCOO<int32_t>(CPU);
+  _TestCSRToCOO<int64_t>(CPU);
+#if DGL_USE_CUDA
+  _TestCSRToCOO<int32_t>(GPU);
+  _TestCSRToCOO<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRSliceRows(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  auto x = aten::CSRSliceRows(csr, 1, 4);
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [3, 1, 4]
+  ASSERT_EQ(x.num_rows, 3);
+  ASSERT_EQ(x.num_cols, 5);
+  auto tp =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 3, 3}), sizeof(IDX) * 8, ctx);
+  auto ti =
+      aten::VecToIdArray(std::vector<IDX>({0, 2, 3}), sizeof(IDX) * 8, ctx);
+  auto td =
+      aten::VecToIdArray(std::vector<IDX>({3, 1, 4}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+
+  auto r =
+      aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRSliceRows(csr, r);
+  // [[0, 1, 2, 0, 0],
+  //  [1, 0, 0, 0, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: [0, 2, 5, 3]
+  tp = aten::VecToIdArray(std::vector<IDX>({0, 3, 4, 4}), sizeof(IDX) * 8, ctx);
+  ti = aten::VecToIdArray(std::vector<IDX>({1, 2, 2, 0}), sizeof(IDX) * 8, ctx);
+  td = aten::VecToIdArray(std::vector<IDX>({0, 2, 5, 3}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+
+  // Testing non-increasing row id based slicing
+  r = aten::VecToIdArray(std::vector<IDX>({3, 2, 1}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRSliceRows(csr, r);
+  // [[0, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [1, 0, 0, 0, 0]]
+  // data: [1, 4, 3]
+  tp = aten::VecToIdArray(std::vector<IDX>({0, 0, 2, 3}), sizeof(IDX) * 8, ctx);
+  ti = aten::VecToIdArray(std::vector<IDX>({2, 3, 0}), sizeof(IDX) * 8, ctx);
+  td = aten::VecToIdArray(std::vector<IDX>({1, 4, 3}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+
+  // Testing zero-degree row slicing with different rows
+  r = aten::VecToIdArray(
+      std::vector<IDX>({1, 3, 0, 3, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRSliceRows(csr, r);
+  // [[1, 0, 0, 0, 0],
+  //  [0, 0, 0, 0, 0],
+  //  [0, 1, 2, 0, 0],
+  //  [0, 0, 0, 0, 0],
+  //  [0, 0, 1, 1, 0]]
+  // data: [3, 0, 2, 5, 1, 4]
+  tp = aten::VecToIdArray(
+      std::vector<IDX>({0, 1, 1, 4, 4, 6}), sizeof(IDX) * 8, ctx);
+  ti = aten::VecToIdArray(
+      std::vector<IDX>({0, 1, 2, 2, 2, 3}), sizeof(IDX) * 8, ctx);
+  td = aten::VecToIdArray(
+      std::vector<IDX>({3, 0, 2, 5, 1, 4}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+
+  // Testing empty output (i.e. sliced rows will be zero-degree)
+  r = aten::VecToIdArray(std::vector<IDX>({3, 3, 3}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRSliceRows(csr, r);
+  // [[0, 0, 0, 0, 0],
+  //  [0, 0, 0, 0, 0],
+  //  [0, 0, 0, 0, 0]]
+  // data: []
+  tp = aten::VecToIdArray(std::vector<IDX>({0, 0, 0, 0}), sizeof(IDX) * 8, ctx);
+  ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+
+  // Testing constant output: we pick last row with at least one nnz
+  r = aten::VecToIdArray(std::vector<IDX>({2, 2, 2}), sizeof(IDX) * 8, ctx);
+  x = aten::CSRSliceRows(csr, r);
+  // [[0, 0, 1, 1, 0],
+  //  [0, 0, 1, 1, 0],
+  //  [0, 0, 1, 1, 0]]
+  // data: [1, 4, 1, 4, 1, 4]
+  tp = aten::VecToIdArray(std::vector<IDX>({0, 2, 4, 6}), sizeof(IDX) * 8, ctx);
+  ti = aten::VecToIdArray(
+      std::vector<IDX>({2, 3, 2, 3, 2, 3}), sizeof(IDX) * 8, ctx);
+  td = aten::VecToIdArray(
+      std::vector<IDX>({1, 4, 1, 4, 1, 4}), sizeof(IDX) * 8, ctx);
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+  ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+}
+
+TEST(SpmatTest, TestCSRSliceRows) {
+  _TestCSRSliceRows<int32_t>(CPU);
+  _TestCSRSliceRows<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRSliceRows<int32_t>(GPU);
+  _TestCSRSliceRows<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRSliceMatrix1(DGLContext ctx) {
+  auto csr = CSR2<IDX>(ctx);
+  {
+    // square
+    auto r =
+        aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
+    auto c =
+        aten::VecToIdArray(std::vector<IDX>({1, 2, 3}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[1, 2, 0],
+    //  [0, 0, 0],
+    //  [0, 0, 0]]
+    // data: [0, 2, 5]
+    ASSERT_EQ(x.num_rows, 3);
+    ASSERT_EQ(x.num_cols, 3);
+    auto tp = aten::VecToIdArray(
+        std::vector<IDX>({0, 3, 3, 3}), sizeof(IDX) * 8, ctx);
+    auto ti =
+        aten::VecToIdArray(std::vector<IDX>({0, 1, 1}), sizeof(IDX) * 8, ctx);
+    auto td =
+        aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+  {
+    // non-square
+    auto r =
+        aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[0, 1],
+    //  [1, 0],
+    //  [0, 0]]
+    // data: [0, 3]
+    ASSERT_EQ(x.num_rows, 3);
+    ASSERT_EQ(x.num_cols, 2);
+    auto tp = aten::VecToIdArray(
+        std::vector<IDX>({0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
+    auto ti =
+        aten::VecToIdArray(std::vector<IDX>({1, 0}), sizeof(IDX) * 8, ctx);
+    auto td =
+        aten::VecToIdArray(std::vector<IDX>({0, 3}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+  {
+    // empty slice
+    auto r = aten::VecToIdArray(std::vector<IDX>({2, 3}), sizeof(IDX) * 8, ctx);
+    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[0, 0],
+    //  [0, 0]]
+    // data: []
+    ASSERT_EQ(x.num_rows, 2);
+    ASSERT_EQ(x.num_cols, 2);
+    auto tp =
+        aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+    auto ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+    auto td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+}
+
+template <typename IDX>
+void _TestCSRSliceMatrix2(DGLContext ctx) {
+  auto csr = CSR3<IDX>(ctx);
+  {
+    // square
+    auto r =
+        aten::VecToIdArray(std::vector<IDX>({0, 1, 3}), sizeof(IDX) * 8, ctx);
+    auto c =
+        aten::VecToIdArray(std::vector<IDX>({1, 2, 3}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[1, 1, 1],
+    //  [0, 0, 0],
+    //  [0, 0, 0]]
+    // data: [5, 2, 0]
+    ASSERT_EQ(x.num_rows, 3);
+    ASSERT_EQ(x.num_cols, 3);
+    auto tp = aten::VecToIdArray(
+        std::vector<IDX>({0, 3, 3, 3}), sizeof(IDX) * 8, ctx);
+    // indexes are in reverse order in CSR3
+    auto ti =
+        aten::VecToIdArray(std::vector<IDX>({2, 1, 0}), sizeof(IDX) * 8, ctx);
+    auto td =
+        aten::VecToIdArray(std::vector<IDX>({0, 2, 5}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+  {
+    // non-square
+    auto r =
+        aten::VecToIdArray(std::vector<IDX>({0, 1, 2}), sizeof(IDX) * 8, ctx);
+    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[0, 1],
+    //  [1, 0],
+    //  [0, 0]]
+    // data: [0, 3]
+    ASSERT_EQ(x.num_rows, 3);
+    ASSERT_EQ(x.num_cols, 2);
+    auto tp = aten::VecToIdArray(
+        std::vector<IDX>({0, 1, 2, 2}), sizeof(IDX) * 8, ctx);
+    auto ti =
+        aten::VecToIdArray(std::vector<IDX>({1, 0}), sizeof(IDX) * 8, ctx);
+    auto td =
+        aten::VecToIdArray(std::vector<IDX>({5, 3}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+  {
+    // empty slice
+    auto r = aten::VecToIdArray(std::vector<IDX>({2, 3}), sizeof(IDX) * 8, ctx);
+    auto c = aten::VecToIdArray(std::vector<IDX>({0, 1}), sizeof(IDX) * 8, ctx);
+    auto x = aten::CSRSliceMatrix(csr, r, c);
+    // [[0, 0],
+    //  [0, 0]]
+    // data: []
+    ASSERT_EQ(x.num_rows, 2);
+    ASSERT_EQ(x.num_cols, 2);
+    auto tp =
+        aten::VecToIdArray(std::vector<IDX>({0, 0, 0}), sizeof(IDX) * 8, ctx);
+    auto ti = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+    auto td = aten::VecToIdArray(std::vector<IDX>({}), sizeof(IDX) * 8, ctx);
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indptr, tp));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.indices, ti));
+    ASSERT_TRUE(ArrayEQ<IDX>(x.data, td));
+  }
+}
+
+TEST(SpmatTest, CSRSliceMatrix) {
+  _TestCSRSliceMatrix1<int32_t>(CPU);
+  _TestCSRSliceMatrix1<int64_t>(CPU);
+  _TestCSRSliceMatrix2<int32_t>(CPU);
+  _TestCSRSliceMatrix2<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRSliceMatrix1<int32_t>(GPU);
+  _TestCSRSliceMatrix1<int64_t>(GPU);
+  _TestCSRSliceMatrix2<int32_t>(GPU);
+  _TestCSRSliceMatrix2<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRHasDuplicate(DGLContext ctx) {
+  auto csr = CSR1<IDX>(ctx);
+  ASSERT_FALSE(aten::CSRHasDuplicate(csr));
+  csr = CSR2<IDX>(ctx);
+  ASSERT_TRUE(aten::CSRHasDuplicate(csr));
+}
+
+TEST(SpmatTest, CSRHasDuplicate) {
+  _TestCSRHasDuplicate<int32_t>(CPU);
+  _TestCSRHasDuplicate<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRHasDuplicate<int32_t>(GPU);
+  _TestCSRHasDuplicate<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRSort(DGLContext ctx) {
+  auto csr = CSR1<IDX>(ctx);
+  ASSERT_FALSE(aten::CSRIsSorted(csr));
+  auto csr1 = aten::CSRSort(csr);
+  ASSERT_FALSE(aten::CSRIsSorted(csr));
+  ASSERT_TRUE(aten::CSRIsSorted(csr1));
+  ASSERT_TRUE(csr1.sorted);
+  aten::CSRSort_(&csr);
+  ASSERT_TRUE(aten::CSRIsSorted(csr));
+  ASSERT_TRUE(csr.sorted);
+  csr = CSR2<IDX>(ctx);
+  ASSERT_TRUE(aten::CSRIsSorted(csr));
+}
+
+TEST(SpmatTest, CSRSort) {
+  _TestCSRSort<int32_t>(CPU);
+  _TestCSRSort<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestCSRSort<int32_t>(GPU);
+  _TestCSRSort<int64_t>(GPU);
+#endif
+}
+
+template <typename IDX>
+void _TestCSRReorder() {
+  auto csr = CSR2<IDX>();
+  auto new_row =
+      aten::VecToIdArray(std::vector<IDX>({2, 0, 3, 1}), sizeof(IDX) * 8, CTX);
+  auto new_col = aten::VecToIdArray(
+      std::vector<IDX>({2, 0, 4, 3, 1}), sizeof(IDX) * 8, CTX);
+  auto new_csr = CSRReorder(csr, new_row, new_col);
+  ASSERT_EQ(new_csr.num_rows, csr.num_rows);
+  ASSERT_EQ(new_csr.num_cols, csr.num_cols);
+}
+
+TEST(SpmatTest, TestCSRReorder) {
+  _TestCSRReorder<int32_t>();
+  _TestCSRReorder<int64_t>();
+}

tests/cpp/test_spmm.cc

@@ -0,0 +1,244 @@
+#if !defined(_WIN32)
+#include <../src/array/cpu/spmm.h>
+#include <dgl/array.h>
+#include <gtest/gtest.h>
+#include <time.h>
+
+#include <random>
+
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+
+int sizes[] = {1, 7, 8, 9, 31, 32, 33, 54, 63, 64, 65, 256, 257};
+namespace ns_op = dgl::aten::cpu::op;
+namespace {
+
+template <class T>
+void GenerateData(T* data, int dim, T mul) {
+  for (int i = 0; i < dim; i++) {
+    data[i] = (i + 1) * mul;
+  }
+}
+
+template <class T>
+void GenerateRandomData(T* data, int dim) {
+  std::mt19937 rng(std::random_device{}());
+  std::uniform_int_distribution<> dist(0, 10000);
+  for (int i = 0; i < dim; i++) {
+    data[i] = (dist(rng) / 100);
+  }
+}
+
+template <class T>
+void GenerateZeroData(T* data, int dim) {
+  for (int i = 0; i < dim; i++) {
+    data[i] = 0;
+  }
+}
+
+template <class T>
+void Copy(T* exp, T* out, T* hs, int dim) {
+  for (int i = 0; i < dim; i++) {
+    exp[i] = out[i] + hs[i];
+  }
+}
+
+template <class T>
+void Add(T* exp, T* out, T* lhs, T* rhs, int dim) {
+  for (int i = 0; i < dim; i++) {
+    exp[i] = out[i] + lhs[i] + rhs[i];
+  }
+}
+
+template <class T>
+void Sub(T* exp, T* out, T* lhs, T* rhs, int dim) {
+  for (int i = 0; i < dim; i++) {
+    exp[i] = out[i] + lhs[i] - rhs[i];
+  }
+}
+
+template <class T>
+void Mul(T* exp, T* out, T* lhs, T* rhs, int dim) {
+  for (int i = 0; i < dim; i++) {
+    exp[i] = (out[i] + (lhs[i] * rhs[i]));
+  }
+}
+
+template <class T>
+void Div(T* exp, T* out, T* lhs, T* rhs, int dim) {
+  for (int i = 0; i < dim; i++) {
+    exp[i] = (out[i] + (lhs[i] / rhs[i]));
+  }
+}
+
+template <class T>
+void CheckResult(T* exp, T* out, int dim) {
+  for (int i = 0; i < dim; i++) {
+    ASSERT_TRUE(exp[i] == out[i]);
+  }
+}
+
+}  // namespace
+
+template <typename IDX>
+void _TestSpmmCopyLhs() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], lhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateRandomData(lhs, dim);
+
+    // Calculation of expected output - 'exp'
+    Copy(exp, out, lhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::CopyLhs<IDX>::Call(lhs + k, nullptr);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmCopyLhs) {
+  _TestSpmmCopyLhs<float>();
+  _TestSpmmCopyLhs<double>();
+  _TestSpmmCopyLhs<BFloat16>();
+}
+
+template <typename IDX>
+void _TestSpmmCopyRhs() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], rhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateRandomData(rhs, dim);
+
+    // Calculation of expected output - 'exp'
+    Copy(exp, out, rhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::CopyRhs<IDX>::Call(nullptr, rhs + k);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmCopyRhs) {
+  _TestSpmmCopyRhs<float>();
+  _TestSpmmCopyRhs<double>();
+  _TestSpmmCopyRhs<BFloat16>();
+}
+
+template <typename IDX>
+void _TestSpmmAdd() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], lhs[dim], rhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateRandomData(lhs, dim);
+    GenerateRandomData(rhs, dim);
+
+    // Calculation of expected output - 'exp'
+    Add(exp, out, lhs, rhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::Add<IDX>::Call(lhs + k, rhs + k);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmAdd) {
+  _TestSpmmAdd<float>();
+  _TestSpmmAdd<double>();
+  _TestSpmmAdd<BFloat16>();
+}
+
+template <typename IDX>
+void _TestSpmmSub() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], lhs[dim], rhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateRandomData(lhs, dim);
+    GenerateRandomData(rhs, dim);
+
+    // Calculation of expected output - 'exp'
+    Sub(exp, out, lhs, rhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::Sub<IDX>::Call(lhs + k, rhs + k);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmSub) {
+  _TestSpmmSub<float>();
+  _TestSpmmSub<double>();
+  _TestSpmmSub<BFloat16>();
+}
+
+template <typename IDX>
+void _TestSpmmMul() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], lhs[dim], rhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateRandomData(lhs, dim);
+    GenerateRandomData(rhs, dim);
+
+    // Calculation of expected output - 'exp'
+    Mul(exp, out, lhs, rhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::Mul<IDX>::Call(lhs + k, rhs + k);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmMul) {
+  _TestSpmmMul<float>();
+  _TestSpmmMul<double>();
+  _TestSpmmMul<BFloat16>();
+}
+
+template <typename IDX>
+void _TestSpmmDiv() {
+  for (size_t i = 0; i < sizeof(sizes) / sizeof(int); i++) {
+    int dim = sizes[i];
+    IDX out[dim], exp[dim], lhs[dim], rhs[dim];
+    GenerateZeroData(out, dim);
+    GenerateData(lhs, dim, (IDX)15);
+    GenerateData(rhs, dim, (IDX)1);
+
+    // Calculation of expected output - 'exp'
+    Div(exp, out, lhs, rhs, dim);
+
+    // Calculation of output using legacy path - 'out'
+    for (int k = 0; k < dim; k++) {
+      out[k] += ns_op::Div<IDX>::Call(lhs + k, rhs + k);
+    }
+
+    CheckResult(exp, out, dim);
+  }
+}
+
+TEST(SpmmTest, TestSpmmDiv) {
+  _TestSpmmDiv<float>();
+  _TestSpmmDiv<double>();
+  _TestSpmmDiv<BFloat16>();
+}
+#endif  // _WIN32

tests/cpp/test_unit_graph.cc

@@ -0,0 +1,434 @@
+/**
+ *  Copyright (c) 2019 by Contributors
+ * @file test_unit_graph.cc
+ * @brief Test UnitGraph
+ */
+#include <dgl/array.h>
+#include <dgl/immutable_graph.h>
+#include <dgl/runtime/device_api.h>
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <vector>
+
+#include "../../src/graph/unit_graph.h"
+#include "./../src/graph/heterograph.h"
+#include "./common.h"
+
+using namespace dgl;
+using namespace dgl::runtime;
+
+template <typename IdType>
+aten::CSRMatrix CSR1(DGLContext ctx) {
+  /**
+   * G = [[0, 0, 1],
+   *      [1, 0, 1],
+   *      [0, 1, 0],
+   *      [1, 0, 1]]
+   */
+  IdArray g_indptr = aten::VecToIdArray(
+      std::vector<IdType>({0, 1, 3, 4, 6}), sizeof(IdType) * 8, CTX);
+  IdArray g_indices = aten::VecToIdArray(
+      std::vector<IdType>({2, 0, 2, 1, 0, 2}), sizeof(IdType) * 8, CTX);
+
+  const aten::CSRMatrix &csr_a =
+      aten::CSRMatrix(4, 3, g_indptr, g_indices, aten::NullArray(), false);
+  return csr_a;
+}
+
+template aten::CSRMatrix CSR1<int32_t>(DGLContext ctx);
+template aten::CSRMatrix CSR1<int64_t>(DGLContext ctx);
+
+template <typename IdType>
+aten::COOMatrix COO1(DGLContext ctx) {
+  /**
+   * G = [[1, 1, 0],
+   *      [0, 1, 0]]
+   */
+  IdArray g_row = aten::VecToIdArray(
+      std::vector<IdType>({0, 0, 1}), sizeof(IdType) * 8, CTX);
+  IdArray g_col = aten::VecToIdArray(
+      std::vector<IdType>({0, 1, 1}), sizeof(IdType) * 8, CTX);
+  const aten::COOMatrix &coo =
+      aten::COOMatrix(2, 3, g_row, g_col, aten::NullArray(), true, true);
+
+  return coo;
+}
+
+template aten::COOMatrix COO1<int32_t>(DGLContext ctx);
+template aten::COOMatrix COO1<int64_t>(DGLContext ctx);
+
+template <typename IdType>
+void _TestUnitGraph_InOutDegrees(DGLContext ctx) {
+  /**
+  InDegree(s) is available only if COO or CSC formats permitted.
+  OutDegree(s) is available only if COO or CSR formats permitted.
+  */
+
+  // COO
+  {
+    const aten::COOMatrix &coo = COO1<IdType>(ctx);
+    auto &&g = CreateFromCOO(2, coo, COO_CODE);
+    ASSERT_EQ(g->InDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->InDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2}, g->NumBits(), g->Context())));
+    ASSERT_EQ(g->OutDegree(0, 0), 2);
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->OutDegrees(0, nids),
+        aten::VecToIdArray<IdType>({2, 1}, g->NumBits(), g->Context())));
+  }
+  // CSC
+  {
+    const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+    auto &&g = CreateFromCSC(2, csr, CSC_CODE);
+    ASSERT_EQ(g->InDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->InDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2, 1}, g->NumBits(), g->Context())));
+    EXPECT_ANY_THROW(g->OutDegree(0, 0));
+    EXPECT_ANY_THROW(g->OutDegrees(0, nids));
+  }
+  // CSR
+  {
+    const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+    auto &&g = CreateFromCSR(2, csr, CSR_CODE);
+    ASSERT_EQ(g->OutDegree(0, 0), 1);
+    auto &&nids = aten::Range(0, g->NumVertices(0), g->NumBits(), g->Context());
+    ASSERT_TRUE(ArrayEQ<IdType>(
+        g->OutDegrees(0, nids),
+        aten::VecToIdArray<IdType>({1, 2, 1, 2}, g->NumBits(), g->Context())));
+    EXPECT_ANY_THROW(g->InDegree(0, 0));
+    EXPECT_ANY_THROW(g->InDegrees(0, nids));
+  }
+}
+
+template <typename IdType>
+void _TestUnitGraph(DGLContext ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(ctx);
+
+  auto g = CreateFromCSC(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+
+  g = CreateFromCSR(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+
+  g = CreateFromCOO(2, coo);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+
+  auto src = aten::VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = aten::VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto mg = dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, COO_CODE);
+  ASSERT_EQ(mg->GetCreatedFormats(), 1);
+  auto hmg = dgl::UnitGraph::CreateFromCOO(1, 8, 8, src, dst, COO_CODE);
+  auto img = std::dynamic_pointer_cast<ImmutableGraph>(hmg->AsImmutableGraph());
+  ASSERT_TRUE(img != nullptr);
+  mg = dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, CSR_CODE | COO_CODE);
+  ASSERT_EQ(mg->GetCreatedFormats(), 1);
+  hmg = dgl::UnitGraph::CreateFromCOO(1, 8, 8, src, dst, CSR_CODE | COO_CODE);
+  img = std::dynamic_pointer_cast<ImmutableGraph>(hmg->AsImmutableGraph());
+  ASSERT_TRUE(img != nullptr);
+  mg = dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst, CSC_CODE | COO_CODE);
+  ASSERT_EQ(mg->GetCreatedFormats(), 1);
+  hmg = dgl::UnitGraph::CreateFromCOO(1, 8, 8, src, dst, CSC_CODE | COO_CODE);
+  img = std::dynamic_pointer_cast<ImmutableGraph>(hmg->AsImmutableGraph());
+  ASSERT_TRUE(img != nullptr);
+
+  g = CreateFromCSC(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+
+  g = CreateFromCSR(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+
+  g = CreateFromCOO(2, coo);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+}
+
+template <typename IdType>
+void _TestUnitGraph_GetInCSR(DGLContext ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(ctx);
+
+  auto g = CreateFromCSC(2, csr);
+  auto in_csr_matrix = g->GetCSCMatrix(0);
+  ASSERT_EQ(in_csr_matrix.num_rows, csr.num_rows);
+  ASSERT_EQ(in_csr_matrix.num_cols, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+
+  // test out csr
+  g = CreateFromCSR(2, csr);
+  auto g_ptr = g->GetGraphInFormat(CSC_CODE);
+  in_csr_matrix = g_ptr->GetCSCMatrix(0);
+  ASSERT_EQ(in_csr_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(in_csr_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+  in_csr_matrix = g->GetCSCMatrix(0);
+  ASSERT_EQ(in_csr_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(in_csr_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 6);
+
+  // test out coo
+  g = CreateFromCOO(2, coo);
+  g_ptr = g->GetGraphInFormat(CSC_CODE);
+  in_csr_matrix = g_ptr->GetCSCMatrix(0);
+  ASSERT_EQ(in_csr_matrix.num_cols, coo.num_rows);
+  ASSERT_EQ(in_csr_matrix.num_rows, coo.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+
+  in_csr_matrix = g->GetCSCMatrix(0);
+  ASSERT_EQ(in_csr_matrix.num_cols, coo.num_rows);
+  ASSERT_EQ(in_csr_matrix.num_rows, coo.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 5);
+}
+
+template <typename IdType>
+void _TestUnitGraph_GetOutCSR(DGLContext ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(ctx);
+
+  auto g = CreateFromCSC(2, csr);
+  auto g_ptr = g->GetGraphInFormat(CSR_CODE);
+  auto out_csr_matrix = g_ptr->GetCSRMatrix(0);
+  ASSERT_EQ(out_csr_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(out_csr_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+  out_csr_matrix = g->GetCSRMatrix(0);
+  ASSERT_EQ(out_csr_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(out_csr_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 6);
+
+  // test out csr
+  g = CreateFromCSR(2, csr);
+  out_csr_matrix = g->GetCSRMatrix(0);
+  ASSERT_EQ(out_csr_matrix.num_rows, csr.num_rows);
+  ASSERT_EQ(out_csr_matrix.num_cols, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+
+  // test out coo
+  g = CreateFromCOO(2, coo);
+  g_ptr = g->GetGraphInFormat(CSR_CODE);
+  out_csr_matrix = g_ptr->GetCSRMatrix(0);
+  ASSERT_EQ(out_csr_matrix.num_rows, coo.num_rows);
+  ASSERT_EQ(out_csr_matrix.num_cols, coo.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+
+  out_csr_matrix = g->GetCSRMatrix(0);
+  ASSERT_EQ(out_csr_matrix.num_rows, coo.num_rows);
+  ASSERT_EQ(out_csr_matrix.num_cols, coo.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 3);
+}
+
+template <typename IdType>
+void _TestUnitGraph_GetCOO(DGLContext ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(ctx);
+
+  auto g = CreateFromCSC(2, csr);
+  auto g_ptr = g->GetGraphInFormat(COO_CODE);
+  auto out_coo_matrix = g_ptr->GetCOOMatrix(0);
+  ASSERT_EQ(out_coo_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(out_coo_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+  out_coo_matrix = g->GetCOOMatrix(0);
+  ASSERT_EQ(out_coo_matrix.num_cols, csr.num_rows);
+  ASSERT_EQ(out_coo_matrix.num_rows, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 5);
+
+  // test out csr
+  g = CreateFromCSR(2, csr);
+  g_ptr = g->GetGraphInFormat(COO_CODE);
+  out_coo_matrix = g_ptr->GetCOOMatrix(0);
+  ASSERT_EQ(out_coo_matrix.num_rows, csr.num_rows);
+  ASSERT_EQ(out_coo_matrix.num_cols, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+  out_coo_matrix = g->GetCOOMatrix(0);
+  ASSERT_EQ(out_coo_matrix.num_rows, csr.num_rows);
+  ASSERT_EQ(out_coo_matrix.num_cols, csr.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 3);
+
+  // test out coo
+  g = CreateFromCOO(2, coo);
+  out_coo_matrix = g->GetCOOMatrix(0);
+  ASSERT_EQ(out_coo_matrix.num_rows, coo.num_rows);
+  ASSERT_EQ(out_coo_matrix.num_cols, coo.num_cols);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+}
+
+template <typename IdType>
+void _TestUnitGraph_Reserve(DGLContext ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(ctx);
+
+  auto g = CreateFromCSC(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+  auto r_g =
+      std::dynamic_pointer_cast<UnitGraph>(g->GetRelationGraph(0))->Reverse();
+  ASSERT_EQ(r_g->GetCreatedFormats(), 2);
+  aten::CSRMatrix g_in_csr = g->GetCSCMatrix(0);
+  aten::CSRMatrix r_g_out_csr = r_g->GetCSRMatrix(0);
+  ASSERT_TRUE(g_in_csr.indptr->data == r_g_out_csr.indptr->data);
+  ASSERT_TRUE(g_in_csr.indices->data == r_g_out_csr.indices->data);
+  aten::CSRMatrix g_out_csr = g->GetCSRMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 6);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 6);
+  aten::CSRMatrix r_g_in_csr = r_g->GetCSCMatrix(0);
+  ASSERT_TRUE(g_out_csr.indptr->data == r_g_in_csr.indptr->data);
+  ASSERT_TRUE(g_out_csr.indices->data == r_g_in_csr.indices->data);
+  aten::COOMatrix g_coo = g->GetCOOMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 7);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 6);
+  aten::COOMatrix r_g_coo = r_g->GetCOOMatrix(0);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 7);
+  ASSERT_EQ(g_coo.num_rows, r_g_coo.num_cols);
+  ASSERT_EQ(g_coo.num_cols, r_g_coo.num_rows);
+  ASSERT_TRUE(ArrayEQ<IdType>(g_coo.row, r_g_coo.col));
+  ASSERT_TRUE(ArrayEQ<IdType>(g_coo.col, r_g_coo.row));
+
+  // test out csr
+  g = CreateFromCSR(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+  r_g = std::dynamic_pointer_cast<UnitGraph>(g->GetRelationGraph(0))->Reverse();
+  ASSERT_EQ(r_g->GetCreatedFormats(), 4);
+  g_out_csr = g->GetCSRMatrix(0);
+  r_g_in_csr = r_g->GetCSCMatrix(0);
+  ASSERT_TRUE(g_out_csr.indptr->data == r_g_in_csr.indptr->data);
+  ASSERT_TRUE(g_out_csr.indices->data == r_g_in_csr.indices->data);
+  g_in_csr = g->GetCSCMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 6);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 6);
+  r_g_out_csr = r_g->GetCSRMatrix(0);
+  ASSERT_TRUE(g_in_csr.indptr->data == r_g_out_csr.indptr->data);
+  ASSERT_TRUE(g_in_csr.indices->data == r_g_out_csr.indices->data);
+  g_coo = g->GetCOOMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 7);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 6);
+  r_g_coo = r_g->GetCOOMatrix(0);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 7);
+  ASSERT_EQ(g_coo.num_rows, r_g_coo.num_cols);
+  ASSERT_EQ(g_coo.num_cols, r_g_coo.num_rows);
+  ASSERT_TRUE(ArrayEQ<IdType>(g_coo.row, r_g_coo.col));
+  ASSERT_TRUE(ArrayEQ<IdType>(g_coo.col, r_g_coo.row));
+
+  // test out coo
+  g = CreateFromCOO(2, coo);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+  r_g = std::dynamic_pointer_cast<UnitGraph>(g->GetRelationGraph(0))->Reverse();
+  ASSERT_EQ(r_g->GetCreatedFormats(), 1);
+  g_coo = g->GetCOOMatrix(0);
+  r_g_coo = r_g->GetCOOMatrix(0);
+  ASSERT_EQ(g_coo.num_rows, r_g_coo.num_cols);
+  ASSERT_EQ(g_coo.num_cols, r_g_coo.num_rows);
+  ASSERT_TRUE(g_coo.row->data == r_g_coo.col->data);
+  ASSERT_TRUE(g_coo.col->data == r_g_coo.row->data);
+  g_in_csr = g->GetCSCMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 5);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 3);
+  r_g_out_csr = r_g->GetCSRMatrix(0);
+  ASSERT_TRUE(g_in_csr.indptr->data == r_g_out_csr.indptr->data);
+  ASSERT_TRUE(g_in_csr.indices->data == r_g_out_csr.indices->data);
+  g_out_csr = g->GetCSRMatrix(0);
+  ASSERT_EQ(g->GetCreatedFormats(), 7);
+  ASSERT_EQ(r_g->GetCreatedFormats(), 7);
+  r_g_in_csr = r_g->GetCSCMatrix(0);
+  ASSERT_TRUE(g_out_csr.indptr->data == r_g_in_csr.indptr->data);
+  ASSERT_TRUE(g_out_csr.indices->data == r_g_in_csr.indices->data);
+}
+
+template <typename IdType>
+void _TestUnitGraph_CopyTo(
+    const DGLContext &src_ctx, const DGLContext &dst_ctx) {
+  const aten::CSRMatrix &csr = CSR1<IdType>(src_ctx);
+  const aten::COOMatrix &coo = COO1<IdType>(src_ctx);
+
+  auto device = dgl::runtime::DeviceAPI::Get(dst_ctx);
+  // We don't allow SetStream in DGL for now.
+  auto stream = nullptr;
+
+  auto g = dgl::UnitGraph::CreateFromCSC(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 4);
+  auto cg = dgl::UnitGraph::CopyTo(g, dst_ctx);
+  device->StreamSync(dst_ctx, stream);
+  ASSERT_EQ(cg->GetCreatedFormats(), 4);
+
+  g = dgl::UnitGraph::CreateFromCSR(2, csr);
+  ASSERT_EQ(g->GetCreatedFormats(), 2);
+  cg = dgl::UnitGraph::CopyTo(g, dst_ctx);
+  device->StreamSync(dst_ctx, stream);
+  ASSERT_EQ(cg->GetCreatedFormats(), 2);
+
+  g = dgl::UnitGraph::CreateFromCOO(2, coo);
+  ASSERT_EQ(g->GetCreatedFormats(), 1);
+  cg = dgl::UnitGraph::CopyTo(g, dst_ctx);
+  device->StreamSync(dst_ctx, stream);
+  ASSERT_EQ(cg->GetCreatedFormats(), 1);
+}
+
+TEST(UniGraphTest, TestUnitGraph_CopyTo) {
+  _TestUnitGraph_CopyTo<int32_t>(CPU, CPU);
+  _TestUnitGraph_CopyTo<int64_t>(CPU, CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_CopyTo<int32_t>(CPU, GPU);
+  _TestUnitGraph_CopyTo<int32_t>(GPU, GPU);
+  _TestUnitGraph_CopyTo<int32_t>(GPU, CPU);
+  _TestUnitGraph_CopyTo<int64_t>(CPU, GPU);
+  _TestUnitGraph_CopyTo<int64_t>(GPU, GPU);
+  _TestUnitGraph_CopyTo<int64_t>(GPU, CPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_InOutDegrees) {
+  _TestUnitGraph_InOutDegrees<int32_t>(CPU);
+  _TestUnitGraph_InOutDegrees<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_InOutDegrees<int32_t>(GPU);
+  _TestUnitGraph_InOutDegrees<int64_t>(GPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_Create) {
+  _TestUnitGraph<int32_t>(CPU);
+  _TestUnitGraph<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph<int32_t>(GPU);
+  _TestUnitGraph<int64_t>(GPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_GetInCSR) {
+  _TestUnitGraph_GetInCSR<int32_t>(CPU);
+  _TestUnitGraph_GetInCSR<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_GetInCSR<int32_t>(GPU);
+  _TestUnitGraph_GetInCSR<int64_t>(GPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_GetOutCSR) {
+  _TestUnitGraph_GetOutCSR<int32_t>(CPU);
+  _TestUnitGraph_GetOutCSR<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_GetOutCSR<int32_t>(GPU);
+  _TestUnitGraph_GetOutCSR<int64_t>(GPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_GetCOO) {
+  _TestUnitGraph_GetCOO<int32_t>(CPU);
+  _TestUnitGraph_GetCOO<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_GetCOO<int32_t>(GPU);
+  _TestUnitGraph_GetCOO<int64_t>(GPU);
+#endif
+}
+
+TEST(UniGraphTest, TestUnitGraph_Reserve) {
+  _TestUnitGraph_Reserve<int32_t>(CPU);
+  _TestUnitGraph_Reserve<int64_t>(CPU);
+#ifdef DGL_USE_CUDA
+  _TestUnitGraph_Reserve<int32_t>(GPU);
+  _TestUnitGraph_Reserve<int64_t>(GPU);
+#endif
+}

tests/cpp/test_zerocopy_serialize.cc

@@ -0,0 +1,167 @@
+#include <dgl/array.h>
+#include <dgl/immutable_graph.h>
+#include <dgl/zerocopy_serializer.h>
+#include <dmlc/memory_io.h>
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <iostream>
+#include <vector>
+
+#include "../../src/graph/heterograph.h"
+#include "../../src/graph/unit_graph.h"
+#include "./common.h"
+
+#ifndef _WIN32
+
+using namespace dgl;
+using namespace dgl::aten;
+using namespace dmlc;
+// Function to convert an idarray to string
+std::string IdArrayToStr(IdArray arr) {
+  arr = arr.CopyTo(DGLContext{kDGLCPU, 0});
+  int64_t len = arr->shape[0];
+  std::ostringstream oss;
+  oss << "(" << len << ")[";
+  if (arr->dtype.bits == 32) {
+    int32_t *data = static_cast<int32_t *>(arr->data);
+    for (int64_t i = 0; i < len; ++i) {
+      oss << data[i] << " ";
+    }
+  } else {
+    int64_t *data = static_cast<int64_t *>(arr->data);
+    for (int64_t i = 0; i < len; ++i) {
+      oss << data[i] << " ";
+    }
+  }
+  oss << "]";
+  return oss.str();
+}
+
+TEST(ZeroCopySerialize, NDArray) {
+  auto tensor1 = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto tensor2 = VecToIdArray<int64_t>({6, 6, 5, 7});
+
+  std::string nonzerocopy_blob;
+  dmlc::MemoryStringStream ifs(&nonzerocopy_blob);
+  static_cast<dmlc::Stream *>(&ifs)->Write(tensor1);
+  static_cast<dmlc::Stream *>(&ifs)->Write(tensor2);
+
+  std::string zerocopy_blob;
+  StreamWithBuffer zc_write_strm(&zerocopy_blob, true);
+  zc_write_strm.Write(tensor1);
+  zc_write_strm.Write(tensor2);
+
+  EXPECT_EQ(nonzerocopy_blob.size() - zerocopy_blob.size(), 126)
+      << "Invalid save";
+
+  std::vector<void *> new_ptr_list;
+  // Use memcpy to mimic remote machine reconstruction
+  for (auto ptr : zc_write_strm.buffer_list()) {
+    auto new_ptr = malloc(ptr.size);
+    memcpy(new_ptr, ptr.data, ptr.size);
+    new_ptr_list.emplace_back(new_ptr);
+  }
+
+  NDArray loadtensor1, loadtensor2;
+  StreamWithBuffer zc_read_strm(&zerocopy_blob, new_ptr_list);
+  zc_read_strm.Read(&loadtensor1);
+  zc_read_strm.Read(&loadtensor2);
+}
+
+TEST(ZeroCopySerialize, ZeroShapeNDArray) {
+  auto tensor1 = VecToIdArray<int64_t>({6, 6, 5, 7});
+  auto tensor2 = VecToIdArray<int64_t>({});
+  auto tensor3 = VecToIdArray<int64_t>({6, 6, 2, 7});
+  std::vector<NDArray> ndvec;
+  ndvec.push_back(tensor1);
+  ndvec.push_back(tensor2);
+  ndvec.push_back(tensor3);
+
+  std::string zerocopy_blob;
+  StreamWithBuffer zc_write_strm(&zerocopy_blob, true);
+  zc_write_strm.Write(ndvec);
+
+  std::vector<void *> new_ptr_list;
+  // Use memcpy to mimic remote machine reconstruction
+  for (auto ptr : zc_write_strm.buffer_list()) {
+    auto new_ptr = malloc(ptr.size);
+    memcpy(new_ptr, ptr.data, ptr.size);
+    new_ptr_list.emplace_back(new_ptr);
+  }
+
+  std::vector<NDArray> ndvec_read;
+  StreamWithBuffer zc_read_strm(&zerocopy_blob, new_ptr_list);
+  zc_read_strm.Read(&ndvec_read);
+  EXPECT_EQ(ndvec_read[1]->ndim, 1);
+  EXPECT_EQ(ndvec_read[1]->shape[0], 0);
+}
+
+TEST(ZeroCopySerialize, SharedMem) {
+  auto tensor1 = VecToIdArray<int64_t>({1, 2, 5, 3});
+  DGLDataType dtype = {kDGLInt, 64, 1};
+  std::vector<int64_t> shape{4};
+  DGLContext cpu_ctx = {kDGLCPU, 0};
+  auto shared_tensor =
+      NDArray::EmptyShared("test", shape, dtype, cpu_ctx, true);
+  shared_tensor.CopyFrom(tensor1);
+
+  std::string nonzerocopy_blob;
+  dmlc::MemoryStringStream ifs(&nonzerocopy_blob);
+  static_cast<dmlc::Stream *>(&ifs)->Write(shared_tensor);
+
+  std::string zerocopy_blob;
+  StreamWithBuffer zc_write_strm(&zerocopy_blob, false);
+  zc_write_strm.Write(shared_tensor);
+
+  EXPECT_EQ(nonzerocopy_blob.size() - zerocopy_blob.size(), 51)
+      << "Invalid save";
+  NDArray loadtensor1;
+
+  StreamWithBuffer zc_read_strm = StreamWithBuffer(&zerocopy_blob, false);
+  zc_read_strm.Read(&loadtensor1);
+}
+
+TEST(ZeroCopySerialize, HeteroGraph) {
+  auto src = VecToIdArray<int64_t>({1, 2, 5, 3});
+  auto dst = VecToIdArray<int64_t>({1, 6, 2, 6});
+  auto mg1 = dgl::UnitGraph::CreateFromCOO(2, 9, 8, src, dst);
+  src = VecToIdArray<int64_t>({6, 2, 5, 1, 8});
+  dst = VecToIdArray<int64_t>({5, 2, 4, 8, 0});
+  auto mg2 = dgl::UnitGraph::CreateFromCOO(1, 9, 9, src, dst);
+  std::vector<HeteroGraphPtr> relgraphs;
+  relgraphs.push_back(mg1);
+  relgraphs.push_back(mg2);
+  src = VecToIdArray<int64_t>({0, 0});
+  dst = VecToIdArray<int64_t>({1, 0});
+  auto meta_gptr = ImmutableGraph::CreateFromCOO(3, src, dst);
+  auto hrptr = std::make_shared<HeteroGraph>(meta_gptr, relgraphs);
+
+  std::string nonzerocopy_blob;
+  dmlc::MemoryStringStream ifs(&nonzerocopy_blob);
+  static_cast<dmlc::Stream *>(&ifs)->Write(hrptr);
+
+  std::string zerocopy_blob;
+  StreamWithBuffer zc_write_strm(&zerocopy_blob, true);
+  zc_write_strm.Write(hrptr);
+
+  EXPECT_EQ(nonzerocopy_blob.size() - zerocopy_blob.size(), 745)
+      << "Invalid save";
+
+  std::vector<void *> new_ptr_list;
+  // Use memcpy to mimic remote machine reconstruction
+  for (auto ptr : zc_write_strm.buffer_list()) {
+    auto new_ptr = malloc(ptr.size);
+    memcpy(new_ptr, ptr.data, ptr.size);
+    new_ptr_list.emplace_back(new_ptr);
+  }
+
+  auto gptr = dgl::Serializer::make_shared<HeteroGraph>();
+  StreamWithBuffer zc_read_strm(&zerocopy_blob, new_ptr_list);
+  zc_read_strm.Read(&gptr);
+
+  EXPECT_EQ(gptr->NumVertices(0), 9);
+  EXPECT_EQ(gptr->NumVertices(1), 8);
+}
+
+#endif  // _WIN32