chore: import upstream snapshot with attribution

tests/python/pytorch/cuda/test_nccl.py

@@ -0,0 +1,117 @@
+import unittest
+
+import backend as F
+import torch
+import torch.distributed as dist
+
+from dgl.cuda import nccl
+from dgl.partition import NDArrayPartition
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="NCCL only runs on GPU."
+)
+def test_nccl_sparse_push_single_remainder():
+    torch.cuda.set_device("cuda:0")
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:12345",
+        world_size=1,
+        rank=0,
+    )
+
+    index = F.randint([10000], F.int32, F.ctx(), 0, 10000)
+    value = F.uniform([10000, 100], F.float32, F.ctx(), -1.0, 1.0)
+
+    part = NDArrayPartition(10000, 1, "remainder")
+
+    ri, rv = nccl.sparse_all_to_all_push(index, value, part)
+    assert F.array_equal(ri, index)
+    assert F.array_equal(rv, value)
+
+    dist.destroy_process_group()
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="NCCL only runs on GPU."
+)
+def test_nccl_sparse_pull_single_remainder():
+    torch.cuda.set_device("cuda:0")
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:12345",
+        world_size=1,
+        rank=0,
+    )
+
+    req_index = F.randint([10000], F.int64, F.ctx(), 0, 100000)
+    value = F.uniform([100000, 100], F.float32, F.ctx(), -1.0, 1.0)
+
+    part = NDArrayPartition(100000, 1, "remainder")
+
+    rv = nccl.sparse_all_to_all_pull(req_index, value, part)
+    exp_rv = F.gather_row(value, req_index)
+    assert F.array_equal(rv, exp_rv)
+
+    dist.destroy_process_group()
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="NCCL only runs on GPU."
+)
+def test_nccl_sparse_push_single_range():
+    torch.cuda.set_device("cuda:0")
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:12345",
+        world_size=1,
+        rank=0,
+    )
+
+    index = F.randint([10000], F.int32, F.ctx(), 0, 10000)
+    value = F.uniform([10000, 100], F.float32, F.ctx(), -1.0, 1.0)
+
+    part_ranges = F.copy_to(
+        F.tensor([0, value.shape[0]], dtype=F.int64), F.ctx()
+    )
+    part = NDArrayPartition(10000, 1, "range", part_ranges=part_ranges)
+
+    ri, rv = nccl.sparse_all_to_all_push(index, value, part)
+    assert F.array_equal(ri, index)
+    assert F.array_equal(rv, value)
+
+    dist.destroy_process_group()
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="NCCL only runs on GPU."
+)
+def test_nccl_sparse_pull_single_range():
+    torch.cuda.set_device("cuda:0")
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:12345",
+        world_size=1,
+        rank=0,
+    )
+
+    req_index = F.randint([10000], F.int64, F.ctx(), 0, 100000)
+    value = F.uniform([100000, 100], F.float32, F.ctx(), -1.0, 1.0)
+
+    part_ranges = F.copy_to(
+        F.tensor([0, value.shape[0]], dtype=F.int64), F.ctx()
+    )
+    part = NDArrayPartition(100000, 1, "range", part_ranges=part_ranges)
+
+    rv = nccl.sparse_all_to_all_pull(req_index, value, part)
+    exp_rv = F.gather_row(value, req_index)
+    assert F.array_equal(rv, exp_rv)
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    test_nccl_sparse_push_single_remainder()
+    test_nccl_sparse_pull_single_remainder()
+    test_nccl_sparse_push_single_range()
+    test_nccl_sparse_pull_single_range()

tests/python/pytorch/dataloading/test_dataloader.py

@@ -0,0 +1,854 @@
+import os
+import unittest
+from collections.abc import Iterator, Mapping
+from functools import partial
+
+import backend as F
+
+import dgl
+import dgl.ops as OPS
+import numpy as np
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from utils import parametrize_idtype
+
+
+@pytest.mark.parametrize("batch_size", [None, 16])
+def test_graph_dataloader(batch_size):
+    num_batches = 2
+    num_samples = num_batches * (batch_size if batch_size is not None else 1)
+    minigc_dataset = dgl.data.MiniGCDataset(num_samples, 10, 20)
+    data_loader = dgl.dataloading.GraphDataLoader(
+        minigc_dataset, batch_size=batch_size, shuffle=True
+    )
+    assert isinstance(iter(data_loader), Iterator)
+    for graph, label in data_loader:
+        assert isinstance(graph, dgl.DGLGraph)
+        if batch_size is not None:
+            assert F.asnumpy(label).shape[0] == batch_size
+        else:
+            # If batch size is None, the label element will be a single scalar following
+            # PyTorch's practice.
+            assert F.asnumpy(label).ndim == 0
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("num_workers", [0, 4])
+def test_cluster_gcn(num_workers):
+    dataset = dgl.data.CoraFullDataset()
+    g = dataset[0]
+    sampler = dgl.dataloading.ClusterGCNSampler(g, 100)
+    dataloader = dgl.dataloading.DataLoader(
+        g, torch.arange(100), sampler, batch_size=4, num_workers=num_workers
+    )
+    assert len(dataloader) == 25
+    for i, sg in enumerate(dataloader):
+        pass
+
+
+@pytest.mark.parametrize("num_workers", [0, 4])
+def test_shadow(num_workers):
+    g = dgl.data.CoraFullDataset()[0]
+    sampler = dgl.dataloading.ShaDowKHopSampler([5, 10, 15])
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        torch.arange(g.num_nodes()),
+        sampler,
+        batch_size=5,
+        shuffle=True,
+        drop_last=False,
+        num_workers=num_workers,
+    )
+    for i, (input_nodes, output_nodes, subgraph) in enumerate(dataloader):
+        assert torch.equal(input_nodes, subgraph.ndata[dgl.NID])
+        assert torch.equal(input_nodes[: output_nodes.shape[0]], output_nodes)
+        assert torch.equal(
+            subgraph.ndata["label"], g.ndata["label"][input_nodes]
+        )
+        assert torch.equal(subgraph.ndata["feat"], g.ndata["feat"][input_nodes])
+        if i == 5:
+            break
+
+
+@pytest.mark.parametrize("num_workers", [0, 4])
+@pytest.mark.parametrize("mode", ["node", "edge", "walk"])
+def test_saint(num_workers, mode):
+    g = dgl.data.CoraFullDataset()[0]
+
+    if mode == "node":
+        budget = 100
+    elif mode == "edge":
+        budget = 200
+    elif mode == "walk":
+        budget = (3, 2)
+
+    sampler = dgl.dataloading.SAINTSampler(mode, budget)
+    dataloader = dgl.dataloading.DataLoader(
+        g, torch.arange(100), sampler, num_workers=num_workers
+    )
+    assert len(dataloader) == 100
+    for sg in dataloader:
+        pass
+
+
+@parametrize_idtype
+@pytest.mark.parametrize(
+    "mode", ["cpu", "uva_cuda_indices", "uva_cpu_indices", "pure_gpu"]
+)
+@pytest.mark.parametrize("use_ddp", [False, True])
+@pytest.mark.parametrize("use_mask", [False, True])
+def test_neighbor_nonuniform(idtype, mode, use_ddp, use_mask):
+    if mode != "cpu" and F.ctx() == F.cpu():
+        pytest.skip("UVA and GPU sampling require a GPU.")
+    if mode != "cpu" and use_mask:
+        pytest.skip("Masked sampling only works on CPU.")
+    if use_ddp:
+        if os.name == "nt":
+            pytest.skip("PyTorch 1.13.0+ has problems in Windows DDP...")
+        dist.init_process_group(
+            "gloo" if F.ctx() == F.cpu() else "nccl",
+            "tcp://127.0.0.1:12347",
+            world_size=1,
+            rank=0,
+        )
+    g = dgl.graph(([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1])).astype(
+        idtype
+    )
+    g.edata["p"] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
+    g.edata["mask"] = g.edata["p"] != 0
+    if mode in ("cpu", "uva_cpu_indices"):
+        indices = F.copy_to(F.tensor([0, 1], idtype), F.cpu())
+    else:
+        indices = F.copy_to(F.tensor([0, 1], idtype), F.cuda())
+    if mode == "pure_gpu":
+        g = g.to(F.cuda())
+    use_uva = mode.startswith("uva")
+
+    if use_mask:
+        prob, mask = None, "mask"
+    else:
+        prob, mask = "p", None
+
+    sampler = dgl.dataloading.MultiLayerNeighborSampler(
+        [2], prob=prob, mask=mask
+    )
+    for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g,
+            indices,
+            sampler,
+            batch_size=1,
+            device=F.ctx(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp,
+        )
+        for input_nodes, output_nodes, blocks in dataloader:
+            seed = output_nodes.item()
+            neighbors = set(input_nodes[1:].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {5, 6}
+            elif seed == 0:
+                assert neighbors == {1, 2}
+
+    g = dgl.heterograph(
+        {
+            ("B", "BA", "A"): (
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                [0, 0, 0, 0, 1, 1, 1, 1],
+            ),
+            ("C", "CA", "A"): (
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                [0, 0, 0, 0, 1, 1, 1, 1],
+            ),
+        }
+    ).astype(idtype)
+    g.edges["BA"].data["p"] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
+    g.edges["BA"].data["mask"] = g.edges["BA"].data["p"] != 0
+    g.edges["CA"].data["p"] = torch.FloatTensor([0, 0, 1, 1, 0, 0, 1, 1])
+    g.edges["CA"].data["mask"] = g.edges["CA"].data["p"] != 0
+    if mode == "pure_gpu":
+        g = g.to(F.cuda())
+    for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g,
+            {"A": indices},
+            sampler,
+            batch_size=1,
+            device=F.ctx(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp,
+        )
+        for input_nodes, output_nodes, blocks in dataloader:
+            seed = output_nodes["A"].item()
+            # Seed and neighbors are of different node types so slicing is not necessary here.
+            neighbors = set(input_nodes["B"].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {5, 6}
+            elif seed == 0:
+                assert neighbors == {1, 2}
+
+            neighbors = set(input_nodes["C"].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {7, 8}
+            elif seed == 0:
+                assert neighbors == {3, 4}
+
+    if use_ddp:
+        dist.destroy_process_group()
+
+
+def _check_dtype(data, dtype, attr_name):
+    if isinstance(data, dict):
+        for k, v in data.items():
+            assert getattr(v, attr_name) == dtype
+    elif isinstance(data, list):
+        for v in data:
+            assert getattr(v, attr_name) == dtype
+    else:
+        assert getattr(data, attr_name) == dtype
+
+
+def _check_device(data):
+    if isinstance(data, dict):
+        for k, v in data.items():
+            assert v.device == F.ctx()
+    elif isinstance(data, list):
+        for v in data:
+            assert v.device == F.ctx()
+    else:
+        assert data.device == F.ctx()
+
+
+@pytest.mark.parametrize("sampler_name", ["full", "neighbor"])
+@pytest.mark.parametrize(
+    "mode", ["cpu", "uva_cuda_indices", "uva_cpu_indices", "pure_gpu"]
+)
+@pytest.mark.parametrize("nprocs", [1, 4])
+@pytest.mark.parametrize("drop_last", [True, False])
+def test_ddp_dataloader_decompose_dataset(
+    sampler_name, mode, nprocs, drop_last
+):
+    if torch.cuda.device_count() < nprocs and mode != "cpu":
+        pytest.skip(
+            "DDP dataloader needs sufficient GPUs for UVA and GPU sampling."
+        )
+    if mode != "cpu" and F.ctx() == F.cpu():
+        pytest.skip("UVA and GPU sampling require a GPU.")
+
+    if os.name == "nt":
+        pytest.skip("PyTorch 1.13.0+ has problems in Windows DDP...")
+    g, _, _, _ = _create_homogeneous()
+    g = g.to(F.cpu())
+
+    sampler = {
+        "full": dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        "neighbor": dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+    }[sampler_name]
+    indices = F.copy_to(F.arange(0, g.num_nodes()), F.cpu())
+    data = indices, sampler
+    arguments = mode, drop_last
+    g.create_formats_()
+    os.environ["OMP_NUM_THREADS"] = str(mp.cpu_count() // 2 // nprocs)
+    mp.spawn(_ddp_runner, args=(nprocs, g, data, arguments), nprocs=nprocs)
+
+
+def _ddp_runner(proc_id, nprocs, g, data, args):
+    mode, drop_last = args
+    indices, sampler = data
+    if mode == "cpu":
+        device = torch.device("cpu")
+    else:
+        device = torch.device(proc_id)
+        torch.cuda.set_device(device)
+    if mode == "pure_gpu":
+        g = g.to(F.cuda())
+    if mode in ("cpu", "uva_cpu_indices"):
+        indices = indices.cpu()
+    else:
+        indices = indices.cuda()
+
+    dist.init_process_group(
+        "nccl" if mode != "cpu" else "gloo",
+        "tcp://127.0.0.1:12347",
+        world_size=nprocs,
+        rank=proc_id,
+    )
+    use_uva = mode.startswith("uva")
+    batch_size = g.num_nodes()
+    shuffle = False
+    for num_workers in [1, 4] if mode == "cpu" else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g,
+            indices,
+            sampler,
+            device=device,
+            batch_size=batch_size,  # g1.num_nodes(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=True,
+            drop_last=drop_last,
+            shuffle=shuffle,
+        )
+        max_nid = [0]
+        for i, (input_nodes, output_nodes, blocks) in enumerate(dataloader):
+            block = blocks[-1]
+            o_src, o_dst = block.edges()
+            src_nodes_id = block.srcdata[dgl.NID][o_src]
+            dst_nodes_id = block.dstdata[dgl.NID][o_dst]
+            max_nid.append(np.max(dst_nodes_id.cpu().numpy()))
+
+        local_max = torch.tensor(np.max(max_nid))
+        if torch.distributed.get_backend() == "nccl":
+            local_max = local_max.cuda()
+        dist.reduce(local_max, 0, op=dist.ReduceOp.MAX)
+        if proc_id == 0:
+            if drop_last and not shuffle and local_max > 0:
+                assert (
+                    local_max.item()
+                    == len(indices)
+                    - len(indices) % nprocs
+                    - 1
+                    - (len(indices) // nprocs) % batch_size
+                )
+            elif not drop_last:
+                assert local_max == len(indices) - 1
+    dist.destroy_process_group()
+
+
+@parametrize_idtype
+@pytest.mark.parametrize(
+    "sampler_name", ["full", "neighbor", "neighbor2", "labor"]
+)
+@pytest.mark.parametrize(
+    "mode", ["cpu", "uva_cuda_indices", "uva_cpu_indices", "pure_gpu"]
+)
+@pytest.mark.parametrize("use_ddp", [False, True])
+def test_node_dataloader(idtype, sampler_name, mode, use_ddp):
+    if mode != "cpu" and F.ctx() == F.cpu():
+        pytest.skip("UVA and GPU sampling require a GPU.")
+    if use_ddp:
+        if os.name == "nt":
+            pytest.skip("PyTorch 1.13.0+ has problems in Windows DDP...")
+        dist.init_process_group(
+            "gloo" if F.ctx() == F.cpu() else "nccl",
+            "tcp://127.0.0.1:12347",
+            world_size=1,
+            rank=0,
+        )
+    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
+    g1.ndata["feat"] = F.copy_to(F.randn((5, 8)), F.cpu())
+    g1.ndata["label"] = F.copy_to(F.randn((g1.num_nodes(),)), F.cpu())
+    if mode in ("cpu", "uva_cpu_indices"):
+        indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cpu())
+    else:
+        indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cuda())
+    if mode == "pure_gpu":
+        g1 = g1.to(F.cuda())
+
+    use_uva = mode.startswith("uva")
+
+    sampler = {
+        "full": dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        "neighbor": dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+        "neighbor2": dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+        "labor": dgl.dataloading.LaborSampler([3, 3]),
+    }[sampler_name]
+    for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g1,
+            indices,
+            sampler,
+            device=F.ctx(),
+            batch_size=g1.num_nodes(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp,
+        )
+        for input_nodes, output_nodes, blocks in dataloader:
+            _check_device(input_nodes)
+            _check_device(output_nodes)
+            _check_device(blocks)
+            _check_dtype(input_nodes, idtype, "dtype")
+            _check_dtype(output_nodes, idtype, "dtype")
+            _check_dtype(blocks, idtype, "idtype")
+
+    g2 = dgl.heterograph(
+        {
+            ("user", "follow", "user"): (
+                [0, 0, 0, 1, 1, 1, 2],
+                [1, 2, 3, 0, 2, 3, 0],
+            ),
+            ("user", "followed-by", "user"): (
+                [1, 2, 3, 0, 2, 3, 0],
+                [0, 0, 0, 1, 1, 1, 2],
+            ),
+            ("user", "play", "game"): ([0, 1, 1, 3, 5], [0, 1, 2, 0, 2]),
+            ("game", "played-by", "user"): ([0, 1, 2, 0, 2], [0, 1, 1, 3, 5]),
+        }
+    ).astype(idtype)
+    for ntype in g2.ntypes:
+        g2.nodes[ntype].data["feat"] = F.copy_to(
+            F.randn((g2.num_nodes(ntype), 8)), F.cpu()
+        )
+    if mode in ("cpu", "uva_cpu_indices"):
+        indices = {nty: F.copy_to(g2.nodes(nty), F.cpu()) for nty in g2.ntypes}
+    else:
+        indices = {nty: F.copy_to(g2.nodes(nty), F.cuda()) for nty in g2.ntypes}
+    if mode == "pure_gpu":
+        g2 = g2.to(F.cuda())
+
+    batch_size = max(g2.num_nodes(nty) for nty in g2.ntypes)
+    sampler = {
+        "full": dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        "neighbor": dgl.dataloading.MultiLayerNeighborSampler(
+            [{etype: 3 for etype in g2.etypes}] * 2
+        ),
+        "neighbor2": dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+        "labor": dgl.dataloading.LaborSampler([3, 3]),
+    }[sampler_name]
+    for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g2,
+            indices,
+            sampler,
+            device=F.ctx(),
+            batch_size=batch_size,
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp,
+        )
+        assert isinstance(iter(dataloader), Iterator)
+        for input_nodes, output_nodes, blocks in dataloader:
+            _check_device(input_nodes)
+            _check_device(output_nodes)
+            _check_device(blocks)
+            _check_dtype(input_nodes, idtype, "dtype")
+            _check_dtype(output_nodes, idtype, "dtype")
+            _check_dtype(blocks, idtype, "idtype")
+
+    if use_ddp:
+        dist.destroy_process_group()
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("sampler_name", ["full", "neighbor"])
+@pytest.mark.parametrize(
+    "neg_sampler",
+    [
+        dgl.dataloading.negative_sampler.Uniform(2),
+        dgl.dataloading.negative_sampler.GlobalUniform(15, False, 3),
+        dgl.dataloading.negative_sampler.GlobalUniform(15, True, 3),
+    ],
+)
+@pytest.mark.parametrize("mode", ["cpu", "uva", "pure_gpu"])
+@pytest.mark.parametrize("use_ddp", [False, True])
+def test_edge_dataloader(idtype, sampler_name, neg_sampler, mode, use_ddp):
+    if mode != "cpu" and F.ctx() == F.cpu():
+        pytest.skip("UVA and GPU sampling require a GPU.")
+    if mode == "uva" and isinstance(
+        neg_sampler, dgl.dataloading.negative_sampler.GlobalUniform
+    ):
+        pytest.skip("GlobalUniform don't support UVA yet.")
+    if use_ddp:
+        if os.name == "nt":
+            pytest.skip("PyTorch 1.13.0+ has problems in Windows DDP...")
+        dist.init_process_group(
+            "gloo" if F.ctx() == F.cpu() else "nccl",
+            "tcp://127.0.0.1:12347",
+            world_size=1,
+            rank=0,
+        )
+    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
+    g1.ndata["feat"] = F.copy_to(F.randn((5, 8)), F.cpu())
+    if mode == "pure_gpu":
+        g1 = g1.to(F.cuda())
+
+    sampler = {
+        "full": dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        "neighbor": dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+    }[sampler_name]
+
+    # no negative sampler
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    dataloader = dgl.dataloading.DataLoader(
+        g1,
+        g1.edges(form="eid"),
+        edge_sampler,
+        device=F.ctx(),
+        batch_size=g1.num_edges(),
+        use_uva=(mode == "uva"),
+        use_ddp=use_ddp,
+    )
+    for input_nodes, pos_pair_graph, blocks in dataloader:
+        _check_device(input_nodes)
+        _check_device(pos_pair_graph)
+        _check_device(blocks)
+
+    # negative sampler
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler, negative_sampler=neg_sampler
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g1,
+        g1.edges(form="eid"),
+        edge_sampler,
+        device=F.ctx(),
+        batch_size=g1.num_edges(),
+        use_uva=(mode == "uva"),
+        use_ddp=use_ddp,
+    )
+    for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+        _check_device(input_nodes)
+        _check_device(pos_pair_graph)
+        _check_device(neg_pair_graph)
+        _check_device(blocks)
+
+    g2 = dgl.heterograph(
+        {
+            ("user", "follow", "user"): (
+                [0, 0, 0, 1, 1, 1, 2],
+                [1, 2, 3, 0, 2, 3, 0],
+            ),
+            ("user", "followed-by", "user"): (
+                [1, 2, 3, 0, 2, 3, 0],
+                [0, 0, 0, 1, 1, 1, 2],
+            ),
+            ("user", "play", "game"): ([0, 1, 1, 3, 5], [0, 1, 2, 0, 2]),
+            ("game", "played-by", "user"): ([0, 1, 2, 0, 2], [0, 1, 1, 3, 5]),
+        }
+    ).astype(idtype)
+    for ntype in g2.ntypes:
+        g2.nodes[ntype].data["feat"] = F.copy_to(
+            F.randn((g2.num_nodes(ntype), 8)), F.cpu()
+        )
+    if mode == "pure_gpu":
+        g2 = g2.to(F.cuda())
+
+    batch_size = max(g2.num_edges(ety) for ety in g2.canonical_etypes)
+    sampler = {
+        "full": dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        "neighbor": dgl.dataloading.MultiLayerNeighborSampler(
+            [{etype: 3 for etype in g2.etypes}] * 2
+        ),
+    }[sampler_name]
+
+    # no negative sampler
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    dataloader = dgl.dataloading.DataLoader(
+        g2,
+        {ety: g2.edges(form="eid", etype=ety) for ety in g2.canonical_etypes},
+        edge_sampler,
+        device=F.ctx(),
+        batch_size=batch_size,
+        use_uva=(mode == "uva"),
+        use_ddp=use_ddp,
+    )
+    for input_nodes, pos_pair_graph, blocks in dataloader:
+        _check_device(input_nodes)
+        _check_device(pos_pair_graph)
+        _check_device(blocks)
+
+    # negative sampler
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler, negative_sampler=neg_sampler
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g2,
+        {ety: g2.edges(form="eid", etype=ety) for ety in g2.canonical_etypes},
+        edge_sampler,
+        device=F.ctx(),
+        batch_size=batch_size,
+        use_uva=(mode == "uva"),
+        use_ddp=use_ddp,
+    )
+
+    assert isinstance(iter(dataloader), Iterator)
+    for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
+        _check_device(input_nodes)
+        _check_device(pos_pair_graph)
+        _check_device(neg_pair_graph)
+        _check_device(blocks)
+
+    if use_ddp:
+        dist.destroy_process_group()
+
+
+def _create_homogeneous():
+    s = torch.randint(0, 200, (1000,), device=F.ctx())
+    d = torch.randint(0, 200, (1000,), device=F.ctx())
+    src = torch.cat([s, d])
+    dst = torch.cat([d, s])
+    g = dgl.graph((s, d), num_nodes=200)
+    reverse_eids = torch.cat(
+        [torch.arange(1000, 2000), torch.arange(0, 1000)]
+    ).to(F.ctx())
+    always_exclude = torch.randint(0, 1000, (50,), device=F.ctx())
+    seed_edges = torch.arange(0, 1000, device=F.ctx())
+    return g, reverse_eids, always_exclude, seed_edges
+
+
+def _create_heterogeneous():
+    edges = {}
+    for utype, etype, vtype in [("A", "AA", "A"), ("A", "AB", "B")]:
+        s = torch.randint(0, 200, (1000,), device=F.ctx())
+        d = torch.randint(0, 200, (1000,), device=F.ctx())
+        edges[utype, etype, vtype] = (s, d)
+        edges[vtype, "rev-" + etype, utype] = (d, s)
+    g = dgl.heterograph(edges, num_nodes_dict={"A": 200, "B": 200})
+    reverse_etypes = {
+        "AA": "rev-AA",
+        "AB": "rev-AB",
+        "rev-AA": "AA",
+        "rev-AB": "AB",
+    }
+    always_exclude = {
+        "AA": torch.randint(0, 1000, (50,), device=F.ctx()),
+        "AB": torch.randint(0, 1000, (50,), device=F.ctx()),
+    }
+    seed_edges = {
+        "AA": torch.arange(0, 1000, device=F.ctx()),
+        "AB": torch.arange(0, 1000, device=F.ctx()),
+    }
+    return g, reverse_etypes, always_exclude, seed_edges
+
+
+def _remove_duplicates(s, d):
+    s, d = list(zip(*list(set(zip(s.tolist(), d.tolist())))))
+    return torch.tensor(s, device=F.ctx()), torch.tensor(d, device=F.ctx())
+
+
+def _find_edges_to_exclude(g, exclude, always_exclude, pair_eids):
+    if exclude == None:
+        return always_exclude
+    elif exclude == "self":
+        return (
+            torch.cat([pair_eids, always_exclude])
+            if always_exclude is not None
+            else pair_eids
+        )
+    elif exclude == "reverse_id":
+        pair_eids = torch.cat([pair_eids, pair_eids + 1000])
+        return (
+            torch.cat([pair_eids, always_exclude])
+            if always_exclude is not None
+            else pair_eids
+        )
+    elif exclude == "reverse_types":
+        pair_eids = {g.to_canonical_etype(k): v for k, v in pair_eids.items()}
+        if ("A", "AA", "A") in pair_eids:
+            pair_eids[("A", "rev-AA", "A")] = pair_eids[("A", "AA", "A")]
+        if ("A", "AB", "B") in pair_eids:
+            pair_eids[("B", "rev-AB", "A")] = pair_eids[("A", "AB", "B")]
+        if always_exclude is not None:
+            always_exclude = {
+                g.to_canonical_etype(k): v for k, v in always_exclude.items()
+            }
+            for k in always_exclude.keys():
+                if k in pair_eids:
+                    pair_eids[k] = torch.cat([pair_eids[k], always_exclude[k]])
+                else:
+                    pair_eids[k] = always_exclude[k]
+        return pair_eids
+
+
+@pytest.mark.parametrize("always_exclude_flag", [False, True])
+@pytest.mark.parametrize(
+    "exclude", [None, "self", "reverse_id", "reverse_types"]
+)
+@pytest.mark.parametrize(
+    "sampler",
+    [
+        dgl.dataloading.MultiLayerFullNeighborSampler(1),
+        dgl.dataloading.ShaDowKHopSampler([5]),
+    ],
+)
+@pytest.mark.parametrize("batch_size", [1, 50])
+def test_edge_dataloader_excludes(
+    exclude, always_exclude_flag, batch_size, sampler
+):
+    if exclude == "reverse_types":
+        g, reverse_etypes, always_exclude, seed_edges = _create_heterogeneous()
+    else:
+        g, reverse_eids, always_exclude, seed_edges = _create_homogeneous()
+    g = g.to(F.ctx())
+    if not always_exclude_flag:
+        always_exclude = None
+
+    kwargs = {}
+    kwargs["exclude"] = (
+        partial(_find_edges_to_exclude, g, exclude, always_exclude)
+        if always_exclude_flag
+        else exclude
+    )
+    kwargs["reverse_eids"] = reverse_eids if exclude == "reverse_id" else None
+    kwargs["reverse_etypes"] = (
+        reverse_etypes if exclude == "reverse_types" else None
+    )
+    sampler = dgl.dataloading.as_edge_prediction_sampler(sampler, **kwargs)
+
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        seed_edges,
+        sampler,
+        batch_size=batch_size,
+        device=F.ctx(),
+        use_prefetch_thread=False,
+    )
+    for i, (input_nodes, pair_graph, blocks) in enumerate(dataloader):
+        if isinstance(blocks, list):
+            subg = blocks[0]
+        else:
+            subg = blocks
+        pair_eids = pair_graph.edata[dgl.EID]
+        block_eids = subg.edata[dgl.EID]
+
+        edges_to_exclude = _find_edges_to_exclude(
+            g, exclude, always_exclude, pair_eids
+        )
+        if edges_to_exclude is None:
+            continue
+        edges_to_exclude = dgl.utils.recursive_apply(
+            edges_to_exclude, lambda x: x.cpu().numpy()
+        )
+        block_eids = dgl.utils.recursive_apply(
+            block_eids, lambda x: x.cpu().numpy()
+        )
+
+        if isinstance(edges_to_exclude, Mapping):
+            for k in edges_to_exclude.keys():
+                assert not np.isin(edges_to_exclude[k], block_eids[k]).any()
+        else:
+            assert not np.isin(edges_to_exclude, block_eids).any()
+
+        if i == 10:
+            break
+
+
+def test_edge_dataloader_exclusion_with_reverse_seed_nodes():
+    utype, etype, vtype = ("A", "AB", "B")
+    s = torch.randint(0, 20, (500,), device=F.ctx())
+    d = torch.randint(0, 20, (500,), device=F.ctx())
+    s, d = _remove_duplicates(s, d)
+    g = dgl.heterograph({("A", "AB", "B"): (s, d), ("B", "BA", "A"): (d, s)})
+    sampler = dgl.dataloading.as_edge_prediction_sampler(
+        dgl.dataloading.NeighborSampler(fanouts=[2, 2, 2]),
+        exclude="reverse_types",
+        reverse_etypes={"AB": "BA", "BA": "AB"},
+    )
+    seed_edges = {
+        "AB": torch.arange(g.number_of_edges("AB"), device=F.ctx()),
+        "BA": torch.arange(g.number_of_edges("BA"), device=F.ctx()),
+    }
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        seed_edges,
+        sampler,
+        batch_size=2,
+        device=F.ctx(),
+        shuffle=True,
+        drop_last=False,
+    )
+    for _, pos_graph, mfgs in dataloader:
+        s, d = pos_graph["AB"].edges()
+        AB_pos = list(zip(s.tolist(), d.tolist()))
+        s, d = pos_graph["BA"].edges()
+        BA_pos = list(zip(s.tolist(), d.tolist()))
+
+        s, d = mfgs[-1]["AB"].edges()
+        AB_mfg = list(zip(s.tolist(), d.tolist()))
+        s, d = mfgs[-1]["BA"].edges()
+        BA_mfg = list(zip(s.tolist(), d.tolist()))
+
+        assert all(edge not in AB_mfg for edge in AB_pos)
+        assert all(edge not in BA_mfg for edge in BA_pos)
+
+
+def test_edge_dataloader_exclusion_without_all_reverses():
+    data_dict = {
+        ("A", "AB", "B"): (torch.tensor([0, 1]), torch.tensor([0, 1])),
+        ("B", "BA", "A"): (torch.tensor([0, 1]), torch.tensor([0, 1])),
+        ("B", "BC", "C"): (torch.tensor([0]), torch.tensor([0])),
+        ("C", "CA", "A"): (torch.tensor([0, 1]), torch.tensor([0, 1])),
+    }
+    g = dgl.heterograph(data_dict=data_dict)
+    block_sampler = dgl.dataloading.MultiLayerNeighborSampler(
+        fanouts=[1], replace=True
+    )
+    block_sampler = dgl.dataloading.as_edge_prediction_sampler(
+        block_sampler,
+        exclude="reverse_types",
+        reverse_etypes={"AB": "BA"},
+    )
+    d = dgl.dataloading.DataLoader(
+        graph=g,
+        indices={
+            "AB": torch.tensor([0]),
+            "BC": torch.tensor([0]),
+        },
+        graph_sampler=block_sampler,
+        batch_size=2,
+        shuffle=True,
+        drop_last=False,
+        num_workers=0,
+        device=F.ctx(),
+        use_ddp=False,
+    )
+
+    next(iter(d))
+
+
+def dummy_worker_init_fn(worker_id):
+    pass
+
+
+def test_dataloader_worker_init_fn():
+    dataset = dgl.data.CoraFullDataset()
+    g = dataset[0]
+    sampler = dgl.dataloading.MultiLayerNeighborSampler([2])
+    dataloader = dgl.dataloading.DataLoader(
+        g,
+        torch.arange(100),
+        sampler,
+        batch_size=4,
+        num_workers=4,
+        worker_init_fn=dummy_worker_init_fn,
+    )
+    for _ in dataloader:
+        pass
+
+
+def test_distributed_dataloaders():
+    # Test distributed dataloaders could be successfully imported.
+    try:
+        from dgl.dataloading import (
+            DistDataLoader,
+            DistEdgeDataLoader,
+            DistNodeDataLoader,
+            EdgeCollator,
+            NodeCollator,
+        )
+    except ImportError:
+        pytest.fail("Distributed DataLoader from dataloading import failed")
+
+    try:
+        from dgl.distributed import (
+            DistDataLoader,
+            DistEdgeDataLoader,
+            DistNodeDataLoader,
+            EdgeCollator,
+            NodeCollator,
+        )
+    except ImportError:
+        pytest.fail("Distributed DataLoader from dataloading import failed")
+
+
+if __name__ == "__main__":
+    # test_node_dataloader(F.int32, 'neighbor', None)
+    test_edge_dataloader_excludes(
+        "reverse_types", False, 1, dgl.dataloading.ShaDowKHopSampler([5])
+    )
+    test_edge_dataloader_exclusion_without_all_reverses()

tests/python/pytorch/dataloading/test_spot_target.py

@@ -0,0 +1,83 @@
+from collections.abc import Mapping
+
+import dgl
+import numpy as np
+import pytest
+import torch
+
+
+def _create_homogeneous():
+    s = torch.randint(0, 200, (1000,))
+    d = torch.randint(0, 200, (1000,))
+    g = dgl.graph((s, d), num_nodes=200)
+    reverse_eids = torch.cat([torch.arange(1000, 2000), torch.arange(0, 1000)])
+    seed_edges = torch.arange(0, 1000)
+    return g, reverse_eids, seed_edges
+
+
+def _find_edges_to_exclude(g, pair_eids, degree_threshold):
+    src, dst = g.find_edges(pair_eids)
+    head_degree = g.in_degrees(src)
+    tail_degree = g.in_degrees(dst)
+    degree = torch.min(head_degree, tail_degree)
+    degree_mask = degree < degree_threshold
+    low_degree_pair_eids = pair_eids[degree_mask]
+    low_degree_pair_eids = torch.cat(
+        [low_degree_pair_eids, low_degree_pair_eids + 1000]
+    )
+    return low_degree_pair_eids
+
+
+@pytest.mark.parametrize("degree_threshold", [1, 2, 3, 4, 5])
+@pytest.mark.parametrize("batch_size", [1, 10, 50])
+def test_spot_target_excludes(degree_threshold, batch_size):
+    g, reverse_eids, seed_edges = _create_homogeneous()
+    sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
+    low_degree_excluder = dgl.dataloading.SpotTarget(
+        g,
+        exclude="reverse_id",
+        degree_threshold=degree_threshold,
+        reverse_eids=reverse_eids,
+    )
+    sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler,
+        exclude=low_degree_excluder,
+        negative_sampler=dgl.dataloading.negative_sampler.Uniform(1),
+    )
+    dataloader = dgl.dataloading.DataLoader(
+        g, seed_edges, sampler, batch_size=batch_size
+    )
+
+    for i, (input_nodes, pair_graph, neg_pair_graph, blocks) in enumerate(
+        dataloader
+    ):
+        if isinstance(blocks, list):
+            subg = blocks[0]
+        else:
+            subg = blocks
+        pair_eids = pair_graph.edata[dgl.EID]
+        block_eids = subg.edata[dgl.EID]
+        edges_to_exclude = _find_edges_to_exclude(
+            g, pair_eids, degree_threshold
+        )
+        if edges_to_exclude is None:
+            continue
+        edges_to_exclude = dgl.utils.recursive_apply(
+            edges_to_exclude, lambda x: x.cpu().numpy()
+        )
+        block_eids = dgl.utils.recursive_apply(
+            block_eids, lambda x: x.cpu().numpy()
+        )
+
+        if isinstance(edges_to_exclude, Mapping):
+            for k in edges_to_exclude.keys():
+                assert not np.isin(edges_to_exclude[k], block_eids[k]).any()
+        else:
+            assert not np.isin(edges_to_exclude, block_eids).any()
+
+        if i == 10:
+            break
+
+
+if __name__ == "__main__":
+    test_spot_target_excludes(degree_threshold=2, batch_size=10)

tests/python/pytorch/distributed/optim/test_dist_optim.py

@@ -0,0 +1,200 @@
+import os
+
+os.environ["OMP_NUM_THREADS"] = "1"
+import multiprocessing as mp
+import pickle
+import random
+import socket
+import sys
+import time
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+import torch as th
+from dgl import function as fn
+from dgl.distributed import (
+    DistEmbedding,
+    DistGraph,
+    DistGraphServer,
+    load_partition_book,
+    partition_graph,
+)
+from dgl.distributed.optim import SparseAdagrad, SparseAdam
+from scipy import sparse as spsp
+
+# Set seeds to make tests fully reproducible.
+SEED = 12345  # random.randint(1, 99999)
+F.seed(SEED)
+
+
+def create_random_graph(n):
+    arr = (
+        spsp.random(n, n, density=0.001, format="coo", random_state=100) != 0
+    ).astype(np.int64)
+    return dgl.from_scipy(arr)
+
+
+def get_local_usable_addr():
+    """Get local usable IP and port
+
+    Returns
+    -------
+    str
+        IP address, e.g., '192.168.8.12:50051'
+    """
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # doesn't even have to be reachable
+        sock.connect(("10.255.255.255", 1))
+        ip_addr = sock.getsockname()[0]
+    except ValueError:
+        ip_addr = "127.0.0.1"
+    finally:
+        sock.close()
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.bind(("", 0))
+    sock.listen(1)
+    port = sock.getsockname()[1]
+    sock.close()
+
+    return ip_addr + " " + str(port)
+
+
+def prepare_dist():
+    ip_config = open("optim_ip_config.txt", "w")
+    ip_addr = get_local_usable_addr()
+    ip_config.write("{}\n".format(ip_addr))
+    ip_config.close()
+
+
+def run_server(graph_name, server_id, server_count, num_clients, shared_mem):
+    g = DistGraphServer(
+        server_id,
+        "optim_ip_config.txt",
+        num_clients,
+        server_count,
+        "/tmp/dist_graph/{}.json".format(graph_name),
+        disable_shared_mem=not shared_mem,
+    )
+    print("start server", server_id)
+    g.start()
+
+
+def initializer(shape, dtype):
+    arr = th.zeros(shape, dtype=dtype)
+    th.manual_seed(0)
+    th.nn.init.uniform_(arr, 0, 1.0)
+    return arr
+
+
+def run_client(graph_name, cli_id, part_id, server_count):
+    device = F.ctx()
+    time.sleep(5)
+    os.environ["DGL_NUM_SERVER"] = str(server_count)
+    dgl.distributed.initialize("optim_ip_config.txt")
+    gpb, graph_name, _, _ = load_partition_book(
+        "/tmp/dist_graph/{}.json".format(graph_name), part_id
+    )
+    g = DistGraph(graph_name, gpb=gpb)
+    policy = dgl.distributed.PartitionPolicy("node", g.get_partition_book())
+    num_nodes = g.num_nodes()
+    emb_dim = 4
+    dgl_emb = DistEmbedding(
+        num_nodes,
+        emb_dim,
+        name="optim",
+        init_func=initializer,
+        part_policy=policy,
+    )
+    dgl_emb_zero = DistEmbedding(
+        num_nodes,
+        emb_dim,
+        name="optim-zero",
+        init_func=initializer,
+        part_policy=policy,
+    )
+    dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
+    dgl_adam._world_size = 1
+    dgl_adam._rank = 0
+
+    torch_emb = th.nn.Embedding(num_nodes, emb_dim, sparse=True)
+    torch_emb_zero = th.nn.Embedding(num_nodes, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb_zero.weight, 0, 1.0)
+    torch_adam = th.optim.SparseAdam(
+        list(torch_emb.parameters()) + list(torch_emb_zero.parameters()),
+        lr=0.01,
+    )
+
+    labels = th.ones((4,)).long()
+    idx = th.randint(0, num_nodes, size=(4,))
+    dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
+    torch_value = torch_emb(idx)
+    torch_adam.zero_grad()
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    torch_loss.backward()
+    torch_adam.step()
+
+    dgl_adam.zero_grad()
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    dgl_loss.backward()
+    dgl_adam.step()
+
+    assert F.allclose(
+        dgl_emb.weight[0 : num_nodes // 2], torch_emb.weight[0 : num_nodes // 2]
+    )
+
+
+def check_sparse_adam(num_trainer=1, shared_mem=True):
+    prepare_dist()
+    g = create_random_graph(2000)
+    num_servers = num_trainer
+    num_clients = num_trainer
+    num_parts = 1
+
+    graph_name = "dist_graph_test"
+    partition_graph(g, graph_name, num_parts, "/tmp/dist_graph")
+
+    # let's just test on one partition for now.
+    # We cannot run multiple servers and clients on the same machine.
+    serv_ps = []
+    ctx = mp.get_context("spawn")
+    for serv_id in range(num_servers):
+        p = ctx.Process(
+            target=run_server,
+            args=(graph_name, serv_id, num_servers, num_clients, shared_mem),
+        )
+        serv_ps.append(p)
+        p.start()
+
+    cli_ps = []
+    for cli_id in range(num_clients):
+        print("start client", cli_id)
+        p = ctx.Process(
+            target=run_client, args=(graph_name, cli_id, 0, num_servers)
+        )
+        p.start()
+        cli_ps.append(p)
+
+    for p in cli_ps:
+        p.join()
+
+    for p in serv_ps:
+        p.join()
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+def test_sparse_opt():
+    os.environ["DGL_DIST_MODE"] = "distributed"
+    check_sparse_adam(1, True)
+    check_sparse_adam(1, False)
+
+
+if __name__ == "__main__":
+    os.makedirs("/tmp/dist_graph", exist_ok=True)
+    test_sparse_opt()

tests/python/pytorch/geometry/test_geometry.py

@@ -0,0 +1,254 @@
+import backend as F
+
+import dgl
+import dgl.nn
+import numpy as np
+import pytest
+import torch as th
+from dgl import DGLError
+from dgl.base import DGLWarning
+from dgl.geometry import farthest_point_sampler, neighbor_matching
+from utils import parametrize_idtype
+from utils.graph_cases import get_cases
+
+
+def test_fps():
+    N = 1000
+    batch_size = 5
+    sample_points = 10
+    x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
+    ctx = F.ctx()
+    if F.gpu_ctx():
+        x = x.to(ctx)
+    res = farthest_point_sampler(x, sample_points)
+    assert res.shape[0] == batch_size
+    assert res.shape[1] == sample_points
+    assert res.sum() > 0
+
+
+def test_fps_start_idx():
+    N = 1000
+    batch_size = 5
+    sample_points = 10
+    x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
+    ctx = F.ctx()
+    if F.gpu_ctx():
+        x = x.to(ctx)
+    res = farthest_point_sampler(x, sample_points, start_idx=0)
+    assert th.any(res[:, 0] == 0)
+
+
+def _test_knn_common(device, algorithm, dist, exclude_self):
+    x = th.randn(8, 3).to(device)
+    kg = dgl.nn.KNNGraph(3)
+    if dist == "euclidean":
+        d = th.cdist(x, x).to(F.cpu())
+    else:
+        x = x + th.randn(1).item()
+        tmp_x = x / (1e-5 + F.sqrt(F.sum(x * x, dim=1, keepdims=True)))
+        d = 1 - F.matmul(tmp_x, tmp_x.T).to(F.cpu())
+
+    def check_knn(g, x, start, end, k, exclude_self, check_indices=True):
+        assert g.device == x.device
+        g = g.to(F.cpu())
+        for v in range(start, end):
+            src, _ = g.in_edges(v)
+            src = set(src.numpy())
+            assert len(src) == k
+            if check_indices:
+                i = v - start
+                src_ans = set(
+                    th.topk(
+                        d[start:end, start:end][i],
+                        k + (1 if exclude_self else 0),
+                        largest=False,
+                    )[1].numpy()
+                    + start
+                )
+                if exclude_self:
+                    # remove self
+                    src_ans.remove(v)
+                assert src == src_ans
+
+    def check_batch(g, k, expected_batch_info):
+        assert F.array_equal(g.batch_num_nodes(), F.tensor(expected_batch_info))
+        assert F.array_equal(
+            g.batch_num_edges(), k * F.tensor(expected_batch_info)
+        )
+
+    # check knn with 2d input
+    g = kg(x, algorithm, dist, exclude_self)
+    check_knn(g, x, 0, 8, 3, exclude_self)
+    check_batch(g, 3, [8])
+
+    # check knn with 3d input
+    g = kg(x.view(2, 4, 3), algorithm, dist, exclude_self)
+    check_knn(g, x, 0, 4, 3, exclude_self)
+    check_knn(g, x, 4, 8, 3, exclude_self)
+    check_batch(g, 3, [4, 4])
+
+    # check segmented knn
+    # there are only 2 edges per node possible when exclude_self with 3 nodes in the segment
+    # and this test case isn't supposed to warn, so limit it when exclude_self is True
+    adjusted_k = 3 - (1 if exclude_self else 0)
+    kg = dgl.nn.SegmentedKNNGraph(adjusted_k)
+    g = kg(x, [3, 5], algorithm, dist, exclude_self)
+    check_knn(g, x, 0, 3, adjusted_k, exclude_self)
+    check_knn(g, x, 3, 8, adjusted_k, exclude_self)
+    check_batch(g, adjusted_k, [3, 5])
+
+    # check k > num_points
+    kg = dgl.nn.KNNGraph(10)
+    with pytest.warns(DGLWarning):
+        g = kg(x, algorithm, dist, exclude_self)
+    # there are only 7 edges per node possible when exclude_self with 8 nodes total
+    adjusted_k = 8 - (1 if exclude_self else 0)
+    check_knn(g, x, 0, 8, adjusted_k, exclude_self)
+    check_batch(g, adjusted_k, [8])
+
+    with pytest.warns(DGLWarning):
+        g = kg(x.view(2, 4, 3), algorithm, dist, exclude_self)
+    # there are only 3 edges per node possible when exclude_self with 4 nodes per segment
+    adjusted_k = 4 - (1 if exclude_self else 0)
+    check_knn(g, x, 0, 4, adjusted_k, exclude_self)
+    check_knn(g, x, 4, 8, adjusted_k, exclude_self)
+    check_batch(g, adjusted_k, [4, 4])
+
+    kg = dgl.nn.SegmentedKNNGraph(5)
+    with pytest.warns(DGLWarning):
+        g = kg(x, [3, 5], algorithm, dist, exclude_self)
+    # there are only 2 edges per node possible when exclude_self in the segment with
+    # only 3 nodes, and the current implementation reduces k for all segments
+    # in that case
+    adjusted_k = 3 - (1 if exclude_self else 0)
+    check_knn(g, x, 0, 3, adjusted_k, exclude_self)
+    check_knn(g, x, 3, 8, adjusted_k, exclude_self)
+    check_batch(g, adjusted_k, [3, 5])
+
+    # check k == 0
+    # that's valid for exclude_self, but -1 is not, so check -1 instead for exclude_self
+    adjusted_k = 0 - (1 if exclude_self else 0)
+    kg = dgl.nn.KNNGraph(adjusted_k)
+    with pytest.raises(DGLError):
+        g = kg(x, algorithm, dist, exclude_self)
+    kg = dgl.nn.SegmentedKNNGraph(adjusted_k)
+    with pytest.raises(DGLError):
+        g = kg(x, [3, 5], algorithm, dist, exclude_self)
+
+    # check empty
+    x_empty = th.tensor([])
+    kg = dgl.nn.KNNGraph(3)
+    with pytest.raises(DGLError):
+        g = kg(x_empty, algorithm, dist, exclude_self)
+    kg = dgl.nn.SegmentedKNNGraph(3)
+    with pytest.raises(DGLError):
+        g = kg(x_empty, [3, 5], algorithm, dist, exclude_self)
+
+    # check all coincident points
+    x = th.zeros((20, 3)).to(device)
+    kg = dgl.nn.KNNGraph(3)
+    g = kg(x, algorithm, dist, exclude_self)
+    # different algorithms may break the tie differently, so don't check the indices
+    check_knn(g, x, 0, 20, 3, exclude_self, False)
+    check_batch(g, 3, [20])
+
+    # check all coincident points
+    kg = dgl.nn.SegmentedKNNGraph(3)
+    g = kg(x, [4, 7, 5, 4], algorithm, dist, exclude_self)
+    # different algorithms may break the tie differently, so don't check the indices
+    check_knn(g, x, 0, 4, 3, exclude_self, False)
+    check_knn(g, x, 4, 11, 3, exclude_self, False)
+    check_knn(g, x, 11, 16, 3, exclude_self, False)
+    check_knn(g, x, 16, 20, 3, exclude_self, False)
+    check_batch(g, 3, [4, 7, 5, 4])
+
+
+@pytest.mark.parametrize(
+    "algorithm", ["bruteforce-blas", "bruteforce", "kd-tree"]
+)
+@pytest.mark.parametrize("dist", ["euclidean", "cosine"])
+@pytest.mark.parametrize("exclude_self", [False, True])
+def test_knn_cpu(algorithm, dist, exclude_self):
+    _test_knn_common(F.cpu(), algorithm, dist, exclude_self)
+
+
+@pytest.mark.parametrize(
+    "algorithm", ["bruteforce-blas", "bruteforce", "bruteforce-sharemem"]
+)
+@pytest.mark.parametrize("dist", ["euclidean", "cosine"])
+@pytest.mark.parametrize("exclude_self", [False, True])
+def test_knn_cuda(algorithm, dist, exclude_self):
+    if not th.cuda.is_available():
+        return
+    _test_knn_common(F.cuda(), algorithm, dist, exclude_self)
+
+
+@pytest.mark.parametrize("num_points", [8, 64, 256, 1024])
+def test_knn_sharedmem_large(num_points):
+    if not th.cuda.is_available():
+        return
+    x = th.randn(num_points, 5, device="cuda")
+    y = th.randn(num_points, 5, device="cuda")
+    k = 4
+
+    def ground_truth(x, y, k):
+        dist = (
+            th.sum(x * x, dim=1)
+            + th.sum(y * y, dim=1).unsqueeze(-1)
+            - 2 * th.mm(y, x.T)
+        )
+        ret = th.topk(dist, k, dim=-1, largest=False)[1]
+        return th.sort(ret, dim=-1)[0]
+
+    gt = ground_truth(x, y, k)
+    actual = th.sort(
+        dgl.functional.knn(
+            k, x, [num_points], y, [num_points], algorithm="bruteforce-sharemem"
+        )[1].reshape(-1, k),
+        -1,
+    )[0]
+    assert th.all(actual == gt).item()
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+@pytest.mark.parametrize("weight", [True, False])
+@pytest.mark.parametrize("relabel", [True, False])
+def test_edge_coarsening(idtype, g, weight, relabel):
+    num_nodes = g.num_nodes()
+    g = dgl.to_bidirected(g)
+    g = g.astype(idtype).to(F.ctx())
+    edge_weight = None
+    if weight:
+        edge_weight = F.abs(F.randn((g.num_edges(),))).to(F.ctx())
+    node_labels = neighbor_matching(g, edge_weight, relabel_idx=relabel)
+    unique_ids, counts = th.unique(node_labels, return_counts=True)
+    num_result_ids = unique_ids.size(0)
+
+    # shape correct
+    assert node_labels.shape == (g.num_nodes(),)
+
+    # all nodes marked
+    assert F.reduce_sum(node_labels < 0).item() == 0
+
+    # number of unique node ids correct.
+    assert num_result_ids >= num_nodes // 2 and num_result_ids <= num_nodes
+
+    # each unique id has <= 2 nodes
+    assert F.reduce_sum(counts > 2).item() == 0
+
+    # if two nodes have the same id, they must be neighbors
+    idxs = F.arange(0, num_nodes, idtype)
+    for l in unique_ids:
+        l = l.item()
+        idx = idxs[(node_labels == l)]
+        if idx.size(0) == 2:
+            u, v = idx[0].item(), idx[1].item()
+            assert g.has_edges_between(u, v)
+
+
+if __name__ == "__main__":
+    test_fps()
+    test_fps_start_idx()
+    test_knn()
+    test_knn_sharedmem_large()

tests/python/pytorch/graphbolt/__init__.py

@@ -0,0 +1 @@
+""" DGL graphbolt API tests"""

tests/python/pytorch/graphbolt/gb_test_utils.py

@@ -0,0 +1,374 @@
+import os
+
+import dgl
+import dgl.graphbolt as gb
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+
+
+def rand_csc_graph(N, density, bidirection_edge=False):
+    adj = sp.random(N, N, density)
+    if bidirection_edge:
+        adj = adj + adj.T
+    adj = adj.tocsc()
+
+    indptr = torch.LongTensor(adj.indptr)
+    indices = torch.LongTensor(adj.indices)
+
+    graph = gb.fused_csc_sampling_graph(indptr, indices)
+
+    return graph
+
+
+def random_homo_graph(num_nodes, num_edges):
+    csc_indptr = torch.randint(0, num_edges, (num_nodes + 1,))
+    csc_indptr = torch.sort(csc_indptr)[0]
+    csc_indptr[0] = 0
+    csc_indptr[-1] = num_edges
+    indices = torch.randint(0, num_nodes, (num_edges,))
+    return csc_indptr, indices
+
+
+def get_type_to_id(num_ntypes, num_etypes):
+    ntypes = {f"n{i}": i for i in range(num_ntypes)}
+    etypes = {}
+    count = 0
+    for n1 in range(num_ntypes):
+        for n2 in range(n1, num_ntypes):
+            if count >= num_etypes:
+                break
+            etypes.update({f"n{n1}:e{count}:n{n2}": count})
+            count += 1
+    return ntypes, etypes
+
+
+def get_ntypes_and_etypes(num_nodes, num_ntypes, num_etypes):
+    ntypes = {f"n{i}": num_nodes // num_ntypes for i in range(num_ntypes)}
+    if num_nodes % num_ntypes != 0:
+        ntypes["n0"] += num_nodes % num_ntypes
+    etypes = []
+    count = 0
+    while count < num_etypes:
+        for n1 in range(num_ntypes):
+            for n2 in range(num_ntypes):
+                if count >= num_etypes:
+                    break
+                etypes.append((f"n{n1}", f"e{count}", f"n{n2}"))
+                count += 1
+    return ntypes, etypes
+
+
+def random_hetero_graph(num_nodes, num_edges, num_ntypes, num_etypes):
+    ntypes, etypes = get_ntypes_and_etypes(num_nodes, num_ntypes, num_etypes)
+    edges = {}
+    for step, etype in enumerate(etypes):
+        src_ntype, _, dst_ntype = etype
+        num_e = num_edges // num_etypes + (
+            0 if step != 0 else num_edges % num_etypes
+        )
+        if ntypes[src_ntype] == 0 or ntypes[dst_ntype] == 0:
+            continue
+        src = torch.randint(0, ntypes[src_ntype], (num_e,))
+        dst = torch.randint(0, ntypes[dst_ntype], (num_e,))
+
+        edges[etype] = (src, dst)
+
+    gb_g = gb.from_dglgraph(dgl.heterograph(edges, ntypes))
+    return (
+        gb_g.csc_indptr,
+        gb_g.indices,
+        gb_g.node_type_offset,
+        gb_g.type_per_edge,
+        gb_g.node_type_to_id,
+        gb_g.edge_type_to_id,
+    )
+
+
+def random_homo_graphbolt_graph(
+    test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt="csv"
+):
+    """Generate random graphbolt version homograph"""
+    # Generate random edges.
+    nodes = np.repeat(np.arange(num_nodes, dtype=np.int64), 5)
+    neighbors = np.random.randint(
+        0, num_nodes, size=(num_edges), dtype=np.int64
+    )
+    edges = np.stack([nodes, neighbors], axis=1)
+    os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
+    assert edge_fmt in [
+        "numpy",
+        "csv",
+    ], "Only numpy and csv are supported for edges."
+    if edge_fmt == "csv":
+        # Write into edges/edge.csv
+        edges_DataFrame = pd.DataFrame(edges, columns=["src", "dst"])
+        edge_path = os.path.join("edges", "edge.csv")
+        edges_DataFrame.to_csv(
+            os.path.join(test_dir, edge_path),
+            index=False,
+            header=False,
+        )
+    else:
+        # Write into edges/edge.npy
+        edges = edges.T
+        edge_path = os.path.join("edges", "edge.npy")
+        np.save(os.path.join(test_dir, edge_path), edges)
+
+    # Generate random graph edge-feats.
+    edge_feats = np.random.rand(num_edges, num_classes)
+    os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
+    edge_feat_path = os.path.join("data", "edge-feat.npy")
+    np.save(os.path.join(test_dir, edge_feat_path), edge_feats)
+
+    # Generate random node-feats.
+    if num_classes == 1:
+        node_feats = np.random.rand(num_nodes)
+    else:
+        node_feats = np.random.rand(num_nodes, num_classes)
+    node_feat_path = os.path.join("data", "node-feat.npy")
+    np.save(os.path.join(test_dir, node_feat_path), node_feats)
+
+    # Generate train/test/valid set.
+    assert num_nodes % 4 == 0, "num_nodes must be divisible by 4"
+    each_set_size = num_nodes // 4
+    os.makedirs(os.path.join(test_dir, "set"), exist_ok=True)
+    train_pairs = (
+        np.arange(each_set_size),
+        np.arange(each_set_size, 2 * each_set_size),
+    )
+    train_data = np.vstack(train_pairs).T.astype(edges.dtype)
+    train_path = os.path.join("set", "train.npy")
+    np.save(os.path.join(test_dir, train_path), train_data)
+
+    validation_pairs = (
+        np.arange(each_set_size, 2 * each_set_size),
+        np.arange(2 * each_set_size, 3 * each_set_size),
+    )
+    validation_data = np.vstack(validation_pairs).T.astype(edges.dtype)
+    validation_path = os.path.join("set", "validation.npy")
+    np.save(os.path.join(test_dir, validation_path), validation_data)
+
+    test_pairs = (
+        np.arange(2 * each_set_size, 3 * each_set_size),
+        np.arange(3 * each_set_size, 4 * each_set_size),
+    )
+    test_data = np.vstack(test_pairs).T.astype(edges.dtype)
+    test_path = os.path.join("set", "test.npy")
+    np.save(os.path.join(test_dir, test_path), test_data)
+
+    yaml_content = f"""
+        dataset_name: {dataset_name}
+        graph: # Graph structure and required attributes.
+            nodes:
+                - num: {num_nodes}
+            edges:
+                - format: {edge_fmt}
+                  path: {edge_path}
+            feature_data:
+                - domain: node
+                  type: null
+                  name: feat
+                  format: numpy
+                  in_memory: true
+                  path: {node_feat_path}
+                - domain: edge
+                  type: null
+                  name: feat
+                  format: numpy
+                  in_memory: true
+                  path: {edge_feat_path}
+        feature_data:
+            - domain: node
+              type: null
+              name: feat
+              format: numpy
+              in_memory: true
+              path: {node_feat_path}
+            - domain: edge
+              type: null
+              name: feat
+              format: numpy
+              path: {edge_feat_path}
+        tasks:
+          - name: link_prediction
+            num_classes: {num_classes}
+            train_set:
+              - type: null
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {train_path}
+            validation_set:
+              - type: null
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {validation_path}
+            test_set:
+              - type: null
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {test_path}
+    """
+    return yaml_content
+
+
+def generate_raw_data_for_hetero_dataset(
+    test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt="csv"
+):
+    # Generate edges.
+    edges_path = {}
+    for etype, num_edge in num_edges.items():
+        src_ntype, etype_str, dst_ntype = etype
+        src = torch.randint(0, num_nodes[src_ntype], (num_edge,))
+        dst = torch.randint(0, num_nodes[dst_ntype], (num_edge,))
+        os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
+        assert edge_fmt in [
+            "numpy",
+            "csv",
+        ], "Only numpy and csv are supported for edges."
+        if edge_fmt == "csv":
+            # Write into edges/edge.csv
+            edges = pd.DataFrame(
+                np.stack([src, dst], axis=1), columns=["src", "dst"]
+            )
+            edge_path = os.path.join("edges", f"{etype_str}.csv")
+            edges.to_csv(
+                os.path.join(test_dir, edge_path),
+                index=False,
+                header=False,
+            )
+        else:
+            edges = np.stack([src, dst], axis=1).T
+            edge_path = os.path.join("edges", f"{etype_str}.npy")
+            np.save(os.path.join(test_dir, edge_path), edges)
+        edges_path[etype_str] = edge_path
+
+    # Generate node features.
+    node_feats_path = {}
+    os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
+    for ntype, num_node in num_nodes.items():
+        node_feat_path = os.path.join("data", f"{ntype}-feat.npy")
+        node_feats = np.random.rand(num_node, num_classes)
+        np.save(os.path.join(test_dir, node_feat_path), node_feats)
+        node_feats_path[ntype] = node_feat_path
+
+    # Generate edge features.
+    edge_feats_path = {}
+    os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
+    for etype, num_edge in num_edges.items():
+        src_ntype, etype_str, dst_ntype = etype
+        edge_feat_path = os.path.join("data", f"{etype_str}-feat.npy")
+        edge_feats = np.random.rand(num_edge, num_classes)
+        np.save(os.path.join(test_dir, edge_feat_path), edge_feats)
+        edge_feats_path[etype_str] = edge_feat_path
+
+    # Generate train/test/valid set.
+    os.makedirs(os.path.join(test_dir, "set"), exist_ok=True)
+    user_ids = torch.arange(num_nodes["user"])
+    np.random.shuffle(user_ids.numpy())
+    num_train = int(num_nodes["user"] * 0.6)
+    num_validation = int(num_nodes["user"] * 0.2)
+    num_test = num_nodes["user"] - num_train - num_validation
+    train_path = os.path.join("set", "train.npy")
+    np.save(os.path.join(test_dir, train_path), user_ids[:num_train])
+    validation_path = os.path.join("set", "validation.npy")
+    np.save(
+        os.path.join(test_dir, validation_path),
+        user_ids[num_train : num_train + num_validation],
+    )
+    test_path = os.path.join("set", "test.npy")
+    np.save(
+        os.path.join(test_dir, test_path),
+        user_ids[num_train + num_validation :],
+    )
+
+    yaml_content = f"""
+        dataset_name: {dataset_name}
+        graph: # Graph structure and required attributes.
+          nodes:
+            - type: user
+              num: {num_nodes["user"]}
+            - type: item
+              num: {num_nodes["item"]}
+          edges:
+            - type: "user:follow:user"
+              format: {edge_fmt}
+              path: {edges_path["follow"]}
+            - type: "user:click:item"
+              format: {edge_fmt}
+              path: {edges_path["click"]}
+          feature_data:
+            - domain: node
+              type: user
+              name: feat
+              format: numpy
+              in_memory: true
+              path: {node_feats_path["user"]}
+            - domain: node
+              type: item
+              name: feat
+              format: numpy
+              in_memory: true
+              path: {node_feats_path["item"]}
+            - domain: edge
+              type: "user:follow:user"
+              name: feat
+              format: numpy
+              in_memory: true
+              path: {edge_feats_path["follow"]}
+            - domain: edge
+              type: "user:click:item"
+              name: feat
+              format: numpy
+              in_memory: true
+              path: {edge_feats_path["click"]}
+        feature_data:
+          - domain: node
+            type: user
+            name: feat
+            format: numpy
+            in_memory: true
+            path: {node_feats_path["user"]}
+          - domain: node
+            type: item
+            name: feat
+            format: numpy
+            in_memory: true
+            path: {node_feats_path["item"]}
+        tasks:
+          - name: node_classification
+            num_classes: {num_classes}
+            train_set:
+              - type: user
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {train_path}
+            validation_set:
+              - type: user
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {validation_path}
+            test_set:
+              - type: user
+                data:
+                  - name: seeds
+                    format: numpy
+                    in_memory: true
+                    path: {test_path}
+    """
+
+    yaml_file = os.path.join(test_dir, "metadata.yaml")
+    with open(yaml_file, "w") as f:
+        f.write(yaml_content)

tests/python/pytorch/graphbolt/impl/__init__.py

@@ -0,0 +1 @@
+""" DGL graphbolt/impl tests"""

tests/python/pytorch/graphbolt/impl/test_basic_feature_store.py

@@ -0,0 +1,151 @@
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def test_basic_feature_store_homo():
+    a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+    b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [4, 3]]])
+    metadata = {"max_value": 3}
+
+    features = {}
+    features[("node", None, "a")] = gb.TorchBasedFeature(a, metadata=metadata)
+    features[("node", None, "b")] = gb.TorchBasedFeature(b)
+
+    feature_store = gb.BasicFeatureStore(features)
+
+    # Test __getitem__ to access the stored Feature.
+    feature = feature_store[("node", None, "a")]
+    assert isinstance(feature, gb.Feature)
+    assert torch.equal(
+        feature.read(),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+
+    # Test read the entire feature.
+    assert torch.equal(
+        feature_store.read("node", None, "a"),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+    assert torch.equal(
+        feature_store.read("node", None, "b"),
+        torch.tensor([[[1, 2], [3, 4]], [[2, 5], [4, 3]]]),
+    )
+
+    # Test read with ids.
+    assert torch.equal(
+        feature_store.read("node", None, "a", torch.tensor([0])),
+        torch.tensor([[1, 2, 4]]),
+    )
+    assert torch.equal(
+        feature_store.read("node", None, "b", torch.tensor([0])),
+        torch.tensor([[[1, 2], [3, 4]]]),
+    )
+
+    # Test get the size and count of the entire feature.
+    assert feature_store.size("node", None, "a") == torch.Size([3])
+    assert feature_store.size("node", None, "b") == torch.Size([2, 2])
+    assert feature_store.count("node", None, "a") == a.size(0)
+    assert feature_store.count("node", None, "b") == b.size(0)
+
+    # Test get metadata of the feature.
+    assert feature_store.metadata("node", None, "a") == metadata
+    assert feature_store.metadata("node", None, "b") == {}
+
+    # Test __setitem__ and __contains__ of FeatureStore.
+    assert ("node", None, "c") not in feature_store
+    feature_store[("node", None, "c")] = feature_store[("node", None, "a")]
+    assert ("node", None, "c") in feature_store
+
+    # Test get keys of the features.
+    assert feature_store.keys() == [
+        ("node", None, "a"),
+        ("node", None, "b"),
+        ("node", None, "c"),
+    ]
+
+
+def test_basic_feature_store_hetero():
+    a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+    b = torch.tensor([[[6], [8]], [[8], [9]]])
+    metadata = {"max_value": 3}
+
+    features = {}
+    features[("node", "author", "a")] = gb.TorchBasedFeature(
+        a, metadata=metadata
+    )
+    features[("edge", "paper:cites", "b")] = gb.TorchBasedFeature(b)
+
+    feature_store = gb.BasicFeatureStore(features)
+
+    # Test __getitem__ to access the stored Feature.
+    feature = feature_store[("node", "author", "a")]
+    assert isinstance(feature, gb.Feature)
+    assert torch.equal(
+        feature.read(),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+
+    # Test read the entire feature.
+    assert torch.equal(
+        feature_store.read("node", "author", "a"),
+        torch.tensor([[1, 2, 4], [2, 5, 3]]),
+    )
+    assert torch.equal(
+        feature_store.read("edge", "paper:cites", "b"),
+        torch.tensor([[[6], [8]], [[8], [9]]]),
+    )
+
+    # Test read with ids.
+    assert torch.equal(
+        feature_store.read("node", "author", "a", torch.tensor([0])),
+        torch.tensor([[1, 2, 4]]),
+    )
+
+    # Test get the size of the entire feature.
+    assert feature_store.size("node", "author", "a") == torch.Size([3])
+    assert feature_store.size("edge", "paper:cites", "b") == torch.Size([2, 1])
+
+    # Test get metadata of the feature.
+    assert feature_store.metadata("node", "author", "a") == metadata
+    assert feature_store.metadata("edge", "paper:cites", "b") == {}
+
+    # Test __setitem__ and __contains__ of FeatureStore.
+    assert ("node", "author", "c") not in feature_store
+    feature_store[("node", "author", "c")] = feature_store[
+        ("node", "author", "a")
+    ]
+    assert ("node", "author", "c") in feature_store
+
+    # Test get keys of the features.
+    assert feature_store.keys() == [
+        ("node", "author", "a"),
+        ("edge", "paper:cites", "b"),
+        ("node", "author", "c"),
+    ]
+
+
+def test_basic_feature_store_errors():
+    a = torch.tensor([3, 2, 1])
+    b = torch.tensor([[1, 2, 4], [2, 5, 3]])
+
+    features = {}
+    # Test error when dimension of the value is illegal.
+    with pytest.raises(
+        AssertionError,
+        match=rf"dimension of torch_feature in TorchBasedFeature must be "
+        rf"greater than 1, but got {a.dim()} dimension.",
+    ):
+        features[("node", "paper", "a")] = gb.TorchBasedFeature(a)
+    features[("node", "author", "b")] = gb.TorchBasedFeature(b)
+
+    feature_store = gb.BasicFeatureStore(features)
+
+    # Test error when key does not exist.
+    with pytest.raises(KeyError):
+        feature_store.read("node", "paper", "b")
+
+    # Test error when at least one id is out of bound.
+    with pytest.raises(IndexError):
+        feature_store.read("node", "author", "b", torch.tensor([0, 3]))

tests/python/pytorch/graphbolt/impl/test_cooperative_minibatching_utils.py

@@ -0,0 +1,68 @@
+import unittest
+
+from functools import partial
+
+import backend as F
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+WORLD_SIZE = 7
+
+assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
+
+
+@unittest.skipIf(
+    F._default_context_str != "gpu",
+    reason="This test requires an NVIDIA GPU.",
+)
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("rank", list(range(WORLD_SIZE)))
+def test_rank_sort_and_unique_and_compact(dtype, rank):
+    torch.manual_seed(7)
+    nodes_list1 = [
+        torch.randint(0, 2111111111, [777], dtype=dtype, device=F.ctx())
+        for _ in range(10)
+    ]
+    nodes_list2 = [nodes.sort()[0] for nodes in nodes_list1]
+
+    res1 = torch.ops.graphbolt.rank_sort(nodes_list1, rank, WORLD_SIZE)
+    res2 = torch.ops.graphbolt.rank_sort(nodes_list2, rank, WORLD_SIZE)
+
+    for i, ((nodes1, idx1, offsets1), (nodes2, idx2, offsets2)) in enumerate(
+        zip(res1, res2)
+    ):
+        assert_equal(nodes_list1[i], nodes1[idx1])
+        assert_equal(nodes_list2[i], nodes2[idx2])
+        assert_equal(offsets1, offsets2)
+        assert offsets1.is_pinned() and offsets2.is_pinned()
+
+    res3 = torch.ops.graphbolt.rank_sort(nodes_list1, rank, WORLD_SIZE)
+
+    # This function is deterministic. Call with identical arguments and check.
+    for (nodes1, idx1, offsets1), (nodes3, idx3, offsets3) in zip(res1, res3):
+        assert_equal(nodes1, nodes3)
+        assert_equal(idx1, idx3)
+        assert_equal(offsets1, offsets3)
+
+    # The dependency on the rank argument is simply a permutation.
+    res4 = torch.ops.graphbolt.rank_sort(nodes_list1, 0, WORLD_SIZE)
+    for (nodes1, idx1, offsets1), (nodes4, idx4, offsets4) in zip(res1, res4):
+        off1 = offsets1.tolist()
+        off4 = offsets4.tolist()
+        assert_equal(nodes1[idx1], nodes4[idx4])
+        for i in range(WORLD_SIZE):
+            j = (i - rank + WORLD_SIZE) % WORLD_SIZE
+            assert_equal(
+                nodes1[off1[j] : off1[j + 1]], nodes4[off4[i] : off4[i + 1]]
+            )
+
+    unique, compacted, offsets = gb.unique_and_compact(
+        nodes_list1[:1], rank, WORLD_SIZE
+    )
+
+    nodes1, idx1, offsets1 = res1[0]
+
+    assert_equal(unique, nodes1)
+    assert_equal(compacted[0], idx1)
+    assert_equal(offsets, offsets1)

tests/python/pytorch/graphbolt/impl/test_cpu_cached_feature.py

@@ -0,0 +1,184 @@
+import os
+import tempfile
+import unittest
+
+import backend as F
+import numpy as np
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def to_on_disk_numpy(test_dir, name, t):
+    path = os.path.join(test_dir, name + ".npy")
+    np.save(path, t.numpy())
+    return path
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("policy", ["s3-fifo", "sieve", "lru", "clock"])
+def test_cpu_cached_feature(dtype, policy):
+    cache_size_a = 32
+    cache_size_b = 64
+    a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype)
+    b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]], dtype=dtype)
+
+    pin_memory = F._default_context_str == "gpu"
+
+    cache_size_a *= a[:1].nbytes
+    cache_size_b *= b[:1].nbytes
+
+    feat_store_a = gb.cpu_cached_feature(
+        gb.TorchBasedFeature(a), cache_size_a, policy, pin_memory
+    )
+    feat_store_b = gb.cpu_cached_feature(
+        gb.TorchBasedFeature(b), cache_size_b, policy, pin_memory
+    )
+
+    # Test read the entire feature.
+    assert torch.equal(feat_store_a.read(), a)
+    assert torch.equal(feat_store_b.read(), b)
+
+    # Test read with ids.
+    assert torch.equal(
+        # Test read when ids are on a different device.
+        feat_store_a.read(torch.tensor([0], device=F.ctx())),
+        torch.tensor([[1, 2, 3]], dtype=dtype, device=F.ctx()),
+    )
+    assert torch.equal(
+        feat_store_b.read(torch.tensor([1, 1])),
+        torch.tensor([[[4, 5], [6, 7]], [[4, 5], [6, 7]]], dtype=dtype),
+    )
+    assert torch.equal(
+        feat_store_a.read(torch.tensor([1, 1])),
+        torch.tensor([[4, 5, 6], [4, 5, 6]], dtype=dtype),
+    )
+    assert torch.equal(
+        feat_store_b.read(torch.tensor([0])),
+        torch.tensor([[[1, 2], [3, 4]]], dtype=dtype),
+    )
+    # The cache should be full now for the large cache sizes, %100 hit expected.
+    total_miss = feat_store_a._feature.total_miss
+    feat_store_a.read(torch.tensor([0, 1]))
+    assert total_miss == feat_store_a._feature.total_miss
+    total_miss = feat_store_b._feature.total_miss
+    feat_store_b.read(torch.tensor([0, 1]))
+    assert total_miss == feat_store_b._feature.total_miss
+    assert feat_store_a._feature.miss_rate == feat_store_a.miss_rate
+
+    # Test get the size and count of the entire feature.
+    assert feat_store_a.size() == torch.Size([3])
+    assert feat_store_b.size() == torch.Size([2, 2])
+    assert feat_store_a.count() == a.size(0)
+    assert feat_store_b.count() == b.size(0)
+
+    # Test update the entire feature.
+    feat_store_a.update(torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype))
+    assert torch.equal(
+        feat_store_a.read(),
+        torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype),
+    )
+
+    # Test update with ids.
+    feat_store_a.update(
+        torch.tensor([[2, 0, 1]], dtype=dtype),
+        torch.tensor([0]),
+    )
+    assert torch.equal(
+        feat_store_a.read(),
+        torch.tensor([[2, 0, 1], [3, 5, 2]], dtype=dtype),
+    )
+
+    # Test with different dimensionality
+    feat_store_a.update(b)
+    assert torch.equal(feat_store_a.read(), b)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+def test_cpu_cached_feature_read_async(dtype):
+    a = torch.randint(0, 2, [1000, 13], dtype=dtype)
+
+    cache_size = 256 * a[:1].nbytes
+
+    feat_store = gb.cpu_cached_feature(gb.TorchBasedFeature(a), cache_size)
+
+    # Test read with ids.
+    ids1 = torch.tensor([0, 15, 71, 101])
+    ids2 = torch.tensor([71, 101, 202, 303])
+    for ids in [ids1, ids2]:
+        reader = feat_store.read_async(ids)
+        for _ in range(feat_store.read_async_num_stages(ids.device)):
+            values = next(reader)
+        assert torch.equal(values.wait(), a[ids])
+
+
+@unittest.skipIf(
+    not torch.ops.graphbolt.detect_io_uring(),
+    reason="DiskBasedFeature is not available on this system.",
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+def test_cpu_cached_disk_feature_read_async(dtype):
+    a = torch.randint(0, 2, [1000, 13], dtype=dtype)
+
+    cache_size = 256 * a[:1].nbytes
+
+    ids1 = torch.tensor([0, 15, 71, 101])
+    ids2 = torch.tensor([71, 101, 202, 303])
+
+    with tempfile.TemporaryDirectory() as test_dir:
+        path = to_on_disk_numpy(test_dir, "tensor", a)
+
+        feat_store = gb.cpu_cached_feature(
+            gb.DiskBasedFeature(path=path), cache_size
+        )
+
+        # Test read feature.
+        for ids in [ids1, ids2]:
+            reader = feat_store.read_async(ids)
+            for _ in range(feat_store.read_async_num_stages(ids.device)):
+                values = next(reader)
+            assert torch.equal(values.wait(), a[ids])
+
+        feat_store = None

tests/python/pytorch/graphbolt/impl/test_disk_based_feature_store.py

@@ -0,0 +1,192 @@
+import os
+import tempfile
+import unittest
+from functools import partial
+
+import backend as F
+
+import numpy as np
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def to_on_disk_numpy(test_dir, name, t):
+    path = os.path.join(test_dir, name + ".npy")
+    t = t.numpy()
+    np.save(path, t)
+    return path
+
+
+assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
+
+
+@unittest.skipIf(
+    not torch.ops.graphbolt.detect_io_uring(),
+    reason="DiskBasedFeature is not available on this system.",
+)
+def test_disk_based_feature():
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        c = torch.randn([4111, 47])
+        metadata = {"max_value": 3}
+        path_a = to_on_disk_numpy(test_dir, "a", a)
+        path_b = to_on_disk_numpy(test_dir, "b", b)
+        path_c = to_on_disk_numpy(test_dir, "c", c)
+
+        feature_a = gb.DiskBasedFeature(path=path_a, metadata=metadata)
+        feature_b = gb.DiskBasedFeature(path=path_b)
+        feature_c = gb.DiskBasedFeature(path=path_c)
+
+        # Read the entire feature.
+        assert_equal(feature_a.read(), torch.tensor([[1, 2, 3], [4, 5, 6]]))
+
+        assert_equal(
+            feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        )
+
+        # Test read the feature with ids.
+        assert_equal(
+            feature_a.read(torch.tensor([0])),
+            torch.tensor([[1, 2, 3]]),
+        )
+        assert_equal(
+            feature_b.read(torch.tensor([1])),
+            torch.tensor([[[4, 5], [6, 7]]]),
+        )
+
+        # Test reading into pin_memory
+        if F._default_context_str == "gpu":
+            res = feature_a.read(torch.tensor([0], pin_memory=True))
+            assert res.is_pinned()
+
+        # Test when the index tensor is large.
+        torch_based_feature_a = gb.TorchBasedFeature(a)
+        ind_a = torch.randint(low=0, high=a.size(0), size=(4111,))
+        assert_equal(
+            feature_a.read(ind_a),
+            torch_based_feature_a.read(ind_a),
+        )
+
+        # Test converting to torch_based_feature with read_into_memory()
+        torch_based_feature_b = feature_b.read_into_memory()
+        ind_b = torch.randint(low=0, high=b.size(0), size=(4111,))
+        assert_equal(
+            feature_b.read(ind_b),
+            torch_based_feature_b.read(ind_b),
+        )
+
+        # Test with larger stored feature tensor
+        ind_c = torch.randint(low=0, high=c.size(0), size=(4111,))
+        assert_equal(feature_c.read(ind_c), c[ind_c])
+
+        # Test get the size and count of the entire feature.
+        assert feature_a.size() == torch.Size([3])
+        assert feature_b.size() == torch.Size([2, 2])
+        assert feature_a.count() == a.size(0)
+        assert feature_b.count() == b.size(0)
+
+        # Test get metadata of the feature.
+        assert feature_a.metadata() == metadata
+        assert feature_b.metadata() == {}
+
+        with pytest.raises(IndexError):
+            feature_a.read(torch.tensor([0, 1, 2, 3]))
+
+        # Test loading a Fortran contiguous ndarray.
+        a_T = np.asfortranarray(a)
+        path_a_T = test_dir + "a_T.npy"
+        np.save(path_a_T, a_T)
+        with pytest.raises(
+            AssertionError,
+            match="DiskBasedFeature only supports C_CONTIGUOUS array.",
+        ):
+            gb.DiskBasedFeature(path=path_a_T, metadata=metadata)
+
+        # For windows, the file is locked by the numpy.load. We need to delete
+        # it before closing the temporary directory.
+        a = b = c = None
+        feature_a = feature_b = feature_c = None
+
+
+@unittest.skipIf(
+    not torch.ops.graphbolt.detect_io_uring(),
+    reason="DiskBasedFeature is not available on this system.",
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.float32,
+        torch.float64,
+        torch.int32,
+        torch.int64,
+        torch.int8,
+        torch.float16,
+        torch.complex128,
+    ],
+)
+@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize(
+    "shape", [(10, 20), (20, 10), (20, 25, 10), (137, 50, 30)]
+)
+@pytest.mark.parametrize("index", [[0], [1, 2, 3], [0, 6, 2, 8]])
+def test_more_disk_based_feature(dtype, idtype, shape, index):
+    if dtype == torch.complex128:
+        tensor = torch.complex(
+            torch.randint(0, 127, shape, dtype=torch.float64),
+            torch.randint(0, 127, shape, dtype=torch.float64),
+        )
+    else:
+        tensor = torch.randint(0, 127, shape, dtype=dtype)
+    test_tensor = tensor.clone()
+    idx = torch.tensor(index, dtype=idtype)
+
+    with tempfile.TemporaryDirectory() as test_dir:
+        path = to_on_disk_numpy(test_dir, "tensor", tensor)
+
+        feature = gb.DiskBasedFeature(path=path)
+
+        # Test read feature.
+        assert_equal(feature.read(idx), test_tensor[idx.long()])
+
+
+@unittest.skipIf(
+    not torch.ops.graphbolt.detect_io_uring(),
+    reason="DiskBasedFeature is not available on this system.",
+)
+def test_disk_based_feature_repr():
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        metadata = {"max_value": 3}
+
+        path_a = to_on_disk_numpy(test_dir, "a", a)
+        path_b = to_on_disk_numpy(test_dir, "b", b)
+
+        feature_a = gb.DiskBasedFeature(path=path_a, metadata=metadata)
+        feature_b = gb.DiskBasedFeature(path=path_b)
+
+        expected_str_feature_a = str(
+            "DiskBasedFeature(\n"
+            "    feature=tensor([[1, 2, 3],\n"
+            "                    [4, 5, 6]]),\n"
+            "    metadata={'max_value': 3},\n"
+            ")"
+        )
+        expected_str_feature_b = str(
+            "DiskBasedFeature(\n"
+            "    feature=tensor([[[1, 2],\n"
+            "                     [3, 4]],\n"
+            "\n"
+            "                    [[4, 5],\n"
+            "                     [6, 7]]]),\n"
+            "    metadata={},\n"
+            ")"
+        )
+        assert str(feature_a) == expected_str_feature_a
+        assert str(feature_b) == expected_str_feature_b
+        a = b = metadata = None
+        feature_a = feature_b = None
+        expected_str_feature_a = expected_str_feature_b = None

tests/python/pytorch/graphbolt/impl/test_feature_cache.py

@@ -0,0 +1,172 @@
+import backend as F
+
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def _test_query_and_replace(policy1, policy2, keys, offset):
+    # Testing query_and_replace equivalence to query and then replace.
+    (
+        _,
+        index,
+        pointers,
+        missing_keys,
+        found_offsets,
+        missing_offsets,
+    ) = policy1.query_and_replace(keys, offset)
+    found_cnt = keys.size(0) - missing_keys.size(0)
+    found_pointers = pointers[:found_cnt]
+    policy1.reading_completed(found_pointers, found_offsets)
+    missing_pointers = pointers[found_cnt:]
+    policy1.writing_completed(missing_pointers, missing_offsets)
+
+    (
+        _,
+        index2,
+        missing_keys2,
+        found_pointers2,
+        found_offsets2,
+        missing_offsets2,
+    ) = policy2.query(keys + offset, 0)
+    policy2.reading_completed(found_pointers2, found_offsets2)
+    (_, missing_pointers2, missing_offsets2) = policy2.replace(
+        missing_keys2, missing_offsets2, 0
+    )
+    policy2.writing_completed(missing_pointers2, missing_offsets2)
+
+    assert torch.equal(index, index2)
+    assert torch.equal(missing_keys, missing_keys2 - offset)
+
+
+@pytest.mark.parametrize("offsets", [False, True])
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("feature_size", [2, 16])
+@pytest.mark.parametrize("num_parts", [1, 2, None])
+@pytest.mark.parametrize("policy", ["s3-fifo", "sieve", "lru", "clock"])
+@pytest.mark.parametrize("offset", [0, 1111111])
+def test_feature_cache(offsets, dtype, feature_size, num_parts, policy, offset):
+    cache_size = 32 * (
+        torch.get_num_threads() if num_parts is None else num_parts
+    )
+    a = torch.randint(0, 2, [1024, feature_size], dtype=dtype)
+    cache = gb.impl.CPUFeatureCache(
+        (cache_size,) + a.shape[1:], a.dtype, policy, num_parts
+    )
+    cache2 = gb.impl.CPUFeatureCache(
+        (cache_size,) + a.shape[1:], a.dtype, policy, num_parts
+    )
+    policy1 = gb.impl.CPUFeatureCache(
+        (cache_size,) + a.shape[1:], a.dtype, policy, num_parts
+    )._policy
+    policy2 = gb.impl.CPUFeatureCache(
+        (cache_size,) + a.shape[1:], a.dtype, policy, num_parts
+    )._policy
+    reader_fn = lambda keys: a[keys]
+
+    keys = torch.tensor([0, 1])
+    values, missing_index, missing_keys, missing_offsets = cache.query(
+        keys, offset
+    )
+    if not offsets:
+        missing_offsets = None
+    assert torch.equal(
+        missing_keys.flip([0]) if num_parts == 1 else missing_keys.sort()[0],
+        keys,
+    )
+
+    missing_values = a[missing_keys]
+    cache.replace(missing_keys, missing_values, missing_offsets, offset)
+    values[missing_index] = missing_values
+    assert torch.equal(values, a[keys])
+    assert torch.equal(
+        cache2.query_and_replace(keys, reader_fn, offset), a[keys]
+    )
+
+    _test_query_and_replace(policy1, policy2, keys, offset)
+
+    pin_memory = F._default_context_str == "gpu"
+
+    keys = torch.arange(1, 33, pin_memory=pin_memory)
+    values, missing_index, missing_keys, missing_offsets = cache.query(
+        keys, offset
+    )
+    if not offsets:
+        missing_offsets = None
+    assert torch.equal(
+        missing_keys.flip([0]) if num_parts == 1 else missing_keys.sort()[0],
+        torch.arange(2, 33),
+    )
+    assert not pin_memory or values.is_pinned()
+
+    missing_values = a[missing_keys]
+    cache.replace(missing_keys, missing_values, missing_offsets, offset)
+    values[missing_index] = missing_values
+    assert torch.equal(values, a[keys])
+    assert torch.equal(
+        cache2.query_and_replace(keys, reader_fn, offset), a[keys]
+    )
+
+    _test_query_and_replace(policy1, policy2, keys, offset)
+
+    values, missing_index, missing_keys, missing_offsets = cache.query(
+        keys, offset
+    )
+    if not offsets:
+        missing_offsets = None
+    assert torch.equal(missing_keys.flip([0]), torch.tensor([]))
+
+    missing_values = a[missing_keys]
+    cache.replace(missing_keys, missing_values, missing_offsets, offset)
+    values[missing_index] = missing_values
+    assert torch.equal(values, a[keys])
+    assert torch.equal(
+        cache2.query_and_replace(keys, reader_fn, offset), a[keys]
+    )
+
+    _test_query_and_replace(policy1, policy2, keys, offset)
+
+    values, missing_index, missing_keys, missing_offsets = cache.query(
+        keys, offset
+    )
+    if not offsets:
+        missing_offsets = None
+    assert torch.equal(missing_keys.flip([0]), torch.tensor([]))
+
+    missing_values = a[missing_keys]
+    cache.replace(missing_keys, missing_values, missing_offsets, offset)
+    values[missing_index] = missing_values
+    assert torch.equal(values, a[keys])
+    assert torch.equal(
+        cache2.query_and_replace(keys, reader_fn, offset), a[keys]
+    )
+
+    _test_query_and_replace(policy1, policy2, keys, offset)
+
+    assert cache.miss_rate == cache2.miss_rate
+
+    raw_feature_cache = torch.ops.graphbolt.feature_cache(
+        (cache_size,) + a.shape[1:], a.dtype, pin_memory
+    )
+    idx = torch.tensor([0, 1, 2])
+    raw_feature_cache.replace(idx, a[idx])
+    val = raw_feature_cache.index_select(idx)
+    assert torch.equal(val, a[idx])
+    if pin_memory:
+        val = raw_feature_cache.index_select(idx.to(F.ctx()))
+        assert torch.equal(val, a[idx].to(F.ctx()))

tests/python/pytorch/graphbolt/impl/test_gpu_cached_feature.py

@@ -0,0 +1,219 @@
+import os
+import tempfile
+import unittest
+
+import backend as F
+
+import numpy as np
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def to_on_disk_numpy(test_dir, name, t):
+    path = os.path.join(test_dir, name + ".npy")
+    np.save(path, t.cpu().numpy())
+    return path
+
+
+def _skip_condition_cached_feature():
+    return (F._default_context_str != "gpu") or (
+        torch.cuda.get_device_capability()[0] < 7
+    )
+
+
+def _reason_to_skip_cached_feature():
+    if F._default_context_str != "gpu":
+        return "GPUCachedFeature tests are available only when testing the GPU backend."
+
+    return "GPUCachedFeature requires a Volta or later generation NVIDIA GPU."
+
+
+@unittest.skipIf(
+    _skip_condition_cached_feature(),
+    reason=_reason_to_skip_cached_feature(),
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("cache_size_a", [1, 1024])
+@pytest.mark.parametrize("cache_size_b", [1, 1024])
+def test_gpu_cached_feature(dtype, cache_size_a, cache_size_b):
+    a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype, pin_memory=True)
+    b = torch.tensor(
+        [[[1, 2], [3, 4]], [[4, 5], [6, 7]]], dtype=dtype, pin_memory=True
+    )
+
+    cache_size_a *= a[:1].element_size() * a[:1].numel()
+    cache_size_b *= b[:1].element_size() * b[:1].numel()
+
+    feat_store_a = gb.gpu_cached_feature(gb.TorchBasedFeature(a), cache_size_a)
+    feat_store_b = gb.gpu_cached_feature(gb.TorchBasedFeature(b), cache_size_b)
+
+    # Test read the entire feature.
+    assert torch.equal(feat_store_a.read(), a.to("cuda"))
+    assert torch.equal(feat_store_b.read(), b.to("cuda"))
+
+    # Test read with ids.
+    assert torch.equal(
+        feat_store_a.read(torch.tensor([0]).to("cuda")),
+        torch.tensor([[1, 2, 3]], dtype=dtype).to("cuda"),
+    )
+    assert torch.equal(
+        feat_store_b.read(torch.tensor([1, 1]).to("cuda")),
+        torch.tensor([[[4, 5], [6, 7]], [[4, 5], [6, 7]]], dtype=dtype).to(
+            "cuda"
+        ),
+    )
+    assert torch.equal(
+        feat_store_a.read(torch.tensor([1, 1]).to("cuda")),
+        torch.tensor([[4, 5, 6], [4, 5, 6]], dtype=dtype).to("cuda"),
+    )
+    assert torch.equal(
+        feat_store_b.read(torch.tensor([0]).to("cuda")),
+        torch.tensor([[[1, 2], [3, 4]]], dtype=dtype).to("cuda"),
+    )
+    # The cache should be full now for the large cache sizes, %100 hit expected.
+    if cache_size_a >= 1024:
+        total_miss = feat_store_a._feature.total_miss
+        feat_store_a.read(torch.tensor([0, 1]).to("cuda"))
+        assert total_miss == feat_store_a._feature.total_miss
+    if cache_size_b >= 1024:
+        total_miss = feat_store_b._feature.total_miss
+        feat_store_b.read(torch.tensor([0, 1]).to("cuda"))
+        assert total_miss == feat_store_b._feature.total_miss
+    assert feat_store_a._feature.miss_rate == feat_store_a.miss_rate
+
+    # Test get the size and count of the entire feature.
+    assert feat_store_a.size() == torch.Size([3])
+    assert feat_store_b.size() == torch.Size([2, 2])
+    assert feat_store_a.count() == a.size(0)
+    assert feat_store_b.count() == b.size(0)
+
+    # Test update the entire feature.
+    feat_store_a.update(
+        torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype).to("cuda")
+    )
+    assert torch.equal(
+        feat_store_a.read(),
+        torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype).to("cuda"),
+    )
+
+    # Test update with ids.
+    feat_store_a.update(
+        torch.tensor([[2, 0, 1]], dtype=dtype).to("cuda"),
+        torch.tensor([0]).to("cuda"),
+    )
+    assert torch.equal(
+        feat_store_a.read(),
+        torch.tensor([[2, 0, 1], [3, 5, 2]], dtype=dtype).to("cuda"),
+    )
+
+    # Test with different dimensionality
+    feat_store_a.update(b)
+    assert torch.equal(feat_store_a.read(), b.to("cuda"))
+
+
+@unittest.skipIf(
+    _skip_condition_cached_feature(),
+    reason=_reason_to_skip_cached_feature(),
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("pin_memory", [False, True])
+def test_gpu_cached_feature_read_async(dtype, pin_memory):
+    a = torch.randint(0, 2, [1000, 13], dtype=dtype, pin_memory=pin_memory)
+    a_cuda = a.to(F.ctx())
+
+    cache_size = 256 * a[:1].nbytes
+
+    feat_store = gb.gpu_cached_feature(gb.TorchBasedFeature(a), cache_size)
+
+    # Test read with ids.
+    ids1 = torch.tensor([0, 15, 71, 101], device=F.ctx())
+    ids2 = torch.tensor([71, 101, 202, 303], device=F.ctx())
+    for ids in [ids1, ids2]:
+        reader = feat_store.read_async(ids)
+        for _ in range(feat_store.read_async_num_stages(ids.device)):
+            values = next(reader)
+        assert torch.equal(values.wait(), a_cuda[ids])
+
+
+@unittest.skipIf(
+    _skip_condition_cached_feature(),
+    reason=_reason_to_skip_cached_feature(),
+)
+@unittest.skipIf(
+    not torch.ops.graphbolt.detect_io_uring(),
+    reason="DiskBasedFeature is not available on this system.",
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+def test_gpu_cached_nested_feature_async(dtype):
+    a = torch.randint(0, 2, [1000, 13], dtype=dtype, device=F.ctx())
+
+    cache_size = 256 * a[:1].nbytes
+
+    ids1 = torch.tensor([0, 15, 71, 101], device=F.ctx())
+    ids2 = torch.tensor([71, 101, 202, 303], device=F.ctx())
+
+    with tempfile.TemporaryDirectory() as test_dir:
+        path = to_on_disk_numpy(test_dir, "tensor", a)
+
+        disk_store = gb.DiskBasedFeature(path=path)
+        feat_store1 = gb.gpu_cached_feature(disk_store, cache_size)
+        feat_store2 = gb.gpu_cached_feature(
+            gb.cpu_cached_feature(disk_store, cache_size * 2), cache_size
+        )
+        feat_store3 = gb.gpu_cached_feature(
+            gb.cpu_cached_feature(disk_store, cache_size * 2, pin_memory=True),
+            cache_size,
+        )
+
+        # Test read feature.
+        for feat_store in [feat_store1, feat_store2, feat_store3]:
+            for ids in [ids1, ids2]:
+                reader = feat_store.read_async(ids)
+                for _ in range(feat_store.read_async_num_stages(ids.device)):
+                    values = next(reader)
+                assert torch.equal(values.wait(), a[ids])
+
+        feat_store1 = feat_store2 = feat_store3 = disk_store = None

tests/python/pytorch/graphbolt/impl/test_gpu_graph_cache.py

@@ -0,0 +1,82 @@
+import unittest
+
+import backend as F
+
+import dgl.graphbolt as gb
+
+import pytest
+import torch
+
+
+@unittest.skipIf(
+    F._default_context_str != "gpu"
+    or torch.cuda.get_device_capability()[0] < 7,
+    reason="GPUCachedFeature tests are available only when testing the GPU backend."
+    if F._default_context_str != "gpu"
+    else "GPUCachedFeature requires a Volta or later generation NVIDIA GPU.",
+)
+@pytest.mark.parametrize(
+    "indptr_dtype",
+    [
+        torch.int32,
+        torch.int64,
+    ],
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("cache_size", [4, 9, 11])
+@pytest.mark.parametrize("with_edge_ids", [True, False])
+def test_gpu_graph_cache(indptr_dtype, dtype, cache_size, with_edge_ids):
+    indices_dtype = torch.int32
+    indptr = torch.tensor([0, 3, 6, 10], dtype=indptr_dtype, pin_memory=True)
+    indices = torch.arange(0, indptr[-1], dtype=indices_dtype, pin_memory=True)
+    probs_or_mask = indices.to(dtype).pin_memory()
+    edge_tensors = [indices, probs_or_mask]
+
+    g = gb.GPUGraphCache(
+        cache_size,
+        2,
+        indptr.dtype,
+        [e.dtype for e in edge_tensors],
+        not with_edge_ids,
+    )
+
+    for i in range(10):
+        keys = (
+            torch.arange(2, dtype=indices_dtype, device=F.ctx()) + i * 2
+        ) % (indptr.size(0) - 1)
+        missing_keys, replace = g.query(keys)
+        (
+            missing_indptr,
+            missing_edge_tensors,
+        ) = torch.ops.graphbolt.index_select_csc_batched(
+            indptr, edge_tensors, missing_keys, with_edge_ids, None
+        )
+        output_indptr, output_edge_tensors = replace(
+            missing_indptr, missing_edge_tensors
+        )
+
+        (
+            reference_indptr,
+            reference_edge_tensors,
+        ) = torch.ops.graphbolt.index_select_csc_batched(
+            indptr, edge_tensors, keys, with_edge_ids, None
+        )
+
+        assert torch.equal(output_indptr, reference_indptr)
+        assert len(output_edge_tensors) == len(reference_edge_tensors)
+        for e, ref in zip(output_edge_tensors, reference_edge_tensors):
+            assert torch.equal(e, ref)

tests/python/pytorch/graphbolt/impl/test_hetero_cached_feature.py

@@ -0,0 +1,42 @@
+import backend as F
+
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+@pytest.mark.parametrize(
+    "cached_feature_type", [gb.cpu_cached_feature, gb.gpu_cached_feature]
+)
+def test_hetero_cached_feature(cached_feature_type):
+    if cached_feature_type == gb.gpu_cached_feature and (
+        F._default_context_str != "gpu"
+        or torch.cuda.get_device_capability()[0] < 7
+    ):
+        pytest.skip(
+            "GPUCachedFeature tests are available only when testing the GPU backend."
+            if F._default_context_str != "gpu"
+            else "GPUCachedFeature requires a Volta or later generation NVIDIA GPU."
+        )
+    device = F.ctx() if cached_feature_type == gb.gpu_cached_feature else None
+    pin_memory = cached_feature_type == gb.gpu_cached_feature
+
+    a = {
+        ("node", str(i), "feat"): gb.TorchBasedFeature(
+            torch.randn([(i + 1) * 10, 5], pin_memory=pin_memory)
+        )
+        for i in range(75)
+    }
+    cached_a = cached_feature_type(a, 2**18)
+
+    for i in range(1024):
+        etype = i % len(a)
+        ids = torch.randint(
+            0, (etype + 1) * 10 - 1, ((etype + 1) * 4,), device=device
+        )
+        feature_key = ("node", str(etype), "feat")
+        ref = a[feature_key].read(ids)
+        val = cached_a[feature_key].read(ids)
+        torch.testing.assert_close(ref, val, rtol=0, atol=0)
+    assert cached_a[feature_key].miss_rate < 0.69

tests/python/pytorch/graphbolt/impl/test_in_subgraph_sampler.py

@@ -0,0 +1,287 @@
+import unittest
+
+import backend as F
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+from .. import gb_test_utils
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Tests for pinned memory are only meaningful on GPU.",
+)
+@pytest.mark.parametrize(
+    "indptr_dtype",
+    [torch.int32, torch.int64],
+)
+@pytest.mark.parametrize(
+    "indices_dtype",
+    [
+        torch.int8,
+        torch.uint8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("is_pinned", [False, True])
+@pytest.mark.parametrize("with_edge_ids", [False, True])
+@pytest.mark.parametrize("output_size", [None, True])
+def test_index_select_csc(
+    indptr_dtype, indices_dtype, idtype, is_pinned, with_edge_ids, output_size
+):
+    """Original graph in COO:
+    1   0   1   0   1   0
+    1   0   0   1   0   1
+    0   1   0   1   0   0
+    0   1   0   0   1   0
+    1   0   0   0   0   1
+    0   0   1   0   1   0
+    """
+    indptr = torch.tensor([0, 3, 5, 7, 9, 12, 14], dtype=indptr_dtype)
+    indices = torch.tensor(
+        [0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4], dtype=indices_dtype
+    )
+    index = torch.tensor([0, 5, 3], dtype=idtype)
+
+    cpu_indptr, cpu_indices = torch.ops.graphbolt.index_select_csc(
+        indptr, indices, index, None
+    )
+    if is_pinned:
+        indptr = indptr.pin_memory()
+        indices = indices.pin_memory()
+    else:
+        indptr = indptr.cuda()
+        indices = indices.cuda()
+    index = index.cuda()
+    edge_ids = torch.tensor(
+        [0, 1, 2, 12, 13, 7, 8], dtype=indptr_dtype, device=index.device
+    )
+
+    if output_size:
+        output_size = len(cpu_indices)
+
+    gpu_indptr, gpu_indices = torch.ops.graphbolt.index_select_csc(
+        indptr, indices, index, output_size
+    )
+    assert not cpu_indptr.is_cuda
+    assert not cpu_indices.is_cuda
+
+    assert gpu_indptr.is_cuda
+    assert gpu_indices.is_cuda
+
+    assert torch.equal(cpu_indptr, gpu_indptr.cpu())
+    assert torch.equal(cpu_indices, gpu_indices.cpu())
+
+    for output_size_selection in [None, output_size]:
+        indices_list = [
+            indices,
+            indices.int().pin_memory() if is_pinned else indices.int(),
+        ]
+        (
+            gpu_indptr2,
+            gpu_indices_list,
+        ) = torch.ops.graphbolt.index_select_csc_batched(
+            indptr, indices_list, index, with_edge_ids, output_size_selection
+        )
+
+        assert torch.equal(gpu_indptr, gpu_indptr2)
+        assert torch.equal(gpu_indices_list[0], gpu_indices)
+        assert torch.equal(gpu_indices_list[1], gpu_indices.int())
+        if with_edge_ids:
+            assert torch.equal(gpu_indices_list[2], edge_ids)
+
+
+def test_InSubgraphSampler_homo():
+    """Original graph in COO:
+    1   0   1   0   1   0
+    1   0   0   1   0   1
+    0   1   0   1   0   0
+    0   1   0   0   1   0
+    1   0   0   0   0   1
+    0   0   1   0   1   0
+    """
+    indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
+    indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
+    graph = gb.fused_csc_sampling_graph(indptr, indices).to(F.ctx())
+
+    seed_nodes = torch.LongTensor([0, 5, 3])
+    item_set = gb.ItemSet(seed_nodes, names="seeds")
+    batch_size = 1
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+
+    in_subgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
+
+    it = iter(in_subgraph_sampler)
+
+    def original_indices(minibatch):
+        sampled_subgraph = minibatch.sampled_subgraphs[0]
+        _indices = sampled_subgraph.original_row_node_ids[
+            sampled_subgraph.sampled_csc.indices
+        ]
+        return _indices
+
+    mn = next(it)
+    assert torch.equal(mn.seeds, torch.LongTensor([0]).to(F.ctx()))
+    assert torch.equal(
+        mn.sampled_subgraphs[0].sampled_csc.indptr,
+        torch.tensor([0, 3]).to(F.ctx()),
+    )
+
+    mn = next(it)
+    assert torch.equal(mn.seeds, torch.LongTensor([5]).to(F.ctx()))
+    assert torch.equal(
+        mn.sampled_subgraphs[0].sampled_csc.indptr,
+        torch.tensor([0, 2]).to(F.ctx()),
+    )
+    assert torch.equal(original_indices(mn), torch.tensor([1, 4]).to(F.ctx()))
+
+    mn = next(it)
+    assert torch.equal(mn.seeds, torch.LongTensor([3]).to(F.ctx()))
+    assert torch.equal(
+        mn.sampled_subgraphs[0].sampled_csc.indptr,
+        torch.tensor([0, 2]).to(F.ctx()),
+    )
+    assert torch.equal(original_indices(mn), torch.tensor([1, 2]).to(F.ctx()))
+
+
+def test_InSubgraphSampler_hetero():
+    """Original graph in COO:
+    1   0   1   0   1   0
+    1   0   0   1   0   1
+    0   1   0   1   0   0
+    0   1   0   0   1   0
+    1   0   0   0   0   1
+    0   0   1   0   1   0
+    node_type_0: [0, 1, 2]
+    node_type_1: [3, 4, 5]
+    edge_type_0: node_type_0 -> node_type_0
+    edge_type_1: node_type_0 -> node_type_1
+    edge_type_2: node_type_1 -> node_type_0
+    edge_type_3: node_type_1 -> node_type_1
+    """
+    ntypes = {
+        "N0": 0,
+        "N1": 1,
+    }
+    etypes = {
+        "N0:R0:N0": 0,
+        "N0:R1:N1": 1,
+        "N1:R2:N0": 2,
+        "N1:R3:N1": 3,
+    }
+    indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
+    indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
+    node_type_offset = torch.LongTensor([0, 3, 6])
+    type_per_edge = torch.LongTensor([0, 0, 2, 0, 2, 0, 2, 1, 1, 1, 3, 3, 1, 3])
+    graph = gb.fused_csc_sampling_graph(
+        csc_indptr=indptr,
+        indices=indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+    ).to(F.ctx())
+
+    item_set = gb.HeteroItemSet(
+        {
+            "N0": gb.ItemSet(torch.LongTensor([1, 0, 2]), names="seeds"),
+            "N1": gb.ItemSet(torch.LongTensor([0, 2, 1]), names="seeds"),
+        }
+    )
+    batch_size = 2
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+
+    in_subgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
+
+    it = iter(in_subgraph_sampler)
+
+    mn = next(it)
+    assert torch.equal(mn.seeds["N0"], torch.LongTensor([1, 0]).to(F.ctx()))
+    expected_sampled_csc = {
+        "N0:R0:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1, 3]),
+            indices=torch.LongTensor([2, 1, 0]),
+        ),
+        "N0:R1:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
+        ),
+        "N1:R2:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([0, 1])
+        ),
+        "N1:R3:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
+        ),
+    }
+    for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
+        assert torch.equal(
+            pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
+        )
+        assert torch.equal(
+            pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
+        )
+
+    mn = next(it)
+    assert mn.seeds == {
+        "N0": torch.LongTensor([2]).to(F.ctx()),
+        "N1": torch.LongTensor([0]).to(F.ctx()),
+    }
+    expected_sampled_csc = {
+        "N0:R0:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1]), indices=torch.LongTensor([1])
+        ),
+        "N0:R1:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 2]), indices=torch.LongTensor([2, 0])
+        ),
+        "N1:R2:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1]), indices=torch.LongTensor([1])
+        ),
+        "N1:R3:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 0]), indices=torch.LongTensor([])
+        ),
+    }
+    for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
+        assert torch.equal(
+            pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
+        )
+        assert torch.equal(
+            pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
+        )
+
+    mn = next(it)
+    assert torch.equal(mn.seeds["N1"], torch.LongTensor([2, 1]).to(F.ctx()))
+    expected_sampled_csc = {
+        "N0:R0:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
+        ),
+        "N0:R1:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([0, 1])
+        ),
+        "N1:R2:N0": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
+        ),
+        "N1:R3:N1": gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1, 3]),
+            indices=torch.LongTensor([1, 2, 0]),
+        ),
+    }
+    if graph.csc_indptr.is_cuda and torch.cuda.get_device_capability()[0] < 7:
+        expected_sampled_csc["N0:R1:N1"] = gb.CSCFormatBase(
+            indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([1, 0])
+        )
+    for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
+        assert torch.equal(
+            pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
+        )
+        assert torch.equal(
+            pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
+        )

tests/python/pytorch/graphbolt/impl/test_legacy_dataset.py

@@ -0,0 +1,36 @@
+import dgl.graphbolt as gb
+import pytest
+import torch
+from dgl import AddSelfLoop
+from dgl.data import AsNodePredDataset, CoraGraphDataset
+
+
+def test_LegacyDataset_homo_node_pred():
+    cora = CoraGraphDataset(transform=AddSelfLoop())
+    dataset = gb.LegacyDataset(cora)
+
+    # Check tasks.
+    assert len(dataset.tasks) == 1
+    task = dataset.tasks[0]
+    assert task.train_set.names == ("seeds", "labels")
+    assert len(task.train_set) == 140
+    assert task.validation_set.names == ("seeds", "labels")
+    assert len(task.validation_set) == 500
+    assert task.test_set.names == ("seeds", "labels")
+    assert len(task.test_set) == 1000
+    assert task.metadata["num_classes"] == 7
+
+    num_nodes = 2708
+    assert dataset.graph.num_nodes == num_nodes
+    assert len(dataset.all_nodes_set) == num_nodes
+    assert dataset.feature.size("node", None, "feat") == torch.Size([1433])
+    assert (
+        dataset.feature.read(
+            "node", None, "feat", torch.tensor([num_nodes - 1])
+        ).size(dim=0)
+        == 1
+    )
+    # Out of bound indexing results in segmentation fault instead of exception
+    # in CI. This may be related to docker env. Skip it for now.
+    # with pytest.raises(IndexError):
+    #    dataset.feature.read("node", None, "feat", torch.Tensor([num_nodes]))

tests/python/pytorch/graphbolt/impl/test_negative_sampler.py

@@ -0,0 +1,257 @@
+import re
+
+import backend as F
+
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+from .. import gb_test_utils
+
+
+def test_NegativeSampler_invoke():
+    # Instantiate graph and required datapipes.
+    num_seeds = 30
+    item_set = gb.ItemSet(
+        torch.arange(0, 2 * num_seeds).reshape(-1, 2), names="seeds"
+    )
+    batch_size = 10
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    negative_ratio = 2
+
+    # Invoke NegativeSampler via class constructor.
+    negative_sampler = gb.NegativeSampler(
+        item_sampler,
+        negative_ratio,
+    )
+    with pytest.raises(NotImplementedError):
+        next(iter(negative_sampler))
+
+    # Invoke NegativeSampler via functional form.
+    negative_sampler = item_sampler.sample_negative(
+        negative_ratio,
+    )
+    with pytest.raises(NotImplementedError):
+        next(iter(negative_sampler))
+
+
+def test_UniformNegativeSampler_invoke():
+    # Instantiate graph and required datapipes.
+    graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
+        F.ctx()
+    )
+    num_seeds = 30
+    item_set = gb.ItemSet(
+        torch.arange(0, 2 * num_seeds).reshape(-1, 2), names="seeds"
+    )
+    batch_size = 10
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    negative_ratio = 2
+
+    def _verify(negative_sampler):
+        for data in negative_sampler:
+            # Assertation
+            seeds_len = batch_size + batch_size * negative_ratio
+            assert data.seeds.size(0) == seeds_len
+            assert data.labels.size(0) == seeds_len
+            assert data.indexes.size(0) == seeds_len
+
+    # Invoke UniformNegativeSampler via class constructor.
+    negative_sampler = gb.UniformNegativeSampler(
+        item_sampler,
+        graph,
+        negative_ratio,
+    )
+    _verify(negative_sampler)
+
+    # Invoke UniformNegativeSampler via functional form.
+    negative_sampler = item_sampler.sample_uniform_negative(
+        graph,
+        negative_ratio,
+    )
+    _verify(negative_sampler)
+
+
+@pytest.mark.parametrize("negative_ratio", [1, 5, 10, 20])
+def test_Uniform_NegativeSampler(negative_ratio):
+    # Construct FusedCSCSamplingGraph.
+    graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
+        F.ctx()
+    )
+    num_seeds = 30
+    item_set = gb.ItemSet(
+        torch.arange(0, num_seeds * 2).reshape(-1, 2), names="seeds"
+    )
+    batch_size = 10
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    # Construct NegativeSampler.
+    negative_sampler = gb.UniformNegativeSampler(
+        item_sampler,
+        graph,
+        negative_ratio,
+    )
+    # Perform Negative sampling.
+    for data in negative_sampler:
+        seeds_len = batch_size + batch_size * negative_ratio
+        # Assertation
+        assert data.seeds.size(0) == seeds_len
+        assert data.labels.size(0) == seeds_len
+        assert data.indexes.size(0) == seeds_len
+        # Check negative seeds value.
+        pos_src = data.seeds[:batch_size, 0]
+        neg_src = data.seeds[batch_size:, 0]
+        assert torch.equal(pos_src.repeat_interleave(negative_ratio), neg_src)
+        # Check labels.
+        assert torch.equal(
+            data.labels[:batch_size], torch.ones(batch_size).to(F.ctx())
+        )
+        assert torch.equal(
+            data.labels[batch_size:],
+            torch.zeros(batch_size * negative_ratio).to(F.ctx()),
+        )
+        # Check indexes.
+        pos_indexes = torch.arange(0, batch_size).to(F.ctx())
+        neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
+        expected_indexes = torch.cat((pos_indexes, neg_indexes))
+        assert torch.equal(data.indexes, expected_indexes)
+
+
+def test_Uniform_NegativeSampler_error_shape():
+    # 1. seeds with shape N*3.
+    # Construct FusedCSCSamplingGraph.
+    graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
+        F.ctx()
+    )
+    num_seeds = 30
+    item_set = gb.ItemSet(
+        torch.arange(0, num_seeds * 3).reshape(-1, 3), names="seeds"
+    )
+    batch_size = 10
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    negative_ratio = 2
+    # Construct NegativeSampler.
+    negative_sampler = gb.UniformNegativeSampler(
+        item_sampler,
+        graph,
+        negative_ratio,
+    )
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Only tensor with shape N*2 is "
+            + "supported for negative sampling, but got torch.Size([10, 3])."
+        ),
+    ):
+        next(iter(negative_sampler))
+
+    # 2. seeds with shape N*2*1.
+    # Construct FusedCSCSamplingGraph.
+    item_set = gb.ItemSet(
+        torch.arange(0, num_seeds * 2).reshape(-1, 2, 1), names="seeds"
+    )
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    # Construct NegativeSampler.
+    negative_sampler = gb.UniformNegativeSampler(
+        item_sampler,
+        graph,
+        negative_ratio,
+    )
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Only tensor with shape N*2 is "
+            + "supported for negative sampling, but got torch.Size([10, 2, 1])."
+        ),
+    ):
+        next(iter(negative_sampler))
+
+    # 3. seeds with shape N.
+    # Construct FusedCSCSamplingGraph.
+    item_set = gb.ItemSet(torch.arange(0, num_seeds), names="seeds")
+    item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    # Construct NegativeSampler.
+    negative_sampler = gb.UniformNegativeSampler(
+        item_sampler,
+        graph,
+        negative_ratio,
+    )
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Only tensor with shape N*2 is "
+            + "supported for negative sampling, but got torch.Size([10])."
+        ),
+    ):
+        next(iter(negative_sampler))
+
+
+def get_hetero_graph():
+    # COO graph:
+    # [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    # [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
+    # [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
+    # num_nodes = 5, num_n1 = 2, num_n2 = 3
+    ntypes = {"n1": 0, "n2": 1}
+    etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
+    indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
+    indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
+    type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
+    node_type_offset = torch.LongTensor([0, 2, 5])
+    return gb.fused_csc_sampling_graph(
+        indptr,
+        indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+    )
+
+
+def test_NegativeSampler_Hetero_Data():
+    graph = get_hetero_graph().to(F.ctx())
+    itemset = gb.HeteroItemSet(
+        {
+            "n1:e1:n2": gb.ItemSet(
+                torch.LongTensor([[0, 0, 1, 1], [0, 2, 0, 1]]).T,
+                names="seeds",
+            ),
+            "n2:e2:n1": gb.ItemSet(
+                torch.LongTensor([[0, 0, 1, 1, 2, 2], [0, 1, 1, 0, 0, 1]]).T,
+                names="seeds",
+            ),
+        }
+    )
+    batch_size = 2
+    negative_ratio = 1
+    item_sampler = gb.ItemSampler(itemset, batch_size=batch_size).copy_to(
+        F.ctx()
+    )
+    negative_dp = gb.UniformNegativeSampler(item_sampler, graph, negative_ratio)
+    assert len(list(negative_dp)) == 5
+    # Perform negative sampling.
+    expected_neg_src = [
+        {"n1:e1:n2": torch.tensor([0, 0])},
+        {"n1:e1:n2": torch.tensor([1, 1])},
+        {"n2:e2:n1": torch.tensor([0, 0])},
+        {"n2:e2:n1": torch.tensor([1, 1])},
+        {"n2:e2:n1": torch.tensor([2, 2])},
+    ]
+    for i, data in enumerate(negative_dp):
+        # Check negative seeds value.
+        for etype, seeds_data in data.seeds.items():
+            neg_src = seeds_data[batch_size:, 0]
+            neg_dst = seeds_data[batch_size:, 1]
+            assert torch.equal(expected_neg_src[i][etype].to(F.ctx()), neg_src)
+            assert (neg_dst < 3).all(), neg_dst

tests/python/pytorch/graphbolt/impl/test_neighbor_sampler.py

@@ -0,0 +1,161 @@
+import unittest
+from functools import partial
+
+import backend as F
+
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+
+def get_hetero_graph(include_original_edge_ids):
+    # COO graph:
+    # [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    # [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
+    # [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
+    # num_nodes = 5, num_n1 = 2, num_n2 = 3
+    ntypes = {"n1": 0, "n2": 1, "n3": 2}
+    etypes = {"n2:e1:n3": 0, "n3:e2:n2": 1}
+    indptr = torch.LongTensor([0, 0, 2, 4, 6, 8, 10])
+    indices = torch.LongTensor([3, 5, 3, 4, 1, 2, 2, 1, 1, 2])
+    type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
+    edge_attributes = {
+        "weight": torch.FloatTensor(
+            [2.5, 0, 8.4, 0, 0.4, 1.2, 2.5, 0, 8.4, 0.5]
+        ),
+        "mask": torch.BoolTensor([1, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
+    }
+    if include_original_edge_ids:
+        edge_attributes[gb.ORIGINAL_EDGE_ID] = (
+            torch.arange(indices.size(0), 0, -1) - 1
+        )
+    node_type_offset = torch.LongTensor([0, 1, 3, 6])
+    return gb.fused_csc_sampling_graph(
+        indptr,
+        indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+        edge_attributes=edge_attributes,
+    )
+
+
+@unittest.skipIf(F._default_context_str != "gpu", reason="Enabled only on GPU.")
+@pytest.mark.parametrize("hetero", [False, True])
+@pytest.mark.parametrize("prob_name", [None, "weight", "mask"])
+@pytest.mark.parametrize("sorted", [False, True])
+@pytest.mark.parametrize("num_cached_edges", [0, 10])
+@pytest.mark.parametrize("is_pinned", [False, True])
+@pytest.mark.parametrize("has_orig_edge_ids", [False, True])
+def test_NeighborSampler_GraphFetch(
+    hetero, prob_name, sorted, num_cached_edges, is_pinned, has_orig_edge_ids
+):
+    if sorted:
+        items = torch.arange(3)
+    else:
+        items = torch.tensor([2, 0, 1])
+    names = "seeds"
+    itemset = gb.ItemSet(items, names=names)
+    graph = get_hetero_graph(has_orig_edge_ids)
+    graph = graph.pin_memory_() if is_pinned else graph.to(F.ctx())
+    if hetero:
+        itemset = gb.HeteroItemSet({"n3": itemset})
+    else:
+        graph.type_per_edge = None
+    item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
+    fanout = torch.LongTensor([2])
+    preprocess_fn = partial(
+        gb.SubgraphSampler._preprocess, cooperative=False, async_op=False
+    )
+    datapipe = item_sampler.map(preprocess_fn)
+    datapipe = datapipe.map(
+        partial(gb.NeighborSampler._prepare, graph.node_type_to_id)
+    )
+    sample_per_layer = gb.SamplePerLayer(
+        datapipe, graph.sample_neighbors, fanout, False, prob_name, False
+    )
+    compact_per_layer = sample_per_layer.compact_per_layer(True)
+    gb.seed(123)
+    expected_results = list(compact_per_layer)
+    if num_cached_edges > 0:
+        graph._initialize_gpu_graph_cache(num_cached_edges, 1, prob_name)
+    datapipe = datapipe.sample_per_layer(
+        graph.sample_neighbors, fanout, False, prob_name, True
+    )
+    datapipe = datapipe.compact_per_layer(True)
+    gb.seed(123)
+    new_results = list(datapipe)
+    assert len(expected_results) == len(new_results)
+    for a, b in zip(expected_results, new_results):
+        assert repr(a) == repr(b)
+
+    def remove_input_nodes(minibatch):
+        minibatch.input_nodes = None
+        return minibatch
+
+    datapipe = item_sampler.sample_neighbor(
+        graph, [fanout], False, prob_name=prob_name, overlap_fetch=True
+    )
+    datapipe = datapipe.transform(remove_input_nodes)
+    dataloader = gb.DataLoader(datapipe)
+    gb.seed(123)
+    new_results = list(dataloader)
+    assert len(expected_results) == len(new_results)
+    for a, b in zip(expected_results, new_results):
+        assert repr(a) == repr(b)
+
+
+@pytest.mark.parametrize("layer_dependency", [False, True])
+@pytest.mark.parametrize("overlap_graph_fetch", [False, True])
+def test_labor_dependent_minibatching(layer_dependency, overlap_graph_fetch):
+    if F._default_context_str != "gpu" and overlap_graph_fetch:
+        pytest.skip("overlap_graph_fetch is only available for GPU.")
+    num_edges = 200
+    csc_indptr = torch.cat(
+        (
+            torch.zeros(1, dtype=torch.int64),
+            torch.ones(num_edges + 1, dtype=torch.int64) * num_edges,
+        )
+    )
+    indices = torch.arange(1, num_edges + 1)
+    graph = gb.fused_csc_sampling_graph(
+        csc_indptr.int(),
+        indices.int(),
+    ).to(F.ctx())
+    torch.random.set_rng_state(torch.manual_seed(123).get_state())
+    batch_dependency = 100
+    itemset = gb.ItemSet(torch.zeros(batch_dependency + 1).int(), names="seeds")
+    datapipe = gb.ItemSampler(itemset, batch_size=1).copy_to(F.ctx())
+    fanouts = [5, 5]
+    datapipe = datapipe.sample_layer_neighbor(
+        graph,
+        fanouts,
+        overlap_fetch=overlap_graph_fetch,
+        layer_dependency=layer_dependency,
+        batch_dependency=batch_dependency,
+    )
+    dataloader = gb.DataLoader(datapipe)
+    res = list(dataloader)
+    assert len(res) == batch_dependency + 1
+    if layer_dependency:
+        assert torch.equal(
+            res[0].input_nodes,
+            res[0].sampled_subgraphs[1].original_row_node_ids,
+        )
+    else:
+        assert res[0].input_nodes.size(0) > res[0].sampled_subgraphs[
+            1
+        ].original_row_node_ids.size(0)
+    delta = 0
+    for i in range(batch_dependency):
+        res_current = (
+            res[i].sampled_subgraphs[-1].original_row_node_ids.tolist()
+        )
+        res_next = (
+            res[i + 1].sampled_subgraphs[-1].original_row_node_ids.tolist()
+        )
+        intersect_len = len(set(res_current).intersection(set(res_next)))
+        assert intersect_len >= fanouts[-1]
+        delta += 1 + fanouts[-1] - intersect_len
+    assert delta >= fanouts[-1]

tests/python/pytorch/graphbolt/impl/test_sampled_subgraph_impl.py

@@ -0,0 +1,735 @@
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.graphbolt as gb
+import pytest
+import torch
+from dgl.graphbolt.impl.sampled_subgraph_impl import SampledSubgraphImpl
+
+
+def _assert_container_equal(lhs, rhs):
+    if isinstance(lhs, torch.Tensor):
+        assert isinstance(rhs, torch.Tensor)
+        assert torch.equal(lhs, rhs)
+    elif isinstance(lhs, tuple):
+        assert isinstance(rhs, tuple)
+        assert len(lhs) == len(rhs)
+        for l, r in zip(lhs, rhs):
+            _assert_container_equal(l, r)
+    elif isinstance(lhs, gb.CSCFormatBase):
+        assert isinstance(rhs, gb.CSCFormatBase)
+        assert len(lhs.indptr) == len(rhs.indptr)
+        assert len(lhs.indices) == len(rhs.indices)
+        _assert_container_equal(lhs.indptr, rhs.indptr)
+        _assert_container_equal(lhs.indices, rhs.indices)
+    elif isinstance(lhs, dict):
+        assert isinstance(rhs, dict)
+        assert len(lhs) == len(rhs)
+        for key, value in lhs.items():
+            assert key in rhs
+            _assert_container_equal(value, rhs[key])
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_homo_deduplicated(reverse_row, reverse_column):
+    csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 2, 2, 3]), indices=torch.tensor([0, 3, 2])
+    )
+    if reverse_row:
+        original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        src_to_exclude = torch.tensor([11])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2])
+
+    if reverse_column:
+        original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        dst_to_exclude = torch.tensor([9])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([4])
+    original_edge_ids = torch.Tensor([5, 9, 10])
+    subgraph = SampledSubgraphImpl(
+        csc_formats,
+        original_column_node_ids,
+        original_row_node_ids,
+        original_edge_ids,
+    )
+    edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(2, -1).T
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 2, 2, 2]), indices=torch.tensor([0, 3])
+    )
+    if reverse_row:
+        expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = torch.Tensor([5, 9])
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_homo_duplicated(reverse_row, reverse_column):
+    csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 3, 3, 5]),
+        indices=torch.tensor([0, 3, 3, 2, 2]),
+    )
+    if reverse_row:
+        original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        src_to_exclude = torch.tensor([24])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([3])
+
+    if reverse_column:
+        original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        dst_to_exclude = torch.tensor([11])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([2])
+    original_edge_ids = torch.Tensor([5, 9, 9, 10, 10])
+    subgraph = SampledSubgraphImpl(
+        csc_formats,
+        original_column_node_ids,
+        original_row_node_ids,
+        original_edge_ids,
+    )
+    edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(2, -1).T
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 1, 1, 3]), indices=torch.tensor([0, 2, 2])
+    )
+    if reverse_row:
+        expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = torch.Tensor([5, 10, 10])
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_hetero_deduplicated(reverse_row, reverse_column):
+    csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 2, 3]),
+            indices=torch.tensor([2, 1, 0]),
+        )
+    }
+    if reverse_row:
+        original_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+        src_to_exclude = torch.tensor([15, 13])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2, 0])
+    if reverse_column:
+        original_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+        dst_to_exclude = torch.tensor([10, 12])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([0, 2])
+    original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+    subgraph = SampledSubgraphImpl(
+        sampled_csc=csc_formats,
+        original_column_node_ids=original_column_node_ids,
+        original_row_node_ids=original_row_node_ids,
+        original_edge_ids=original_edge_ids,
+    )
+
+    edges_to_exclude = {
+        "A:relation:B": torch.cat(
+            (
+                src_to_exclude,
+                dst_to_exclude,
+            )
+        )
+        .view(2, -1)
+        .T
+    }
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 0, 1, 1]),
+            indices=torch.tensor([1]),
+        )
+    }
+    if reverse_row:
+        expected_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = {"A:relation:B": torch.tensor([20])}
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_hetero_duplicated(reverse_row, reverse_column):
+    csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 5]),
+            indices=torch.tensor([2, 2, 1, 1, 0]),
+        )
+    }
+    if reverse_row:
+        original_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+        src_to_exclude = torch.tensor([15, 13])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2, 0])
+    if reverse_column:
+        original_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+        dst_to_exclude = torch.tensor([10, 12])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([0, 2])
+    original_edge_ids = {"A:relation:B": torch.tensor([19, 19, 20, 20, 21])}
+    subgraph = SampledSubgraphImpl(
+        sampled_csc=csc_formats,
+        original_column_node_ids=original_column_node_ids,
+        original_row_node_ids=original_row_node_ids,
+        original_edge_ids=original_edge_ids,
+    )
+
+    edges_to_exclude = {
+        "A:relation:B": torch.cat(
+            (
+                src_to_exclude,
+                dst_to_exclude,
+            )
+        )
+        .view(2, -1)
+        .T
+    }
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 0, 2, 2]),
+            indices=torch.tensor([1, 1]),
+        )
+    }
+    if reverse_row:
+        expected_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = {"A:relation:B": torch.tensor([20, 20])}
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_homo_deduplicated_tensor(reverse_row, reverse_column):
+    csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 2, 2, 3]), indices=torch.tensor([0, 3, 2])
+    )
+    if reverse_row:
+        original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        src_to_exclude = torch.tensor([11])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2])
+
+    if reverse_column:
+        original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        dst_to_exclude = torch.tensor([9])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([4])
+    original_edge_ids = torch.Tensor([5, 9, 10])
+    subgraph = SampledSubgraphImpl(
+        csc_formats,
+        original_column_node_ids,
+        original_row_node_ids,
+        original_edge_ids,
+    )
+    edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(1, -1)
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 2, 2, 2]), indices=torch.tensor([0, 3])
+    )
+    if reverse_row:
+        expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = torch.Tensor([5, 9])
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_homo_duplicated_tensor(reverse_row, reverse_column):
+    csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 3, 3, 5]),
+        indices=torch.tensor([0, 3, 3, 2, 2]),
+    )
+    if reverse_row:
+        original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        src_to_exclude = torch.tensor([24])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([3])
+
+    if reverse_column:
+        original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+        dst_to_exclude = torch.tensor([11])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([2])
+    original_edge_ids = torch.Tensor([5, 9, 9, 10, 10])
+    subgraph = SampledSubgraphImpl(
+        csc_formats,
+        original_column_node_ids,
+        original_row_node_ids,
+        original_edge_ids,
+    )
+    edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(1, -1)
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 0, 1, 1, 1, 3]), indices=torch.tensor([0, 2, 2])
+    )
+    if reverse_row:
+        expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = torch.Tensor([5, 10, 10])
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_hetero_deduplicated_tensor(reverse_row, reverse_column):
+    csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 2, 3]),
+            indices=torch.tensor([2, 1, 0]),
+        )
+    }
+    if reverse_row:
+        original_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+        src_to_exclude = torch.tensor([15, 13])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2, 0])
+    if reverse_column:
+        original_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+        dst_to_exclude = torch.tensor([10, 12])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([0, 2])
+    original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+    subgraph = SampledSubgraphImpl(
+        sampled_csc=csc_formats,
+        original_column_node_ids=original_column_node_ids,
+        original_row_node_ids=original_row_node_ids,
+        original_edge_ids=original_edge_ids,
+    )
+
+    edges_to_exclude = {
+        "A:relation:B": torch.cat((src_to_exclude, dst_to_exclude))
+        .view(2, -1)
+        .T
+    }
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 0, 1, 1]),
+            indices=torch.tensor([1]),
+        )
+    }
+    if reverse_row:
+        expected_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = {"A:relation:B": torch.tensor([20])}
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+@pytest.mark.parametrize("reverse_row", [True, False])
+@pytest.mark.parametrize("reverse_column", [True, False])
+def test_exclude_edges_hetero_duplicated_tensor(reverse_row, reverse_column):
+    csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 5]),
+            indices=torch.tensor([2, 2, 1, 1, 0]),
+        )
+    }
+    if reverse_row:
+        original_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+        src_to_exclude = torch.tensor([15, 13])
+    else:
+        original_row_node_ids = None
+        src_to_exclude = torch.tensor([2, 0])
+    if reverse_column:
+        original_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+        dst_to_exclude = torch.tensor([10, 12])
+    else:
+        original_column_node_ids = None
+        dst_to_exclude = torch.tensor([0, 2])
+    original_edge_ids = {"A:relation:B": torch.tensor([19, 19, 20, 20, 21])}
+    subgraph = SampledSubgraphImpl(
+        sampled_csc=csc_formats,
+        original_column_node_ids=original_column_node_ids,
+        original_row_node_ids=original_row_node_ids,
+        original_edge_ids=original_edge_ids,
+    )
+
+    edges_to_exclude = {
+        "A:relation:B": torch.cat((src_to_exclude, dst_to_exclude))
+        .view(2, -1)
+        .T
+    }
+    result = subgraph.exclude_edges(edges_to_exclude)
+    expected_csc_formats = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 0, 2, 2]),
+            indices=torch.tensor([1, 1]),
+        )
+    }
+    if reverse_row:
+        expected_row_node_ids = {
+            "A": torch.tensor([13, 14, 15]),
+        }
+    else:
+        expected_row_node_ids = None
+    if reverse_column:
+        expected_column_node_ids = {
+            "B": torch.tensor([10, 11, 12]),
+        }
+    else:
+        expected_column_node_ids = None
+    expected_edge_ids = {"A:relation:B": torch.tensor([20, 20])}
+
+    _assert_container_equal(result.sampled_csc, expected_csc_formats)
+    _assert_container_equal(
+        result.original_column_node_ids, expected_column_node_ids
+    )
+    _assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
+    _assert_container_equal(result.original_edge_ids, expected_edge_ids)
+
+
+def test_to_pyg_homo():
+    graph = dgl.graph(([5, 0, 7, 7, 2, 4], [0, 1, 2, 2, 3, 4]))
+    graph = gb.from_dglgraph(graph, is_homogeneous=True).to(F.ctx())
+    items = torch.LongTensor([[0, 3], [4, 4]])
+    names = "seeds"
+    itemset = gb.ItemSet(items, names=names)
+    datapipe = gb.ItemSampler(itemset, batch_size=4).copy_to(F.ctx())
+    num_layer = 2
+    fanouts = [torch.LongTensor([-1]) for _ in range(num_layer)]
+    sampler = gb.NeighborSampler
+    datapipe = sampler(
+        datapipe,
+        graph,
+        fanouts,
+        deduplicate=True,
+    )
+    for minibatch in datapipe:
+        x = torch.randn((minibatch.node_ids().size(0), 2), dtype=torch.float32)
+        for subgraph in minibatch.sampled_subgraphs:
+            (x_src, x_dst), edge_index, sizes = subgraph.to_pyg(x)
+            assert torch.equal(x_src, x)
+            dst_size = subgraph.original_column_node_ids.size(0)
+            assert torch.equal(x_dst, x[:dst_size])
+            src_size = subgraph.original_row_node_ids.size(0)
+            assert dst_size == sizes[1]
+            assert src_size == sizes[0]
+            assert torch.equal(edge_index[0], subgraph.sampled_csc.indices)
+            assert torch.equal(
+                edge_index[1],
+                gb.expand_indptr(
+                    subgraph.sampled_csc.indptr,
+                    subgraph.sampled_csc.indices.dtype,
+                ),
+            )
+            x = x_dst
+
+
+def test_to_pyg_hetero():
+    # COO graph:
+    # [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    # [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
+    # [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
+    # num_nodes = 5, num_n1 = 2, num_n2 = 3
+    ntypes = {"n1": 0, "n2": 1}
+    etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
+    indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
+    indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
+    type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
+    node_type_offset = torch.LongTensor([0, 2, 5])
+    graph = gb.fused_csc_sampling_graph(
+        indptr,
+        indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+    ).to(F.ctx())
+    itemset = gb.HeteroItemSet(
+        {"n1:e1:n2": gb.ItemSet(torch.tensor([[0, 1]]), names="seeds")}
+    )
+    item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    Sampler = gb.NeighborSampler
+    datapipe = Sampler(
+        item_sampler,
+        graph,
+        fanouts,
+        deduplicate=True,
+    )
+    for minibatch in datapipe:
+        x = {}
+        for key, ids in minibatch.node_ids().items():
+            x[key] = torch.randn((ids.size(0), 2), dtype=torch.float32)
+        for subgraph in minibatch.sampled_subgraphs:
+            (x_src, x_dst), edge_index, sizes = subgraph.to_pyg(x)
+            assert x_src == x
+            for ntype in x:
+                dst_size = subgraph.original_column_node_ids[ntype].size(0)
+                assert torch.equal(x_dst[ntype], x[ntype][:dst_size])
+            for etype in subgraph.sampled_csc:
+                src_ntype, _, dst_ntype = gb.etype_str_to_tuple(etype)
+                src_size = subgraph.original_row_node_ids[src_ntype].size(0)
+                dst_size = subgraph.original_column_node_ids[dst_ntype].size(0)
+                assert dst_size == sizes[etype][1]
+                assert src_size == sizes[etype][0]
+                assert torch.equal(
+                    edge_index[etype][0], subgraph.sampled_csc[etype].indices
+                )
+                assert torch.equal(
+                    edge_index[etype][1],
+                    gb.expand_indptr(
+                        subgraph.sampled_csc[etype].indptr,
+                        subgraph.sampled_csc[etype].indices.dtype,
+                    ),
+                )
+            x = x_dst
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="`to` function needs GPU to test.",
+)
+def test_sampled_subgraph_to_device():
+    # Initialize data.
+    csc_format = {
+        "A:relation:B": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 2, 3]),
+            indices=torch.tensor([0, 1, 2]),
+        )
+    }
+    original_row_node_ids = {
+        "A": torch.tensor([13, 14, 15]),
+    }
+    src_to_exclude = torch.tensor([15, 13])
+    original_column_node_ids = {
+        "B": torch.tensor([10, 11, 12]),
+    }
+    dst_to_exclude = torch.tensor([10, 12])
+    original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
+    subgraph = SampledSubgraphImpl(
+        sampled_csc=csc_format,
+        original_column_node_ids=original_column_node_ids,
+        original_row_node_ids=original_row_node_ids,
+        original_edge_ids=original_edge_ids,
+    )
+    edges_to_exclude = {
+        "A:relation:B": torch.cat(
+            (
+                src_to_exclude,
+                dst_to_exclude,
+            )
+        )
+        .view(2, -1)
+        .T
+    }
+    graph = subgraph.exclude_edges(edges_to_exclude)
+
+    # Copy to device.
+    graph = graph.to("cuda")
+
+    # Check.
+    for key in graph.sampled_csc:
+        assert graph.sampled_csc[key].indices.device.type == "cuda"
+        assert graph.sampled_csc[key].indptr.device.type == "cuda"
+    for key in graph.original_column_node_ids:
+        assert graph.original_column_node_ids[key].device.type == "cuda"
+    for key in graph.original_row_node_ids:
+        assert graph.original_row_node_ids[key].device.type == "cuda"
+    for key in graph.original_edge_ids:
+        assert graph.original_edge_ids[key].device.type == "cuda"
+
+
+def test_sampled_subgraph_impl_representation_homo():
+    sampled_subgraph_impl = SampledSubgraphImpl(
+        sampled_csc=gb.CSCFormatBase(
+            indptr=torch.arange(0, 101, 10),
+            indices=torch.arange(10, 110),
+        ),
+        original_column_node_ids=torch.arange(0, 10),
+        original_row_node_ids=torch.arange(0, 110),
+        original_edge_ids=None,
+    )
+    expected_result = str(
+        """SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([  0,  10,  20,  30,  40,  50,  60,  70,  80,  90, 100]),
+                                             indices=tensor([ 10,  11,  12,  13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,
+                                                              24,  25,  26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,
+                                                              38,  39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,
+                                                              52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,
+                                                              66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,  79,
+                                                              80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,  92,  93,
+                                                              94,  95,  96,  97,  98,  99, 100, 101, 102, 103, 104, 105, 106, 107,
+                                                             108, 109]),
+                               ),
+                   original_row_node_ids=tensor([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
+                                                  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
+                                                  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
+                                                  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
+                                                  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
+                                                  70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
+                                                  84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
+                                                  98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109]),
+                   original_edge_ids=None,
+                   original_column_node_ids=tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+)"""
+    )
+    assert str(sampled_subgraph_impl) == expected_result, print(
+        sampled_subgraph_impl
+    )
+
+
+def test_sampled_subgraph_impl_representation_hetero():
+    sampled_subgraph_impl = SampledSubgraphImpl(
+        sampled_csc={
+            "n1:e1:n2": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 2, 4]),
+                indices=torch.tensor([4, 5, 6, 7]),
+            ),
+            "n2:e2:n1": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 2, 4, 6, 8]),
+                indices=torch.tensor([2, 3, 4, 5, 6, 7, 8, 9]),
+            ),
+        },
+        original_column_node_ids={
+            "n1": torch.tensor([1, 0, 0, 1]),
+            "n2": torch.tensor([1, 2]),
+        },
+        original_row_node_ids={
+            "n1": torch.tensor([1, 0, 0, 1, 1, 0, 0, 1]),
+            "n2": torch.tensor([1, 2, 0, 1, 0, 2, 0, 2, 0, 1]),
+        },
+        original_edge_ids=None,
+    )
+    expected_result = str(
+        """SampledSubgraphImpl(sampled_csc={'n1:e1:n2': CSCFormatBase(indptr=tensor([0, 2, 4]),
+                                             indices=tensor([4, 5, 6, 7]),
+                               ), 'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 2, 4, 6, 8]),
+                                             indices=tensor([2, 3, 4, 5, 6, 7, 8, 9]),
+                               )},
+                   original_row_node_ids={'n1': tensor([1, 0, 0, 1, 1, 0, 0, 1]), 'n2': tensor([1, 2, 0, 1, 0, 2, 0, 2, 0, 1])},
+                   original_edge_ids=None,
+                   original_column_node_ids={'n1': tensor([1, 0, 0, 1]), 'n2': tensor([1, 2])},
+)"""
+    )
+    assert str(sampled_subgraph_impl) == expected_result, print(
+        sampled_subgraph_impl
+    )

tests/python/pytorch/graphbolt/impl/test_torch_based_feature_store.py

@@ -0,0 +1,458 @@
+import os
+import tempfile
+import unittest
+
+import backend as F
+
+import numpy as np
+import pydantic
+import pytest
+import torch
+
+from dgl import graphbolt as gb
+
+
+def to_on_disk_tensor(test_dir, name, t):
+    path = os.path.join(test_dir, name + ".npy")
+    t = t.numpy()
+    np.save(path, t)
+    # The Pytorch tensor is a view of the numpy array on disk, which does not
+    # consume memory.
+    t = torch.as_tensor(np.load(path, mmap_mode="r+"))
+    return t
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        metadata = {"max_value": 3}
+        if not in_memory:
+            a = to_on_disk_tensor(test_dir, "a", a)
+            b = to_on_disk_tensor(test_dir, "b", b)
+
+        feature_a = gb.TorchBasedFeature(a, metadata=metadata)
+        feature_b = gb.TorchBasedFeature(b)
+
+        # Read the entire feature.
+        assert torch.equal(
+            feature_a.read(), torch.tensor([[1, 2, 3], [4, 5, 6]])
+        )
+
+        # Test read the feature with ids.
+        assert torch.equal(
+            feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        )
+        # Read the feature with ids.
+        assert torch.equal(
+            feature_a.read(torch.tensor([0])),
+            torch.tensor([[1, 2, 3]]),
+        )
+        assert torch.equal(
+            feature_b.read(torch.tensor([1])),
+            torch.tensor([[[4, 5], [6, 7]]]),
+        )
+        # Update the feature with ids.
+        feature_a.update(torch.tensor([[0, 1, 2]]), torch.tensor([0]))
+        assert torch.equal(
+            feature_a.read(), torch.tensor([[0, 1, 2], [4, 5, 6]])
+        )
+        feature_b.update(torch.tensor([[[1, 2], [3, 4]]]), torch.tensor([1]))
+        assert torch.equal(
+            feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[1, 2], [3, 4]]])
+        )
+
+        # Test update the feature.
+        feature_a.update(torch.tensor([[5, 1, 3]]))
+        assert torch.equal(
+            feature_a.read(),
+            torch.tensor([[5, 1, 3]]),
+        ), print(feature_a.read())
+        feature_b.update(
+            torch.tensor([[[1, 3], [5, 7]], [[2, 4], [6, 8]], [[2, 4], [6, 8]]])
+        )
+        assert torch.equal(
+            feature_b.read(),
+            torch.tensor(
+                [[[1, 3], [5, 7]], [[2, 4], [6, 8]], [[2, 4], [6, 8]]]
+            ),
+        )
+
+        # Test get the size and count of the entire feature.
+        assert feature_a.size() == torch.Size([3])
+        assert feature_b.size() == torch.Size([2, 2])
+        assert feature_a.count() == 1
+        assert feature_b.count() == 3
+
+        # Test get metadata of the feature.
+        assert feature_a.metadata() == metadata
+        assert feature_b.metadata() == {}
+
+        with pytest.raises(IndexError):
+            feature_a.read(torch.tensor([0, 1, 2, 3]))
+
+        # For windows, the file is locked by the numpy.load. We need to delete
+        # it before closing the temporary directory.
+        a = b = None
+        feature_a = feature_b = None
+
+        # Test loaded tensors' contiguity from C/Fortran contiguous ndarray.
+        contiguous_numpy = np.array([[1, 2, 3], [4, 5, 6]], order="C")
+        non_contiguous_numpy = np.array([[1, 2, 3], [4, 5, 6]], order="F")
+        assert contiguous_numpy.flags["C_CONTIGUOUS"]
+        assert non_contiguous_numpy.flags["F_CONTIGUOUS"]
+        np.save(
+            os.path.join(test_dir, "contiguous_numpy.npy"), contiguous_numpy
+        )
+        np.save(
+            os.path.join(test_dir, "non_contiguous_numpy.npy"),
+            non_contiguous_numpy,
+        )
+
+        cur_mmap_mode = None
+        if not in_memory:
+            cur_mmap_mode = "r+"
+        feature_a = gb.TorchBasedFeature(
+            torch.from_numpy(
+                np.load(
+                    os.path.join(test_dir, "contiguous_numpy.npy"),
+                    mmap_mode=cur_mmap_mode,
+                )
+            )
+        )
+        feature_b = gb.TorchBasedFeature(
+            torch.from_numpy(
+                np.load(
+                    os.path.join(test_dir, "non_contiguous_numpy.npy"),
+                    mmap_mode=cur_mmap_mode,
+                )
+            )
+        )
+        assert feature_a._tensor.is_contiguous()
+        assert feature_b._tensor.is_contiguous()
+
+        contiguous_numpy = non_contiguous_numpy = None
+        feature_a = feature_b = None
+
+
+def is_feature_store_on_cuda(store):
+    for feature in store._features.values():
+        assert feature._tensor.is_cuda
+
+
+def is_feature_store_on_cpu(store):
+    for feature in store._features.values():
+        assert not feature._tensor.is_cuda
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Tests for pinned memory are only meaningful on GPU.",
+)
+@pytest.mark.parametrize("device", ["pinned", "cuda"])
+def test_feature_store_to_device(device):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
+        write_tensor_to_disk(test_dir, "a", a, fmt="torch")
+        write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
+        feature_data = [
+            gb.OnDiskFeatureData(
+                domain="node",
+                type="paper",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+            ),
+            gb.OnDiskFeatureData(
+                domain="edge",
+                type="paper:cites:paper",
+                name="b",
+                format="numpy",
+                path=os.path.join(test_dir, "b.npy"),
+            ),
+        ]
+        feature_store = gb.TorchBasedFeatureStore(feature_data)
+        feature_store2 = feature_store.to(device)
+        if device == "pinned":
+            assert feature_store2.is_pinned()
+        elif device == "cuda":
+            is_feature_store_on_cuda(feature_store2)
+
+        # The original variable should be untouched.
+        is_feature_store_on_cpu(feature_store)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Tests for pinned memory are only meaningful on GPU.",
+)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.float32,
+        torch.float64,
+        torch.int32,
+        torch.int64,
+        torch.int8,
+        torch.float16,
+        torch.complex128,
+    ],
+)
+@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("shape", [(2, 1), (2, 3), (2, 2, 2), (137, 13, 3)])
+@pytest.mark.parametrize("in_place", [False, True])
+def test_torch_based_pinned_feature(dtype, idtype, shape, in_place):
+    if dtype == torch.complex128:
+        tensor = torch.complex(
+            torch.randint(0, 13, shape, dtype=torch.float64),
+            torch.randint(0, 13, shape, dtype=torch.float64),
+        )
+    else:
+        tensor = torch.randint(0, 13, shape, dtype=dtype)
+    test_tensor = tensor.clone().detach()
+    test_tensor_cuda = test_tensor.cuda()
+
+    feature = gb.TorchBasedFeature(tensor)
+    if in_place:
+        if gb.is_wsl():
+            pytest.skip("In place pinning is not supported on WSL.")
+        feature.pin_memory_()
+
+        # Check if pinning is truly in-place.
+        assert feature._tensor.data_ptr() == tensor.data_ptr()
+    else:
+        feature = feature.to("pinned")
+
+    assert feature.is_pinned()
+
+    # Test read entire pinned feature, the result should be on cuda.
+    assert torch.equal(feature.read(), test_tensor_cuda)
+    assert feature.read().is_cuda
+    assert torch.equal(
+        feature.read(torch.tensor([0], dtype=idtype).cuda()),
+        test_tensor_cuda[[0]],
+    )
+
+    # Test read pinned feature with idx on cuda, the result should be on cuda.
+    assert feature.read(torch.tensor([0], dtype=idtype).cuda()).is_cuda
+
+    # Test read pinned feature with idx on cpu, the result should be on cpu.
+    assert torch.equal(
+        feature.read(torch.tensor([0], dtype=idtype)), test_tensor[[0]]
+    )
+    assert not feature.read(torch.tensor([0], dtype=idtype)).is_cuda
+
+
+def write_tensor_to_disk(dir, name, t, fmt="torch"):
+    if fmt == "torch":
+        torch.save(t, os.path.join(dir, name + ".pt"))
+    elif fmt == "numpy":
+        t = t.numpy()
+        np.save(os.path.join(dir, name + ".npy"), t)
+    else:
+        raise ValueError(f"Unsupported format: {fmt}")
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature_store(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
+        write_tensor_to_disk(test_dir, "a", a, fmt="torch")
+        write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
+        feature_data = [
+            gb.OnDiskFeatureData(
+                domain="node",
+                type="paper",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+                in_memory=True,
+            ),
+            gb.OnDiskFeatureData(
+                domain="edge",
+                type="paper:cites:paper",
+                name="b",
+                format="numpy",
+                path=os.path.join(test_dir, "b.npy"),
+                in_memory=in_memory,
+            ),
+        ]
+        feature_store = gb.TorchBasedFeatureStore(feature_data)
+
+        assert isinstance(
+            feature_store[("node", "paper", "a")], gb.TorchBasedFeature
+        )
+        assert isinstance(
+            feature_store[("edge", "paper:cites:paper", "b")],
+            gb.TorchBasedFeature if in_memory else gb.DiskBasedFeature,
+        )
+
+        # Test read the entire feature.
+        assert torch.equal(
+            feature_store.read("node", "paper", "a"),
+            torch.tensor([[1, 2, 4], [2, 5, 3]]),
+        )
+        assert torch.equal(
+            feature_store.read("edge", "paper:cites:paper", "b"),
+            torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]]),
+        )
+
+        # Test get the size of the entire feature.
+        assert feature_store.size("node", "paper", "a") == torch.Size([3])
+        assert feature_store.size(
+            "edge", "paper:cites:paper", "b"
+        ) == torch.Size([2, 2])
+
+        # Test get the keys of the features.
+        assert feature_store.keys() == [
+            ("node", "paper", "a"),
+            ("edge", "paper:cites:paper", "b"),
+        ]
+
+        # For windows, the file is locked by the numpy.load. We need to delete
+        # it before closing the temporary directory.
+        a = b = None
+        feature_store = None
+
+        # ``domain`` should be enum.
+        with pytest.raises(pydantic.ValidationError):
+            _ = gb.OnDiskFeatureData(
+                domain="invalid",
+                type="paper",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+                in_memory=True,
+            )
+
+        # ``type`` could be null.
+        feature_data = [
+            gb.OnDiskFeatureData(
+                domain="node",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+                in_memory=True,
+            ),
+        ]
+        feature_store = gb.TorchBasedFeatureStore(feature_data)
+        # Test read the entire feature.
+        assert torch.equal(
+            feature_store.read("node", None, "a"),
+            torch.tensor([[1, 2, 4], [2, 5, 3]]),
+        )
+        # Test get the size of the entire feature.
+        assert feature_store.size("node", None, "a") == torch.Size([3])
+
+        feature_store = None
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature_repr(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 3], [4, 5, 6]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
+        metadata = {"max_value": 3}
+        if not in_memory:
+            a = to_on_disk_tensor(test_dir, "a", a)
+            b = to_on_disk_tensor(test_dir, "b", b)
+
+        feature_a = gb.TorchBasedFeature(a, metadata=metadata)
+        feature_b = gb.TorchBasedFeature(b)
+
+        expected_str_feature_a = (
+            "TorchBasedFeature(\n"
+            "    feature=tensor([[1, 2, 3],\n"
+            "                    [4, 5, 6]]),\n"
+            "    metadata={'max_value': 3},\n"
+            ")"
+        )
+        expected_str_feature_b = (
+            "TorchBasedFeature(\n"
+            "    feature=tensor([[[1, 2],\n"
+            "                     [3, 4]],\n"
+            "\n"
+            "                    [[4, 5],\n"
+            "                     [6, 7]]]),\n"
+            "    metadata={},\n"
+            ")"
+        )
+
+        assert repr(feature_a) == expected_str_feature_a, feature_a
+        assert repr(feature_b) == expected_str_feature_b, feature_b
+
+        a = b = metadata = None
+        feature_a = feature_b = None
+        expected_str_feature_a = expected_str_feature_b = None
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_torch_based_feature_store_repr(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        a = torch.tensor([[1, 2, 4], [2, 5, 3]])
+        b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
+        write_tensor_to_disk(test_dir, "a", a, fmt="torch")
+        write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
+        feature_data = [
+            gb.OnDiskFeatureData(
+                domain="node",
+                type="paper",
+                name="a",
+                format="torch",
+                path=os.path.join(test_dir, "a.pt"),
+                in_memory=True,
+            ),
+            gb.OnDiskFeatureData(
+                domain="edge",
+                type="paper:cites:paper",
+                name="b",
+                format="numpy",
+                path=os.path.join(test_dir, "b.npy"),
+                in_memory=in_memory,
+            ),
+        ]
+        feature_store = gb.TorchBasedFeatureStore(feature_data)
+
+        expected_feature_store_str = (
+            (
+                "TorchBasedFeatureStore(\n"
+                "    {(<OnDiskFeatureDataDomain.NODE: 'node'>, 'paper', 'a'): TorchBasedFeature(\n"
+                "        feature=tensor([[1, 2, 4],\n"
+                "                        [2, 5, 3]]),\n"
+                "        metadata={},\n"
+                "    ), (<OnDiskFeatureDataDomain.EDGE: 'edge'>, 'paper:cites:paper', 'b'): TorchBasedFeature(\n"
+                "        feature=tensor([[[1, 2],\n"
+                "                         [3, 4]],\n"
+                "\n"
+                "                        [[2, 5],\n"
+                "                         [3, 4]]]),\n"
+                "        metadata={},\n"
+                "    )}\n"
+                ")"
+            )
+            if in_memory
+            else (
+                "TorchBasedFeatureStore(\n"
+                "    {(<OnDiskFeatureDataDomain.NODE: 'node'>, 'paper', 'a'): TorchBasedFeature(\n"
+                "        feature=tensor([[1, 2, 4],\n"
+                "                        [2, 5, 3]]),\n"
+                "        metadata={},\n"
+                "    ), (<OnDiskFeatureDataDomain.EDGE: 'edge'>, 'paper:cites:paper', 'b'): DiskBasedFeature(\n"
+                "        feature=tensor([[[1, 2],\n"
+                "                         [3, 4]],\n"
+                "\n"
+                "                        [[2, 5],\n"
+                "                         [3, 4]]]),\n"
+                "        metadata={},\n"
+                "    )}\n"
+                ")"
+            )
+        )
+
+        assert repr(feature_store) == expected_feature_store_str, feature_store
+
+        a = b = feature_data = None
+        feature_store = expected_feature_store_str = None

tests/python/pytorch/graphbolt/internal/test_sample_utils.py

@@ -0,0 +1,273 @@
+import backend as F
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+
+def test_unique_and_compact_hetero():
+    N1 = torch.tensor(
+        [0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
+    )
+    N2 = torch.tensor([0, 3, 3, 5, 2, 7, 2, 8, 4, 9, 2, 3], device=F.ctx())
+    N3 = torch.tensor([1, 2, 6, 6, 1, 8, 3, 6, 3, 2], device=F.ctx())
+    expected_unique = {
+        "n1": torch.tensor([0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()),
+        "n2": torch.tensor([0, 3, 5, 2, 7, 8, 4, 9], device=F.ctx()),
+        "n3": torch.tensor([1, 2, 6, 8, 3], device=F.ctx()),
+    }
+    if N1.is_cuda and torch.cuda.get_device_capability()[0] < 7:
+        expected_reverse_id = {
+            k: v.sort()[1] for k, v in expected_unique.items()
+        }
+        expected_unique = {k: v.sort()[0] for k, v in expected_unique.items()}
+    else:
+        expected_reverse_id = {
+            k: torch.arange(0, v.shape[0], device=F.ctx())
+            for k, v in expected_unique.items()
+        }
+    nodes_dict = {
+        "n1": N1.split(5),
+        "n2": N2.split(4),
+        "n3": N3.split(2),
+    }
+    expected_nodes_dict = {
+        "n1": [
+            torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
+            torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
+            torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
+        ],
+        "n2": [
+            torch.tensor([0, 1, 1, 2], device=F.ctx()),
+            torch.tensor([3, 4, 3, 5], device=F.ctx()),
+            torch.tensor([6, 7, 3, 1], device=F.ctx()),
+        ],
+        "n3": [
+            torch.tensor([0, 1], device=F.ctx()),
+            torch.tensor([2, 2], device=F.ctx()),
+            torch.tensor([0, 3], device=F.ctx()),
+            torch.tensor([4, 2], device=F.ctx()),
+            torch.tensor([4, 1], device=F.ctx()),
+        ],
+    }
+
+    unique, compacted, _ = gb.unique_and_compact(nodes_dict)
+    for ntype, nodes in unique.items():
+        expected_nodes = expected_unique[ntype]
+        assert torch.equal(nodes, expected_nodes)
+
+    for ntype, nodes in compacted.items():
+        expected_nodes = expected_nodes_dict[ntype]
+        assert isinstance(nodes, list)
+        for expected_node, node in zip(expected_nodes, nodes):
+            node = expected_reverse_id[ntype][node]
+            assert torch.equal(expected_node, node)
+
+
+def test_unique_and_compact_homo():
+    N = torch.tensor(
+        [0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
+    )
+    expected_unique_N = torch.tensor(
+        [0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()
+    )
+    if N.is_cuda and torch.cuda.get_device_capability()[0] < 7:
+        expected_reverse_id_N = expected_unique_N.sort()[1]
+        expected_unique_N = expected_unique_N.sort()[0]
+    else:
+        expected_reverse_id_N = torch.arange(
+            0, expected_unique_N.shape[0], device=F.ctx()
+        )
+    nodes_list = N.split(5)
+    expected_nodes_list = [
+        torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
+        torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
+        torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
+    ]
+
+    unique, compacted, _ = gb.unique_and_compact(nodes_list)
+
+    assert torch.equal(unique, expected_unique_N)
+    assert isinstance(compacted, list)
+    for expected_node, node in zip(expected_nodes_list, compacted):
+        node = expected_reverse_id_N[node]
+        assert torch.equal(expected_node, node)
+
+
+def test_unique_and_compact_csc_formats_hetero():
+    dst_nodes = {
+        "n2": torch.tensor([2, 4, 1, 3]),
+        "n3": torch.tensor([1, 3, 2, 7]),
+    }
+    csc_formats = {
+        "n1:e1:n2": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 3, 4, 7, 10]),
+            indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
+        ),
+        "n1:e2:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 4, 7, 10]),
+            indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
+        ),
+        "n2:e3:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 6, 8]),
+            indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
+        ),
+    }
+
+    expected_unique_nodes = {
+        "n1": torch.tensor([1, 3, 4, 6, 2, 7, 9, 5, 8, 0]),
+        "n2": torch.tensor([2, 4, 1, 3, 5, 6, 0]),
+        "n3": torch.tensor([1, 3, 2, 7]),
+    }
+    expected_csc_formats = {
+        "n1:e1:n2": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 3, 4, 7, 10]),
+            indices=torch.tensor([0, 1, 2, 3, 4, 5, 6, 2, 4, 3]),
+        ),
+        "n1:e2:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 4, 7, 10]),
+            indices=torch.tensor([7, 4, 3, 2, 5, 4, 8, 0, 1, 9]),
+        ),
+        "n2:e3:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 6, 8]),
+            indices=torch.tensor([0, 4, 1, 2, 1, 3, 5, 6]),
+        ),
+    }
+
+    unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
+        csc_formats, dst_nodes
+    )
+
+    for ntype, nodes in unique_nodes.items():
+        expected_nodes = expected_unique_nodes[ntype]
+        assert torch.equal(nodes, expected_nodes)
+    for etype, pair in compacted_csc_formats.items():
+        indices = pair.indices
+        indptr = pair.indptr
+        expected_indices = expected_csc_formats[etype].indices
+        expected_indptr = expected_csc_formats[etype].indptr
+        assert torch.equal(indices, expected_indices)
+        assert torch.equal(indptr, expected_indptr)
+
+
+def test_unique_and_compact_csc_formats_homo():
+    seeds = torch.tensor([1, 3, 5, 2, 6])
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
+    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
+
+    expected_unique_nodes = torch.tensor([1, 3, 5, 2, 6, 4])
+    expected_indptr = indptr
+    expected_indices = torch.tensor([3, 1, 0, 5, 2, 3, 2, 0, 5, 5, 4])
+
+    unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
+        csc_formats, seeds
+    )
+
+    indptr = compacted_csc_formats.indptr
+    indices = compacted_csc_formats.indices
+    assert torch.equal(indptr, expected_indptr)
+    assert torch.equal(indices, expected_indices)
+    assert torch.equal(unique_nodes, expected_unique_nodes)
+
+
+def test_unique_and_compact_incorrect_indptr():
+    seeds = torch.tensor([1, 3, 5, 2, 6, 7])
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
+    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
+
+    # The number of seeds is not corresponding to indptr.
+    with pytest.raises(AssertionError):
+        gb.unique_and_compact_csc_formats(csc_formats, seeds)
+
+
+def test_compact_csc_format_hetero():
+    dst_nodes = {
+        "n2": torch.tensor([2, 4, 1, 3]),
+        "n3": torch.tensor([1, 3, 2, 7]),
+    }
+    csc_formats = {
+        "n1:e1:n2": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 3, 4, 7, 10]),
+            indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
+        ),
+        "n1:e2:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 4, 7, 10]),
+            indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
+        ),
+        "n2:e3:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 6, 8]),
+            indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
+        ),
+    }
+
+    expected_original_row_ids = {
+        "n1": torch.tensor(
+            [1, 3, 4, 6, 2, 7, 9, 4, 2, 6, 5, 2, 6, 4, 7, 2, 8, 1, 3, 0]
+        ),
+        "n2": torch.tensor([2, 4, 1, 3, 2, 5, 4, 1, 4, 3, 6, 0]),
+        "n3": torch.tensor([1, 3, 2, 7]),
+    }
+    expected_csc_formats = {
+        "n1:e1:n2": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 3, 4, 7, 10]),
+            indices=torch.arange(0, 10),
+        ),
+        "n1:e2:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 4, 7, 10]),
+            indices=torch.arange(0, 10) + 10,
+        ),
+        "n2:e3:n3": gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 4, 6, 8]),
+            indices=torch.arange(0, 8) + 4,
+        ),
+    }
+    original_row_ids, compacted_csc_formats = gb.compact_csc_format(
+        csc_formats, dst_nodes
+    )
+
+    for ntype, nodes in original_row_ids.items():
+        expected_nodes = expected_original_row_ids[ntype]
+        assert torch.equal(nodes, expected_nodes)
+    for etype, csc_format in compacted_csc_formats.items():
+        indptr = csc_format.indptr
+        indices = csc_format.indices
+        expected_indptr = expected_csc_formats[etype].indptr
+        expected_indices = expected_csc_formats[etype].indices
+        assert torch.equal(indptr, expected_indptr)
+        assert torch.equal(indices, expected_indices)
+
+
+def test_compact_csc_format_homo():
+    seeds = torch.tensor([1, 3, 5, 2, 6])
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
+    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
+
+    expected_original_row_ids = torch.tensor(
+        [1, 3, 5, 2, 6, 2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6]
+    )
+    expected_indptr = indptr
+    expected_indices = torch.arange(0, len(indices)) + 5
+
+    original_row_ids, compacted_csc_formats = gb.compact_csc_format(
+        csc_formats, seeds
+    )
+
+    indptr = compacted_csc_formats.indptr
+    indices = compacted_csc_formats.indices
+
+    assert torch.equal(indptr, expected_indptr)
+    assert torch.equal(indices, expected_indices)
+    assert torch.equal(original_row_ids, expected_original_row_ids)
+
+
+def test_compact_incorrect_indptr():
+    seeds = torch.tensor([1, 3, 5, 2, 6, 7])
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
+    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
+
+    # The number of seeds is not corresponding to indptr.
+    with pytest.raises(AssertionError):
+        gb.compact_csc_format(csc_formats, seeds)

tests/python/pytorch/graphbolt/internal/test_utils.py

@@ -0,0 +1,286 @@
+import json
+import os
+import re
+import tempfile
+from functools import partial
+
+import dgl.graphbolt as gb
+
+import dgl.graphbolt.internal as internal
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+
+def test_read_torch_data():
+    with tempfile.TemporaryDirectory() as test_dir:
+        save_tensor = torch.tensor([[1, 2, 4], [2, 5, 3]])
+        file_name = os.path.join(test_dir, "save_tensor.pt")
+        torch.save(save_tensor, file_name)
+        read_tensor = internal.utils._read_torch_data(file_name)
+        assert torch.equal(save_tensor, read_tensor)
+        save_tensor = read_tensor = None
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_read_numpy_data(in_memory):
+    with tempfile.TemporaryDirectory() as test_dir:
+        save_numpy = np.array([[1, 2, 4], [2, 5, 3]])
+        file_name = os.path.join(test_dir, "save_numpy.npy")
+        np.save(file_name, save_numpy)
+        read_tensor = internal.utils._read_numpy_data(file_name, in_memory)
+        assert torch.equal(torch.from_numpy(save_numpy), read_tensor)
+        save_numpy = read_tensor = None
+
+
+@pytest.mark.parametrize("fmt", ["torch", "numpy"])
+def test_read_data(fmt):
+    with tempfile.TemporaryDirectory() as test_dir:
+        data = np.array([[1, 2, 4], [2, 5, 3]])
+        type_name = "pt" if fmt == "torch" else "npy"
+        file_name = os.path.join(test_dir, f"save_data.{type_name}")
+        if fmt == "numpy":
+            np.save(file_name, data)
+        elif fmt == "torch":
+            torch.save(torch.from_numpy(data), file_name)
+        read_tensor = internal.read_data(file_name, fmt)
+        assert torch.equal(torch.from_numpy(data), read_tensor)
+
+
+@pytest.mark.parametrize(
+    "data_fmt, save_fmt, contiguous",
+    [
+        ("torch", "torch", True),
+        ("torch", "torch", False),
+        ("torch", "numpy", True),
+        ("torch", "numpy", False),
+        ("numpy", "torch", True),
+        ("numpy", "torch", False),
+        ("numpy", "numpy", True),
+        ("numpy", "numpy", False),
+    ],
+)
+def test_save_data(data_fmt, save_fmt, contiguous):
+    with tempfile.TemporaryDirectory() as test_dir:
+        data = np.array([[1, 2, 4], [2, 5, 3]])
+        if not contiguous:
+            data = np.asfortranarray(data)
+        tensor_data = torch.from_numpy(data)
+        type_name = "pt" if save_fmt == "torch" else "npy"
+        save_file_name = os.path.join(test_dir, f"save_data.{type_name}")
+        # Step1. Save the data.
+        if data_fmt == "torch":
+            internal.save_data(tensor_data, save_file_name, save_fmt)
+        elif data_fmt == "numpy":
+            internal.save_data(data, save_file_name, save_fmt)
+
+        # Step2. Load the data.
+        if save_fmt == "torch":
+            loaded_data = torch.load(save_file_name, weights_only=False)
+            assert loaded_data.is_contiguous()
+            assert torch.equal(tensor_data, loaded_data)
+        elif save_fmt == "numpy":
+            loaded_data = np.load(save_file_name)
+            # Checks if the loaded data is C-contiguous.
+            assert loaded_data.flags["C_CONTIGUOUS"]
+            assert np.array_equal(tensor_data.numpy(), loaded_data)
+
+        data = tensor_data = loaded_data = None
+
+
+@pytest.mark.parametrize("fmt", ["torch", "numpy"])
+def test_get_npy_dim(fmt):
+    with tempfile.TemporaryDirectory() as test_dir:
+        data = np.array([[1, 2, 4], [2, 5, 3]])
+        type_name = "pt" if fmt == "torch" else "npy"
+        file_name = os.path.join(test_dir, f"save_data.{type_name}")
+        if fmt == "numpy":
+            np.save(file_name, data)
+            assert internal.get_npy_dim(file_name) == 2
+        elif fmt == "torch":
+            torch.save(torch.from_numpy(data), file_name)
+            with pytest.raises(ValueError):
+                internal.get_npy_dim(file_name)
+        data = None
+
+
+@pytest.mark.parametrize("data_fmt", ["numpy", "torch"])
+@pytest.mark.parametrize("save_fmt", ["numpy", "torch"])
+@pytest.mark.parametrize("is_feature", [True, False])
+def test_copy_or_convert_data(data_fmt, save_fmt, is_feature):
+    with tempfile.TemporaryDirectory() as test_dir:
+        data = np.arange(10)
+        tensor_data = torch.from_numpy(data)
+        in_type_name = "npy" if data_fmt == "numpy" else "pt"
+        input_path = os.path.join(test_dir, f"data.{in_type_name}")
+        out_type_name = "npy" if save_fmt == "numpy" else "pt"
+        output_path = os.path.join(test_dir, f"out_data.{out_type_name}")
+        if data_fmt == "numpy":
+            np.save(input_path, data)
+        else:
+            torch.save(tensor_data, input_path)
+        if save_fmt == "torch":
+            with pytest.raises(AssertionError):
+                internal.copy_or_convert_data(
+                    input_path,
+                    output_path,
+                    data_fmt,
+                    save_fmt,
+                    is_feature=is_feature,
+                )
+        else:
+            internal.copy_or_convert_data(
+                input_path,
+                output_path,
+                data_fmt,
+                save_fmt,
+                is_feature=is_feature,
+            )
+        if is_feature:
+            data = data.reshape(-1, 1)
+            tensor_data = tensor_data.reshape(-1, 1)
+        if save_fmt == "numpy":
+            out_data = np.load(output_path)
+            assert (data == out_data).all()
+
+        data = None
+        tensor_data = None
+        out_data = None
+
+
+@pytest.mark.parametrize("edge_fmt", ["csv", "numpy"])
+def test_read_edges(edge_fmt):
+    with tempfile.TemporaryDirectory() as test_dir:
+        num_nodes = 40
+        num_edges = 200
+        nodes = np.repeat(np.arange(num_nodes), 5)
+        neighbors = np.random.randint(0, num_nodes, size=(num_edges))
+        edges = np.stack([nodes, neighbors], axis=1)
+        os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
+        if edge_fmt == "csv":
+            # Wrtie into edges/edge.csv
+            edges = pd.DataFrame(edges, columns=["src", "dst"])
+            edge_path = os.path.join("edges", "edge.csv")
+            edges.to_csv(
+                os.path.join(test_dir, edge_path),
+                index=False,
+                header=False,
+            )
+        else:
+            # Wrtie into edges/edge.npy
+            edges = edges.T
+            edge_path = os.path.join("edges", "edge.npy")
+            np.save(os.path.join(test_dir, edge_path), edges)
+        src, dst = internal.read_edges(test_dir, edge_fmt, edge_path)
+        assert src.all() == nodes.all()
+        assert dst.all() == neighbors.all()
+
+
+def test_read_edges_error():
+    # 1. Unsupported file format.
+    with pytest.raises(
+        AssertionError,
+        match="`numpy` or `csv` is expected when reading edges but got `fake-type`.",
+    ):
+        internal.read_edges("test_dir", "fake-type", "edge_path")
+
+    # 2. Unexpected shape of numpy array
+    with tempfile.TemporaryDirectory() as test_dir:
+        num_nodes = 40
+        num_edges = 200
+        nodes = np.repeat(np.arange(num_nodes), 5)
+        neighbors = np.random.randint(0, num_nodes, size=(num_edges))
+        edges = np.stack([nodes, neighbors, nodes], axis=1)
+        os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
+        # Wrtie into edges/edge.npy
+        edges = edges.T
+        edge_path = os.path.join("edges", "edge.npy")
+        np.save(os.path.join(test_dir, edge_path), edges)
+        with pytest.raises(
+            AssertionError,
+            match=re.escape(
+                "The shape of edges should be (2, N), but got torch.Size([3, 200])."
+            ),
+        ):
+            internal.read_edges(test_dir, "numpy", edge_path)
+
+
+def test_calculate_file_hash():
+    with tempfile.TemporaryDirectory() as test_dir:
+        test_file_path = os.path.join(test_dir, "test.txt")
+        with open(test_file_path, "w") as file:
+            file.write("test content")
+        hash_value = internal.calculate_file_hash(
+            test_file_path, hash_algo="md5"
+        )
+        expected_hash_value = "9473fdd0d880a43c21b7778d34872157"
+        assert expected_hash_value == hash_value
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Hash algorithm must be one of: ['md5', 'sha1', 'sha224', "
+                + "'sha256', 'sha384', 'sha512'], but got `fake`."
+            ),
+        ):
+            hash_value = internal.calculate_file_hash(
+                test_file_path, hash_algo="fake"
+            )
+
+
+def test_calculate_dir_hash():
+    with tempfile.TemporaryDirectory() as test_dir:
+        test_file_path_1 = os.path.join(test_dir, "test_1.txt")
+        test_file_path_2 = os.path.join(test_dir, "test_2.txt")
+        with open(test_file_path_1, "w") as file:
+            file.write("test content")
+        with open(test_file_path_2, "w") as file:
+            file.write("test contents of directory")
+        hash_value = internal.calculate_dir_hash(test_dir, hash_algo="md5")
+        expected_hash_value = [
+            "56e708a2bdf92887d4a7f25cbc13c555",
+            "9473fdd0d880a43c21b7778d34872157",
+        ]
+        assert len(hash_value) == 2
+        for val in hash_value.values():
+            assert val in expected_hash_value
+
+
+def test_check_dataset_change():
+    with tempfile.TemporaryDirectory() as test_dir:
+        # Generate directory and record its hash value.
+        test_file_path_1 = os.path.join(test_dir, "test_1.txt")
+        test_file_path_2 = os.path.join(test_dir, "test_2.txt")
+        with open(test_file_path_1, "w") as file:
+            file.write("test content")
+        with open(test_file_path_2, "w") as file:
+            file.write("test contents of directory")
+        hash_value = internal.calculate_dir_hash(test_dir, hash_algo="md5")
+        hash_value_file = "dataset_hash_value.txt"
+        hash_value_file_paht = os.path.join(
+            test_dir, "preprocessed", hash_value_file
+        )
+        os.makedirs(os.path.join(test_dir, "preprocessed"), exist_ok=True)
+        with open(hash_value_file_paht, "w") as file:
+            file.write(json.dumps(hash_value, indent=4))
+
+        # Modify the content of a file.
+        with open(test_file_path_2, "w") as file:
+            file.write("test contents of directory changed")
+
+        assert internal.check_dataset_change(test_dir, "preprocessed")
+
+
+def test_numpy_save_aligned():
+    assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
+    a = torch.randn(1024, dtype=torch.float32)  # 4096 bytes
+    with tempfile.TemporaryDirectory() as test_dir:
+        aligned_path = os.path.join(test_dir, "aligned.npy")
+        gb.numpy_save_aligned(aligned_path, a.numpy())
+
+        nonaligned_path = os.path.join(test_dir, "nonaligned.npy")
+        np.save(nonaligned_path, a.numpy())
+
+        assert_equal(np.load(aligned_path), np.load(nonaligned_path))
+        # The size of the file should be 4K (aligned header) + 4K (tensor).
+        assert os.path.getsize(aligned_path) == 4096 * 2

tests/python/pytorch/graphbolt/test_base.py

@@ -0,0 +1,413 @@
+import os
+import re
+import unittest
+from collections.abc import Iterable, Mapping
+
+import backend as F
+
+import dgl.graphbolt as gb
+import pytest
+import torch
+from torch.torch_version import TorchVersion
+
+from . import gb_test_utils
+
+
+def test_pytorch_cuda_allocator_conf():
+    env = os.getenv("PYTORCH_CUDA_ALLOC_CONF")
+    assert env is not None
+    config_list = env.split(",")
+    assert "expandable_segments:True" in config_list
+
+
+@unittest.skipIf(F._default_context_str != "gpu", "CopyTo needs GPU to test")
+@pytest.mark.parametrize("non_blocking", [False, True])
+def test_CopyTo(non_blocking):
+    item_sampler = gb.ItemSampler(
+        gb.ItemSet(torch.arange(20), names="seeds"), 4
+    )
+    if non_blocking:
+        item_sampler = item_sampler.transform(lambda x: x.pin_memory())
+
+    # Invoke CopyTo via class constructor.
+    dp = gb.CopyTo(item_sampler, "cuda")
+    for data in dp:
+        assert data.seeds.device.type == "cuda"
+
+    dp = gb.CopyTo(item_sampler, "cuda", non_blocking)
+    for data in dp:
+        assert data.seeds.device.type == "cuda"
+
+    # Invoke CopyTo via functional form.
+    dp = item_sampler.copy_to("cuda", non_blocking)
+    for data in dp:
+        assert data.seeds.device.type == "cuda"
+
+
+@pytest.mark.parametrize(
+    "task",
+    [
+        "node_classification",
+        "node_inference",
+        "link_prediction",
+        "edge_classification",
+    ],
+)
+@unittest.skipIf(F._default_context_str == "cpu", "CopyTo needs GPU to test")
+def test_CopyToWithMiniBatches(task):
+    N = 16
+    B = 2
+    if task == "node_classification":
+        itemset = gb.ItemSet(
+            (torch.arange(N), torch.arange(N)), names=("seeds", "labels")
+        )
+    elif task == "node_inference":
+        itemset = gb.ItemSet(torch.arange(N), names="seeds")
+    elif task == "link_prediction":
+        itemset = gb.ItemSet(
+            (
+                torch.arange(2 * N).reshape(-1, 2),
+                torch.arange(N),
+            ),
+            names=("seeds", "labels"),
+        )
+    elif task == "edge_classification":
+        itemset = gb.ItemSet(
+            (torch.arange(2 * N).reshape(-1, 2), torch.arange(N)),
+            names=("seeds", "labels"),
+        )
+    graph = gb_test_utils.rand_csc_graph(100, 0.15, bidirection_edge=True)
+
+    features = {}
+    keys = [("node", None, "a"), ("node", None, "b")]
+    features[keys[0]] = gb.TorchBasedFeature(torch.randn(200, 4))
+    features[keys[1]] = gb.TorchBasedFeature(torch.randn(200, 4))
+    feature_store = gb.BasicFeatureStore(features)
+
+    datapipe = gb.ItemSampler(itemset, batch_size=B)
+    datapipe = gb.NeighborSampler(
+        datapipe,
+        graph,
+        fanouts=[torch.LongTensor([2]) for _ in range(2)],
+    )
+    if task != "node_inference":
+        datapipe = gb.FeatureFetcher(
+            datapipe,
+            feature_store,
+            ["a"],
+        )
+
+    copied_attrs = [
+        "labels",
+        "compacted_seeds",
+        "sampled_subgraphs",
+        "indexes",
+        "node_features",
+        "edge_features",
+        "blocks",
+        "seeds",
+        "input_nodes",
+    ]
+
+    def test_data_device(datapipe):
+        for data in datapipe:
+            for attr in dir(data):
+                var = getattr(data, attr)
+                if isinstance(var, Mapping):
+                    var = var[next(iter(var))]
+                elif isinstance(var, Iterable):
+                    var = next(iter(var))
+                if (
+                    not callable(var)
+                    and not attr.startswith("__")
+                    and hasattr(var, "device")
+                    and var is not None
+                ):
+                    if attr in copied_attrs:
+                        assert var.device.type == "cuda", attr
+                    else:
+                        assert var.device.type == "cpu", attr
+
+    # Invoke CopyTo via class constructor.
+    test_data_device(gb.CopyTo(datapipe, "cuda"))
+
+    # Invoke CopyTo via functional form.
+    test_data_device(datapipe.copy_to("cuda"))
+
+
+def test_etype_tuple_to_str():
+    """Convert etype from tuple to string."""
+    # Test for expected input.
+    c_etype = ("user", "like", "item")
+    c_etype_str = gb.etype_tuple_to_str(c_etype)
+    assert c_etype_str == "user:like:item"
+
+    # Test for unexpected input: not a tuple.
+    c_etype = "user:like:item"
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Passed-in canonical etype should be in format of (str, str, str). "
+            "But got user:like:item."
+        ),
+    ):
+        _ = gb.etype_tuple_to_str(c_etype)
+
+    # Test for unexpected input: tuple with wrong length.
+    c_etype = ("user", "like")
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Passed-in canonical etype should be in format of (str, str, str). "
+            "But got ('user', 'like')."
+        ),
+    ):
+        _ = gb.etype_tuple_to_str(c_etype)
+
+
+def test_etype_str_to_tuple():
+    """Convert etype from string to tuple."""
+    # Test for expected input.
+    c_etype_str = "user:like:item"
+    c_etype = gb.etype_str_to_tuple(c_etype_str)
+    assert c_etype == ("user", "like", "item")
+
+    # Test for unexpected input: string with wrong format.
+    c_etype_str = "user:like"
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Passed-in canonical etype should be in format of 'str:str:str'. "
+            "But got user:like."
+        ),
+    ):
+        _ = gb.etype_str_to_tuple(c_etype_str)
+
+
+def test_seed_type_str_to_ntypes():
+    """Convert etype from string to tuple."""
+    # Test for node pairs.
+    seed_type_str = "user:like:item"
+    seed_size = 2
+    node_type = gb.seed_type_str_to_ntypes(seed_type_str, seed_size)
+    assert node_type == ["user", "item"]
+
+    # Test for node pairs.
+    seed_type_str = "user:item:user"
+    seed_size = 3
+    node_type = gb.seed_type_str_to_ntypes(seed_type_str, seed_size)
+    assert node_type == ["user", "item", "user"]
+
+    # Test for unexpected input: list.
+    seed_type_str = ["user", "item"]
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Passed-in seed type should be string, but got <class 'list'>"
+        ),
+    ):
+        _ = gb.seed_type_str_to_ntypes(seed_type_str, 2)
+
+
+def test_isin():
+    elements = torch.tensor([2, 3, 5, 5, 20, 13, 11], device=F.ctx())
+    test_elements = torch.tensor([2, 5], device=F.ctx())
+    res = gb.isin(elements, test_elements)
+    expected = torch.tensor(
+        [True, False, True, True, False, False, False], device=F.ctx()
+    )
+    assert torch.equal(res, expected)
+
+
+def test_isin_big_data():
+    elements = torch.randint(0, 10000, (10000000,), device=F.ctx())
+    test_elements = torch.randint(0, 10000, (500000,), device=F.ctx())
+    res = gb.isin(elements, test_elements)
+    expected = torch.isin(elements, test_elements)
+    assert torch.equal(res, expected)
+
+
+def test_isin_non_1D_dim():
+    elements = torch.tensor([[2, 3], [5, 5], [20, 13]], device=F.ctx())
+    test_elements = torch.tensor([2, 5], device=F.ctx())
+    with pytest.raises(Exception):
+        gb.isin(elements, test_elements)
+    elements = torch.tensor([2, 3, 5, 5, 20, 13], device=F.ctx())
+    test_elements = torch.tensor([[2, 5]], device=F.ctx())
+    with pytest.raises(Exception):
+        gb.isin(elements, test_elements)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("pinned", [False, True])
+def test_index_select(dtype, idtype, pinned):
+    if F._default_context_str != "gpu" and pinned:
+        pytest.skip("Pinned tests are available only on GPU.")
+    tensor = torch.tensor([[2, 3], [5, 5], [20, 13]], dtype=dtype)
+    tensor = tensor.pin_memory() if pinned else tensor.to(F.ctx())
+    index = torch.tensor([0, 2], dtype=idtype, device=F.ctx())
+    gb_result = gb.index_select(tensor, index)
+    torch_result = tensor.to(F.ctx())[index.long()]
+    assert torch.equal(torch_result, gb_result)
+    if pinned:
+        gb_result = gb.index_select(tensor.cpu(), index.cpu().pin_memory())
+        assert torch.equal(torch_result.cpu(), gb_result)
+        assert gb_result.is_pinned()
+
+    # Test the internal async API
+    future = torch.ops.graphbolt.index_select_async(tensor.cpu(), index.cpu())
+    assert torch.equal(torch_result.cpu(), future.wait())
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ],
+)
+@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
+def test_scatter_async(dtype, idtype):
+    input = torch.tensor([[2, 3], [5, 5], [20, 13]], dtype=dtype)
+    index = torch.ones([1], dtype=idtype)
+    res = torch.ops.graphbolt.scatter_async(input, index, input[2:3])
+    assert torch.equal(
+        torch.tensor([[2, 3], [20, 13], [20, 13]], dtype=dtype), res.wait()
+    )
+
+
+def torch_expand_indptr(indptr, dtype, nodes=None):
+    if nodes is None:
+        nodes = torch.arange(len(indptr) - 1, dtype=dtype, device=indptr.device)
+    return nodes.to(dtype).repeat_interleave(indptr.diff())
+
+
+@pytest.mark.parametrize("nodes", [None, True])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_expand_indptr(nodes, dtype):
+    if nodes:
+        nodes = torch.tensor([1, 7, 3, 4, 5, 8], dtype=dtype, device=F.ctx())
+    indptr = torch.tensor([0, 2, 2, 7, 10, 12, 20], device=F.ctx())
+    torch_result = torch_expand_indptr(indptr, dtype, nodes)
+    gb_result = gb.expand_indptr(indptr, dtype, nodes)
+    assert torch.equal(torch_result, gb_result)
+    gb_result = gb.expand_indptr(indptr, dtype, nodes, indptr[-1].item())
+    assert torch.equal(torch_result, gb_result)
+
+    if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
+        import torch._dynamo as dynamo
+        from torch.testing._internal.optests import opcheck
+
+        # Tests torch.compile compatibility
+        for output_size in [None, indptr[-1].item()]:
+            kwargs = {"node_ids": nodes, "output_size": output_size}
+            opcheck(
+                torch.ops.graphbolt.expand_indptr,
+                (indptr, dtype),
+                kwargs,
+                test_utils=[
+                    "test_schema",
+                    "test_autograd_registration",
+                    "test_faketensor",
+                    "test_aot_dispatch_dynamic",
+                ],
+                raise_exception=True,
+            )
+
+            explanation = dynamo.explain(gb.expand_indptr)(
+                indptr, dtype, nodes, output_size
+            )
+            expected_breaks = -1 if output_size is None else 0
+            assert explanation.graph_break_count == expected_breaks
+
+
+@unittest.skipIf(
+    F._default_context_str != "gpu", "Only GPU implementation is available."
+)
+@pytest.mark.parametrize("offset", [None, True])
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_indptr_edge_ids(offset, dtype):
+    indptr = torch.tensor([0, 2, 2, 7, 10, 12], device=F.ctx())
+    if offset:
+        offset = indptr[:-1]
+        ref_result = torch.arange(
+            0, indptr[-1].item(), dtype=dtype, device=F.ctx()
+        )
+    else:
+        ref_result = torch.tensor(
+            [0, 1, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1], dtype=dtype, device=F.ctx()
+        )
+    gb_result = gb.indptr_edge_ids(indptr, dtype, offset)
+    assert torch.equal(ref_result, gb_result)
+    gb_result = gb.indptr_edge_ids(indptr, dtype, offset, indptr[-1].item())
+    assert torch.equal(ref_result, gb_result)
+
+    if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
+        import torch._dynamo as dynamo
+        from torch.testing._internal.optests import opcheck
+
+        # Tests torch.compile compatibility
+        for output_size in [None, indptr[-1].item()]:
+            kwargs = {"offset": offset, "output_size": output_size}
+            opcheck(
+                torch.ops.graphbolt.indptr_edge_ids,
+                (indptr, dtype),
+                kwargs,
+                test_utils=[
+                    "test_schema",
+                    "test_autograd_registration",
+                    "test_faketensor",
+                    "test_aot_dispatch_dynamic",
+                ],
+                raise_exception=True,
+            )
+
+            explanation = dynamo.explain(gb.indptr_edge_ids)(
+                indptr, dtype, offset, output_size
+            )
+            expected_breaks = -1 if output_size is None else 0
+            assert explanation.graph_break_count == expected_breaks
+
+
+def test_csc_format_base_representation():
+    csc_format_base = gb.CSCFormatBase(
+        indptr=torch.tensor([0, 2, 4]),
+        indices=torch.tensor([4, 5, 6, 7]),
+    )
+    expected_result = str(
+        """CSCFormatBase(indptr=tensor([0, 2, 4]),
+              indices=tensor([4, 5, 6, 7]),
+)"""
+    )
+    assert str(csc_format_base) == expected_result, print(csc_format_base)
+
+
+def test_csc_format_base_incorrect_indptr():
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
+    with pytest.raises(AssertionError):
+        # The value of last element in indptr is not corresponding to indices.
+        csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)

tests/python/pytorch/graphbolt/test_dataloader.py

@@ -0,0 +1,220 @@
+import os
+import unittest
+from sys import platform
+
+import backend as F
+
+import dgl
+import dgl.graphbolt
+import dgl.graphbolt as gb
+import pytest
+import torch
+import torch.distributed as thd
+
+from dgl.graphbolt.datapipes import find_dps, traverse_dps
+
+from . import gb_test_utils
+
+
+@pytest.mark.parametrize("overlap_feature_fetch", [False, True])
+def test_DataLoader(overlap_feature_fetch):
+    N = 40
+    B = 4
+    itemset = dgl.graphbolt.ItemSet(torch.arange(N), names="seeds")
+    graph = gb_test_utils.rand_csc_graph(200, 0.15, bidirection_edge=True)
+    features = {}
+    keys = [("node", None, "a"), ("node", None, "b"), ("edge", None, "c")]
+    features[keys[0]] = dgl.graphbolt.TorchBasedFeature(torch.randn(200, 4))
+    features[keys[1]] = dgl.graphbolt.TorchBasedFeature(torch.randn(200, 4))
+    M = graph.total_num_edges
+    features[keys[2]] = dgl.graphbolt.TorchBasedFeature(torch.randn(M, 1))
+    feature_store = dgl.graphbolt.BasicFeatureStore(features)
+
+    item_sampler = dgl.graphbolt.ItemSampler(itemset, batch_size=B)
+    subgraph_sampler = dgl.graphbolt.NeighborSampler(
+        item_sampler,
+        graph,
+        fanouts=[torch.LongTensor([2]) for _ in range(2)],
+    )
+    feature_fetcher = dgl.graphbolt.FeatureFetcher(
+        subgraph_sampler,
+        feature_store,
+        ["a", "b"],
+        ["c"],
+        overlap_fetch=overlap_feature_fetch,
+    )
+    device_transferrer = dgl.graphbolt.CopyTo(feature_fetcher, F.ctx())
+
+    dataloader = dgl.graphbolt.DataLoader(
+        device_transferrer,
+        num_workers=4,
+    )
+    for i, minibatch in enumerate(dataloader):
+        assert "a" in minibatch.node_features
+        assert "b" in minibatch.node_features
+        for layer_id in range(minibatch.num_layers()):
+            assert "c" in minibatch.edge_features[layer_id]
+    assert i + 1 == N // B
+
+
+@unittest.skipIf(
+    F._default_context_str != "gpu",
+    reason="This test requires the GPU.",
+)
+@pytest.mark.parametrize(
+    "sampler_name", ["NeighborSampler", "LayerNeighborSampler"]
+)
+@pytest.mark.parametrize("enable_feature_fetch", [True, False])
+@pytest.mark.parametrize("overlap_feature_fetch", [True, False])
+@pytest.mark.parametrize("overlap_graph_fetch", [True, False])
+@pytest.mark.parametrize("cooperative", [True, False])
+@pytest.mark.parametrize("asynchronous", [True, False])
+@pytest.mark.parametrize("num_gpu_cached_edges", [0, 1024])
+@pytest.mark.parametrize("gpu_cache_threshold", [1, 3])
+def test_gpu_sampling_DataLoader(
+    sampler_name,
+    enable_feature_fetch,
+    overlap_feature_fetch,
+    overlap_graph_fetch,
+    cooperative,
+    asynchronous,
+    num_gpu_cached_edges,
+    gpu_cache_threshold,
+):
+    if cooperative and not thd.is_initialized():
+        # On Windows, the init method can only be file.
+        init_method = (
+            f"file:///{os.path.join(os.getcwd(), 'dis_tempfile')}"
+            if platform == "win32"
+            else "tcp://127.0.0.1:12345"
+        )
+        thd.init_process_group(
+            init_method=init_method,
+            world_size=1,
+            rank=0,
+        )
+    N = 40
+    B = 4
+    num_layers = 2
+    itemset = dgl.graphbolt.ItemSet(torch.arange(N), names="seeds")
+    graph = gb_test_utils.rand_csc_graph(200, 0.15, bidirection_edge=True)
+    graph = graph.pin_memory_() if overlap_graph_fetch else graph.to(F.ctx())
+    features = {}
+    keys = [
+        ("node", None, "a"),
+        ("node", None, "b"),
+        ("node", None, "c"),
+        ("edge", None, "d"),
+    ]
+    features[keys[0]] = dgl.graphbolt.TorchBasedFeature(
+        torch.randn(200, 4, pin_memory=True)
+    )
+    features[keys[1]] = dgl.graphbolt.TorchBasedFeature(
+        torch.randn(200, 4, pin_memory=True)
+    )
+    features[keys[2]] = dgl.graphbolt.TorchBasedFeature(
+        torch.randn(200, 4, device=F.ctx())
+    )
+    features[keys[3]] = dgl.graphbolt.TorchBasedFeature(
+        torch.randn(graph.total_num_edges, 1, device=F.ctx())
+    )
+    feature_store = dgl.graphbolt.BasicFeatureStore(features)
+
+    dataloaders = []
+    for i in range(2):
+        datapipe = dgl.graphbolt.ItemSampler(itemset, batch_size=B)
+        datapipe = datapipe.copy_to(F.ctx())
+        kwargs = {
+            "overlap_fetch": overlap_graph_fetch,
+            "num_gpu_cached_edges": num_gpu_cached_edges,
+            "gpu_cache_threshold": gpu_cache_threshold,
+            "cooperative": cooperative,
+            "asynchronous": asynchronous,
+        }
+        if i != 0:
+            kwargs = {}
+        datapipe = getattr(dgl.graphbolt, sampler_name)(
+            datapipe,
+            graph,
+            fanouts=[torch.LongTensor([2]) for _ in range(num_layers)],
+            **kwargs,
+        )
+        if enable_feature_fetch:
+            datapipe = dgl.graphbolt.FeatureFetcher(
+                datapipe,
+                feature_store,
+                ["a", "b", "c"],
+                ["d"],
+                overlap_fetch=overlap_feature_fetch and i == 0,
+                cooperative=asynchronous and cooperative and i == 0,
+            )
+        dataloaders.append(dgl.graphbolt.DataLoader(datapipe))
+    dataloader, dataloader2 = dataloaders
+
+    bufferer_cnt = int(enable_feature_fetch and overlap_feature_fetch)
+    if overlap_graph_fetch:
+        bufferer_cnt += num_layers
+        if num_gpu_cached_edges > 0:
+            bufferer_cnt += 2 * num_layers
+    if asynchronous:
+        bufferer_cnt += 2 * num_layers + 1  # _preprocess stage has 1.
+        if cooperative:
+            bufferer_cnt += 3 * num_layers
+            if enable_feature_fetch:
+                bufferer_cnt += 1  # feature fetch has 1.
+    if cooperative:
+        # _preprocess stage.
+        bufferer_cnt += 4
+    datapipe_graph = traverse_dps(dataloader)
+    bufferers = find_dps(
+        datapipe_graph,
+        dgl.graphbolt.Bufferer,
+    )
+    assert len(bufferers) == bufferer_cnt
+    # Fixes the randomness of LayerNeighborSampler
+    torch.manual_seed(1)
+    minibatches = list(dataloader)
+    assert len(minibatches) == N // B
+
+    for i, _ in enumerate(dataloader):
+        if i >= 1:
+            break
+
+    torch.manual_seed(1)
+
+    for minibatch, minibatch2 in zip(minibatches, dataloader2):
+        if enable_feature_fetch:
+            assert "a" in minibatch.node_features
+            assert "b" in minibatch.node_features
+            assert "c" in minibatch.node_features
+            if sampler_name == "LayerNeighborSampler":
+                assert torch.equal(
+                    minibatch.node_features["a"], minibatch2.node_features["a"]
+                )
+            for layer_id in range(minibatch.num_layers()):
+                assert "d" in minibatch.edge_features[layer_id]
+                edge_feature = minibatch.edge_features[layer_id]["d"]
+                edge_feature_ref = minibatch2.edge_features[layer_id]["d"]
+                if sampler_name == "LayerNeighborSampler":
+                    assert torch.equal(edge_feature, edge_feature_ref)
+    assert len(list(dataloader)) == N // B
+
+    if asynchronous and cooperative:
+        for minibatch in minibatches:
+            x = torch.ones((minibatch.node_ids().size(0), 1), device=F.ctx())
+            for subgraph in minibatch.sampled_subgraphs:
+                x = gb.CooperativeConvFunction.apply(subgraph, x)
+                x, edge_index, size = subgraph.to_pyg(x)
+                x = x[0]
+                one = torch.ones(
+                    edge_index.shape[1], dtype=x.dtype, device=x.device
+                )
+                coo = torch.sparse_coo_tensor(
+                    edge_index.flipud(), one, size=(size[1], size[0])
+                )
+                x = torch.sparse.mm(coo, x)
+            assert x.shape[0] == minibatch.seeds.shape[0]
+            assert x.shape[1] == 1
+
+    if thd.is_initialized():
+        thd.destroy_process_group()

tests/python/pytorch/graphbolt/test_dataset.py

@@ -0,0 +1,15 @@
+import pytest
+
+from dgl import graphbolt as gb
+
+
+def test_Dataset():
+    dataset = gb.Dataset()
+    with pytest.raises(NotImplementedError):
+        _ = dataset.tasks
+    with pytest.raises(NotImplementedError):
+        _ = dataset.graph
+    with pytest.raises(NotImplementedError):
+        _ = dataset.feature
+    with pytest.raises(NotImplementedError):
+        _ = dataset.dataset_name

tests/python/pytorch/graphbolt/test_feature_fetcher.py

@@ -0,0 +1,258 @@
+import random
+from functools import partial
+
+import dgl.graphbolt as gb
+import torch
+from torch.utils.data.datapipes.iter import Mapper
+
+from . import gb_test_utils
+
+
+def test_FeatureFetcher_invoke():
+    # Prepare graph and required datapipes.
+    graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
+    a = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
+    )
+    b = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
+    )
+
+    features = {}
+    keys = [("node", None, "a"), ("edge", None, "b")]
+    features[keys[0]] = gb.TorchBasedFeature(a)
+    features[keys[1]] = gb.TorchBasedFeature(b)
+    feature_store = gb.BasicFeatureStore(features)
+
+    itemset = gb.ItemSet(torch.arange(10), names="seeds")
+    item_sampler = gb.ItemSampler(itemset, batch_size=2)
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+
+    # Invoke FeatureFetcher via class constructor.
+    datapipe = gb.NeighborSampler(item_sampler, graph, fanouts)
+
+    datapipe = gb.FeatureFetcher(datapipe, feature_store, ["a"], ["b"])
+    assert len(list(datapipe)) == 5
+
+    # Invoke FeatureFetcher via functional form.
+    datapipe = item_sampler.sample_neighbor(graph, fanouts).fetch_feature(
+        feature_store, ["a"], ["b"]
+    )
+    assert len(list(datapipe)) == 5
+
+
+def test_FeatureFetcher_homo():
+    graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
+    a = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
+    )
+    b = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
+    )
+
+    features = {}
+    keys = [("node", None, "a"), ("edge", None, "b")]
+    features[keys[0]] = gb.TorchBasedFeature(a)
+    features[keys[1]] = gb.TorchBasedFeature(b)
+    feature_store = gb.BasicFeatureStore(features)
+
+    itemset = gb.ItemSet(torch.arange(10), names="seeds")
+    item_sampler = gb.ItemSampler(itemset, batch_size=2)
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    sampler_dp = gb.NeighborSampler(item_sampler, graph, fanouts)
+    fetcher_dp = gb.FeatureFetcher(sampler_dp, feature_store, ["a"], ["b"])
+
+    assert len(list(fetcher_dp)) == 5
+
+
+def _func(fn, minibatch):
+    return fn(minibatch)
+
+
+def test_FeatureFetcher_with_edges_homo():
+    graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
+    a = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
+    )
+    b = torch.tensor(
+        [[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
+    )
+
+    def add_node_and_edge_ids(minibatch):
+        seeds = minibatch.seeds
+        subgraphs = []
+        for _ in range(3):
+            sampled_csc = gb.CSCFormatBase(
+                indptr=torch.arange(11),
+                indices=torch.arange(10),
+            )
+            subgraphs.append(
+                gb.SampledSubgraphImpl(
+                    sampled_csc=sampled_csc,
+                    original_column_node_ids=torch.arange(10),
+                    original_row_node_ids=torch.arange(10),
+                    original_edge_ids=torch.randint(
+                        0, graph.total_num_edges, (10,)
+                    ),
+                )
+            )
+        data = gb.MiniBatch(input_nodes=seeds, sampled_subgraphs=subgraphs)
+        return data
+
+    features = {}
+    keys = [("node", None, "a"), ("edge", None, "b")]
+    features[keys[0]] = gb.TorchBasedFeature(a)
+    features[keys[1]] = gb.TorchBasedFeature(b)
+    feature_store = gb.BasicFeatureStore(features)
+
+    itemset = gb.ItemSet(torch.arange(10), names="seeds")
+    item_sampler_dp = gb.ItemSampler(itemset, batch_size=2)
+    fn = partial(_func, add_node_and_edge_ids)
+    converter_dp = Mapper(item_sampler_dp, fn)
+    fetcher_dp = gb.FeatureFetcher(converter_dp, feature_store, ["a"], ["b"])
+
+    assert len(list(fetcher_dp)) == 5
+    for data in fetcher_dp:
+        assert data.node_features["a"].size(0) == 2
+        assert len(data.edge_features) == 3
+        for edge_feature in data.edge_features:
+            assert edge_feature["b"].size(0) == 10
+
+
+def get_hetero_graph():
+    # COO graph:
+    # [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    # [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
+    # [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
+    # num_nodes = 5, num_n1 = 2, num_n2 = 3
+    ntypes = {"n1": 0, "n2": 1}
+    etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
+    indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
+    indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
+    type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
+    node_type_offset = torch.LongTensor([0, 2, 5])
+    return gb.fused_csc_sampling_graph(
+        indptr,
+        indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+    )
+
+
+def test_FeatureFetcher_hetero():
+    graph = get_hetero_graph()
+    a = torch.tensor([[random.randint(0, 10)] for _ in range(2)])
+    b = torch.tensor([[random.randint(0, 10)] for _ in range(3)])
+
+    features = {}
+    keys = [("node", "n1", "a"), ("node", "n2", "a")]
+    features[keys[0]] = gb.TorchBasedFeature(a)
+    features[keys[1]] = gb.TorchBasedFeature(b)
+    feature_store = gb.BasicFeatureStore(features)
+
+    itemset = gb.HeteroItemSet(
+        {
+            "n1": gb.ItemSet(torch.LongTensor([0, 1]), names="seeds"),
+            "n2": gb.ItemSet(torch.LongTensor([0, 1, 2]), names="seeds"),
+        }
+    )
+    item_sampler = gb.ItemSampler(itemset, batch_size=2)
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    sampler_dp = gb.NeighborSampler(item_sampler, graph, fanouts)
+    # "n3" is not in the sampled input nodes.
+    node_feature_keys = {"n1": ["a"], "n2": ["a"], "n3": ["a"]}
+    fetcher_dp = gb.FeatureFetcher(
+        sampler_dp, feature_store, node_feature_keys=node_feature_keys
+    )
+    assert len(list(fetcher_dp)) == 3
+
+    # Do not fetch feature for "n1".
+    node_feature_keys = {"n2": ["a"]}
+    fetcher_dp = gb.FeatureFetcher(
+        sampler_dp, feature_store, node_feature_keys=node_feature_keys
+    )
+    for mini_batch in fetcher_dp:
+        assert ("n1", "a") not in mini_batch.node_features
+
+
+def test_FeatureFetcher_with_edges_hetero():
+    a = torch.tensor([[random.randint(0, 10)] for _ in range(20)])
+    b = torch.tensor([[random.randint(0, 10)] for _ in range(50)])
+
+    def add_node_and_edge_ids(minibatch):
+        seeds = minibatch.seeds
+        subgraphs = []
+        original_edge_ids = {
+            "n1:e1:n2": torch.randint(0, 50, (10,)),
+            "n2:e2:n1": torch.randint(0, 50, (10,)),
+        }
+        original_column_node_ids = {
+            "n1": torch.randint(0, 20, (10,)),
+            "n2": torch.randint(0, 20, (10,)),
+        }
+        original_row_node_ids = {
+            "n1": torch.randint(0, 20, (10,)),
+            "n2": torch.randint(0, 20, (10,)),
+        }
+        for _ in range(3):
+            subgraphs.append(
+                gb.SampledSubgraphImpl(
+                    sampled_csc={
+                        "n1:e1:n2": gb.CSCFormatBase(
+                            indptr=torch.arange(11),
+                            indices=torch.arange(10),
+                        ),
+                        "n2:e2:n1": gb.CSCFormatBase(
+                            indptr=torch.arange(11),
+                            indices=torch.arange(10),
+                        ),
+                    },
+                    original_column_node_ids=original_column_node_ids,
+                    original_row_node_ids=original_row_node_ids,
+                    original_edge_ids=original_edge_ids,
+                )
+            )
+        data = gb.MiniBatch(input_nodes=seeds, sampled_subgraphs=subgraphs)
+        return data
+
+    features = {}
+    keys = [
+        ("node", "n1", "a"),
+        ("edge", "n1:e1:n2", "a"),
+        ("edge", "n2:e2:n1", "a"),
+    ]
+    features[keys[0]] = gb.TorchBasedFeature(a)
+    features[keys[1]] = gb.TorchBasedFeature(b)
+    feature_store = gb.BasicFeatureStore(features)
+
+    itemset = gb.HeteroItemSet(
+        {
+            "n1": gb.ItemSet(torch.randint(0, 20, (10,)), names="seeds"),
+        }
+    )
+    item_sampler_dp = gb.ItemSampler(itemset, batch_size=2)
+    fn = partial(_func, add_node_and_edge_ids)
+    converter_dp = Mapper(item_sampler_dp, fn)
+    # "n3:e3:n3" is not in the sampled edges.
+    # Do not fetch feature for "n2:e2:n1".
+    node_feature_keys = {"n1": ["a"]}
+    edge_feature_keys = {"n1:e1:n2": ["a"], "n3:e3:n3": ["a"]}
+    fetcher_dp = gb.FeatureFetcher(
+        converter_dp,
+        feature_store,
+        node_feature_keys=node_feature_keys,
+        edge_feature_keys=edge_feature_keys,
+    )
+
+    assert len(list(fetcher_dp)) == 5
+    for data in fetcher_dp:
+        assert data.node_features[("n1", "a")].size(0) == 2
+        assert len(data.edge_features) == 3
+        for edge_feature in data.edge_features:
+            assert edge_feature[("n1:e1:n2", "a")].size(0) == 10
+            assert ("n2:e2:n1", "a") not in edge_feature

tests/python/pytorch/graphbolt/test_graphbolt_utils.py

@@ -0,0 +1,60 @@
+import backend as F
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+
+def test_find_reverse_edges_homo():
+    edges = torch.tensor([[1, 3, 5], [2, 4, 5]]).T
+    edges = gb.add_reverse_edges(edges)
+    expected_edges = torch.tensor([[1, 3, 5, 2, 4, 5], [2, 4, 5, 1, 3, 5]]).T
+    assert torch.equal(edges, expected_edges)
+    assert torch.equal(edges[1], expected_edges[1])
+
+
+def test_find_reverse_edges_hetero():
+    edges = {
+        "A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
+        "B:rr:A": torch.tensor([[3], [3]]).T,
+    }
+    edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A"})
+    expected_edges = {
+        "A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
+        "B:rr:A": torch.tensor([[3, 2, 5], [3, 1, 5]]).T,
+    }
+    assert torch.equal(edges["A:r:B"], expected_edges["A:r:B"])
+    assert torch.equal(edges["B:rr:A"], expected_edges["B:rr:A"])
+
+
+def test_find_reverse_edges_bi_reverse_types():
+    edges = {
+        "A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
+        "B:rr:A": torch.tensor([[3], [3]]).T,
+    }
+    edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A", "B:rr:A": "A:r:B"})
+    expected_edges = {
+        "A:r:B": torch.tensor([[1, 5, 3], [2, 5, 3]]).T,
+        "B:rr:A": torch.tensor([[3, 2, 5], [3, 1, 5]]).T,
+    }
+    assert torch.equal(edges["A:r:B"], expected_edges["A:r:B"])
+    assert torch.equal(edges["B:rr:A"], expected_edges["B:rr:A"])
+
+
+def test_find_reverse_edges_circual_reverse_types():
+    edges = {
+        "A:r1:B": torch.tensor([[1, 1]]),
+        "B:r2:C": torch.tensor([[2, 2]]),
+        "C:r3:A": torch.tensor([[3, 3]]),
+    }
+    edges = gb.add_reverse_edges(
+        edges, {"A:r1:B": "B:r2:C", "B:r2:C": "C:r3:A", "C:r3:A": "A:r1:B"}
+    )
+    expected_edges = {
+        "A:r1:B": torch.tensor([[1, 3], [1, 3]]).T,
+        "B:r2:C": torch.tensor([[2, 1], [2, 1]]).T,
+        "C:r3:A": torch.tensor([[3, 2], [3, 2]]).T,
+    }
+    assert torch.equal(edges["A:r1:B"], expected_edges["A:r1:B"])
+    assert torch.equal(edges["B:r2:C"], expected_edges["B:r2:C"])
+    assert torch.equal(edges["A:r1:B"], expected_edges["A:r1:B"])
+    assert torch.equal(edges["C:r3:A"], expected_edges["C:r3:A"])

tests/python/pytorch/graphbolt/test_integration.py

@@ -0,0 +1,360 @@
+import dgl
+import dgl.graphbolt as gb
+import dgl.sparse as dglsp
+import torch
+
+
+def test_integration_link_prediction():
+    torch.manual_seed(926)
+
+    indptr = torch.tensor([0, 0, 1, 3, 6, 8, 10])
+    indices = torch.tensor([5, 3, 3, 3, 3, 4, 4, 0, 5, 4])
+
+    matrix_a = dglsp.from_csc(indptr, indices)
+    seeds = torch.t(torch.stack(matrix_a.coo()))
+    node_feature_data = torch.tensor(
+        [
+            [0.9634, 0.2294],
+            [0.6172, 0.7865],
+            [0.2109, 0.1089],
+            [0.8672, 0.2276],
+            [0.5503, 0.8223],
+            [0.5160, 0.2486],
+        ]
+    )
+    edge_feature_data = torch.tensor(
+        [
+            [0.5123, 0.1709, 0.6150],
+            [0.1476, 0.1902, 0.1314],
+            [0.2582, 0.5203, 0.6228],
+            [0.3708, 0.7631, 0.2683],
+            [0.2126, 0.7878, 0.7225],
+            [0.7885, 0.3414, 0.5485],
+            [0.4088, 0.8200, 0.1851],
+            [0.0056, 0.9469, 0.4432],
+            [0.8972, 0.7511, 0.3617],
+            [0.5773, 0.2199, 0.3366],
+        ]
+    )
+
+    item_set = gb.ItemSet(seeds, names="seeds")
+    graph = gb.fused_csc_sampling_graph(indptr, indices)
+
+    node_feature = gb.TorchBasedFeature(node_feature_data)
+    edge_feature = gb.TorchBasedFeature(edge_feature_data)
+    features = {
+        ("node", None, "feat"): node_feature,
+        ("edge", None, "feat"): edge_feature,
+    }
+    feature_store = gb.BasicFeatureStore(features)
+    datapipe = gb.ItemSampler(item_set, batch_size=4)
+    datapipe = datapipe.sample_uniform_negative(graph, 2)
+    fanouts = torch.LongTensor([1])
+    datapipe = datapipe.sample_neighbor(graph, [fanouts, fanouts], replace=True)
+    datapipe = datapipe.transform(gb.exclude_seed_edges)
+    datapipe = datapipe.fetch_feature(
+        feature_store, node_feature_keys=["feat"], edge_feature_keys=["feat"]
+    )
+    dataloader = gb.DataLoader(
+        datapipe,
+    )
+    expected = [
+        str(
+            """MiniBatch(seeds=tensor([[5, 1],
+                        [3, 2],
+                        [3, 2],
+                        [3, 3],
+                        [5, 2],
+                        [5, 1],
+                        [3, 4],
+                        [3, 3],
+                        [3, 5],
+                        [3, 2],
+                        [3, 0],
+                        [3, 4]]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 1, 1, 2, 2], dtype=torch.int32),
+                                                                         indices=tensor([4, 5], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 1, 3, 2, 4, 0]),
+                                               original_edge_ids=tensor([9, 7]),
+                                               original_column_node_ids=tensor([5, 1, 3, 2, 4, 0]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 1, 1, 2, 2], dtype=torch.int32),
+                                                                         indices=tensor([0, 5], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 1, 3, 2, 4, 0]),
+                                               original_edge_ids=tensor([8, 7]),
+                                               original_column_node_ids=tensor([5, 1, 3, 2, 4, 0]),
+                            )],
+          node_features={'feat': tensor([[0.5160, 0.2486],
+                                [0.6172, 0.7865],
+                                [0.8672, 0.2276],
+                                [0.2109, 0.1089],
+                                [0.5503, 0.8223],
+                                [0.9634, 0.2294]])},
+          labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
+          input_nodes=tensor([5, 1, 3, 2, 4, 0]),
+          indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
+          edge_features=[{'feat': tensor([[0.5773, 0.2199, 0.3366],
+                                [0.0056, 0.9469, 0.4432]])},
+                        {'feat': tensor([[0.8972, 0.7511, 0.3617],
+                                [0.0056, 0.9469, 0.4432]])}],
+          compacted_seeds=tensor([[0, 1],
+                                  [2, 3],
+                                  [2, 3],
+                                  [2, 2],
+                                  [0, 3],
+                                  [0, 1],
+                                  [2, 4],
+                                  [2, 2],
+                                  [2, 0],
+                                  [2, 3],
+                                  [2, 5],
+                                  [2, 4]]),
+          blocks=[Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2),
+                 Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2)],
+       )"""
+        ),
+        str(
+            """MiniBatch(seeds=tensor([[3, 3],
+                        [4, 3],
+                        [4, 4],
+                        [0, 4],
+                        [3, 4],
+                        [3, 5],
+                        [4, 1],
+                        [4, 4],
+                        [4, 4],
+                        [4, 5],
+                        [0, 1],
+                        [0, 3]]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 0, 1], dtype=torch.int32),
+                                                                         indices=tensor([3], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([3, 4, 0, 5, 1]),
+                                               original_edge_ids=tensor([0]),
+                                               original_column_node_ids=tensor([3, 4, 0, 5, 1]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([3, 3], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([3, 4, 0, 5, 1]),
+                                               original_edge_ids=tensor([8, 0]),
+                                               original_column_node_ids=tensor([3, 4, 0, 5, 1]),
+                            )],
+          node_features={'feat': tensor([[0.8672, 0.2276],
+                                [0.5503, 0.8223],
+                                [0.9634, 0.2294],
+                                [0.5160, 0.2486],
+                                [0.6172, 0.7865]])},
+          labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
+          input_nodes=tensor([3, 4, 0, 5, 1]),
+          indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
+          edge_features=[{'feat': tensor([[0.5123, 0.1709, 0.6150]])},
+                        {'feat': tensor([[0.8972, 0.7511, 0.3617],
+                                [0.5123, 0.1709, 0.6150]])}],
+          compacted_seeds=tensor([[0, 0],
+                                  [1, 0],
+                                  [1, 1],
+                                  [2, 1],
+                                  [0, 1],
+                                  [0, 3],
+                                  [1, 4],
+                                  [1, 1],
+                                  [1, 1],
+                                  [1, 3],
+                                  [2, 4],
+                                  [2, 0]]),
+          blocks=[Block(num_src_nodes=5, num_dst_nodes=5, num_edges=1),
+                 Block(num_src_nodes=5, num_dst_nodes=5, num_edges=2)],
+       )"""
+        ),
+        str(
+            """MiniBatch(seeds=tensor([[5, 5],
+                        [4, 5],
+                        [5, 5],
+                        [5, 5],
+                        [4, 0],
+                        [4, 0]]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1], dtype=torch.int32),
+                                                                         indices=tensor([1], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 4, 0]),
+                                               original_edge_ids=tensor([6]),
+                                               original_column_node_ids=tensor([5, 4, 0]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1], dtype=torch.int32),
+                                                                         indices=tensor([2], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 4, 0]),
+                                               original_edge_ids=tensor([7]),
+                                               original_column_node_ids=tensor([5, 4, 0]),
+                            )],
+          node_features={'feat': tensor([[0.5160, 0.2486],
+                                [0.5503, 0.8223],
+                                [0.9634, 0.2294]])},
+          labels=tensor([1., 1., 0., 0., 0., 0.]),
+          input_nodes=tensor([5, 4, 0]),
+          indexes=tensor([0, 1, 0, 0, 1, 1]),
+          edge_features=[{'feat': tensor([[0.4088, 0.8200, 0.1851]])},
+                        {'feat': tensor([[0.0056, 0.9469, 0.4432]])}],
+          compacted_seeds=tensor([[0, 0],
+                                  [1, 0],
+                                  [0, 0],
+                                  [0, 0],
+                                  [1, 2],
+                                  [1, 2]]),
+          blocks=[Block(num_src_nodes=3, num_dst_nodes=3, num_edges=1),
+                 Block(num_src_nodes=3, num_dst_nodes=3, num_edges=1)],
+       )"""
+        ),
+    ]
+    for step, data in enumerate(dataloader):
+        assert expected[step] == str(data), print(step, data)
+
+
+def test_integration_node_classification():
+    torch.manual_seed(926)
+
+    indptr = torch.tensor([0, 0, 1, 3, 6, 8, 10])
+    indices = torch.tensor([5, 3, 3, 3, 3, 4, 4, 0, 5, 4])
+
+    seeds = torch.tensor([5, 1, 2, 4, 3, 0])
+    node_feature_data = torch.tensor(
+        [
+            [0.9634, 0.2294],
+            [0.6172, 0.7865],
+            [0.2109, 0.1089],
+            [0.8672, 0.2276],
+            [0.5503, 0.8223],
+            [0.5160, 0.2486],
+        ]
+    )
+    edge_feature_data = torch.tensor(
+        [
+            [0.5123, 0.1709, 0.6150],
+            [0.1476, 0.1902, 0.1314],
+            [0.2582, 0.5203, 0.6228],
+            [0.3708, 0.7631, 0.2683],
+            [0.2126, 0.7878, 0.7225],
+            [0.7885, 0.3414, 0.5485],
+            [0.4088, 0.8200, 0.1851],
+            [0.0056, 0.9469, 0.4432],
+            [0.8972, 0.7511, 0.3617],
+            [0.5773, 0.2199, 0.3366],
+        ]
+    )
+
+    item_set = gb.ItemSet(seeds, names="seeds")
+    graph = gb.fused_csc_sampling_graph(indptr, indices)
+
+    node_feature = gb.TorchBasedFeature(node_feature_data)
+    edge_feature = gb.TorchBasedFeature(edge_feature_data)
+    features = {
+        ("node", None, "feat"): node_feature,
+        ("edge", None, "feat"): edge_feature,
+    }
+    feature_store = gb.BasicFeatureStore(features)
+    datapipe = gb.ItemSampler(item_set, batch_size=2)
+    fanouts = torch.LongTensor([1])
+    datapipe = datapipe.sample_neighbor(graph, [fanouts, fanouts], replace=True)
+    datapipe = datapipe.fetch_feature(
+        feature_store, node_feature_keys=["feat"], edge_feature_keys=["feat"]
+    )
+    dataloader = gb.DataLoader(
+        datapipe,
+    )
+    expected = [
+        str(
+            """MiniBatch(seeds=tensor([5, 1]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([0, 0], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 1]),
+                                               original_edge_ids=tensor([8, 0]),
+                                               original_column_node_ids=tensor([5, 1]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([0, 0], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([5, 1]),
+                                               original_edge_ids=tensor([8, 0]),
+                                               original_column_node_ids=tensor([5, 1]),
+                            )],
+          node_features={'feat': tensor([[0.5160, 0.2486],
+                                [0.6172, 0.7865]])},
+          labels=None,
+          input_nodes=tensor([5, 1]),
+          indexes=None,
+          edge_features=[{'feat': tensor([[0.8972, 0.7511, 0.3617],
+                                [0.5123, 0.1709, 0.6150]])},
+                        {'feat': tensor([[0.8972, 0.7511, 0.3617],
+                                [0.5123, 0.1709, 0.6150]])}],
+          compacted_seeds=None,
+          blocks=[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=2),
+                 Block(num_src_nodes=2, num_dst_nodes=2, num_edges=2)],
+       )"""
+        ),
+        str(
+            """MiniBatch(seeds=tensor([2, 4]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2, 3], dtype=torch.int32),
+                                                                         indices=tensor([2, 1, 2], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([2, 4, 3]),
+                                               original_edge_ids=tensor([1, 6, 3]),
+                                               original_column_node_ids=tensor([2, 4, 3]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([2, 1], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([2, 4, 3]),
+                                               original_edge_ids=tensor([2, 6]),
+                                               original_column_node_ids=tensor([2, 4]),
+                            )],
+          node_features={'feat': tensor([[0.2109, 0.1089],
+                                [0.5503, 0.8223],
+                                [0.8672, 0.2276]])},
+          labels=None,
+          input_nodes=tensor([2, 4, 3]),
+          indexes=None,
+          edge_features=[{'feat': tensor([[0.1476, 0.1902, 0.1314],
+                                [0.4088, 0.8200, 0.1851],
+                                [0.3708, 0.7631, 0.2683]])},
+                        {'feat': tensor([[0.2582, 0.5203, 0.6228],
+                                [0.4088, 0.8200, 0.1851]])}],
+          compacted_seeds=None,
+          blocks=[Block(num_src_nodes=3, num_dst_nodes=3, num_edges=3),
+                 Block(num_src_nodes=3, num_dst_nodes=2, num_edges=2)],
+       )"""
+        ),
+        str(
+            """MiniBatch(seeds=tensor([3, 0]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1], dtype=torch.int32),
+                                                                         indices=tensor([0], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([3, 0]),
+                                               original_edge_ids=tensor([3]),
+                                               original_column_node_ids=tensor([3, 0]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1], dtype=torch.int32),
+                                                                         indices=tensor([0], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([3, 0]),
+                                               original_edge_ids=tensor([3]),
+                                               original_column_node_ids=tensor([3, 0]),
+                            )],
+          node_features={'feat': tensor([[0.8672, 0.2276],
+                                [0.9634, 0.2294]])},
+          labels=None,
+          input_nodes=tensor([3, 0]),
+          indexes=None,
+          edge_features=[{'feat': tensor([[0.3708, 0.7631, 0.2683]])},
+                        {'feat': tensor([[0.3708, 0.7631, 0.2683]])}],
+          compacted_seeds=None,
+          blocks=[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1),
+                 Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1)],
+       )"""
+        ),
+    ]
+    for step, data in enumerate(dataloader):
+        assert expected[step] == str(data), print(step, data)

tests/python/pytorch/graphbolt/test_itemset.py

@@ -0,0 +1,660 @@
+import re
+
+import dgl
+import pytest
+import torch
+from dgl import graphbolt as gb
+
+
+def test_ItemSet_names():
+    # ItemSet with single name.
+    item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
+    assert item_set.names == ("seeds",)
+
+    # ItemSet with multiple names.
+    item_set = gb.ItemSet(
+        (torch.arange(0, 5), torch.arange(5, 10)),
+        names=("seeds", "labels"),
+    )
+    assert item_set.names == ("seeds", "labels")
+
+    # ItemSet without name.
+    item_set = gb.ItemSet(torch.arange(0, 5))
+    assert item_set.names is None
+
+    # Integer-initiated ItemSet with excessive names.
+    with pytest.raises(
+        AssertionError,
+        match=re.escape("Number of items (1) and names (2) don't match."),
+    ):
+        _ = gb.ItemSet(5, names=("seeds", "labels"))
+
+    # ItemSet with mismatched items and names.
+    with pytest.raises(
+        AssertionError,
+        match=re.escape("Number of items (1) and names (2) don't match."),
+    ):
+        _ = gb.ItemSet(torch.arange(0, 5), names=("seeds", "labels"))
+
+
+@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
+def test_ItemSet_scalar_dtype(dtype):
+    item_set = gb.ItemSet(torch.tensor(5, dtype=dtype), names="seeds")
+    for i, item in enumerate(item_set):
+        assert i == item
+        assert item.dtype == dtype
+    assert item_set[2] == torch.tensor(2, dtype=dtype)
+    assert torch.equal(
+        item_set[slice(1, 4, 2)], torch.arange(1, 4, 2, dtype=dtype)
+    )
+
+
+def test_ItemSet_length():
+    # Integer with valid length
+    num = 10
+    item_set = gb.ItemSet(num)
+    assert len(item_set) == 10
+    # Test __iter__() method. Same as below.
+    for i, item in enumerate(item_set):
+        assert i == item
+
+    # Single iterable with valid length.
+    ids = torch.arange(0, 5)
+    item_set = gb.ItemSet(ids)
+    assert len(item_set) == 5
+    for i, item in enumerate(item_set):
+        assert i == item.item()
+
+    # Tuple of iterables with valid length.
+    item_set = gb.ItemSet((torch.arange(0, 5), torch.arange(5, 10)))
+    assert len(item_set) == 5
+    for i, (item1, item2) in enumerate(item_set):
+        assert i == item1.item()
+        assert i + 5 == item2.item()
+
+    class InvalidLength:
+        def __iter__(self):
+            return iter([0, 1, 2])
+
+    # Single iterable with invalid length.
+    with pytest.raises(
+        TypeError, match="object of type 'InvalidLength' has no len()"
+    ):
+        item_set = gb.ItemSet(InvalidLength())
+
+    # Tuple of iterables with invalid length.
+    with pytest.raises(
+        TypeError, match="object of type 'InvalidLength' has no len()"
+    ):
+        item_set = gb.ItemSet((InvalidLength(), InvalidLength()))
+
+
+def test_ItemSet_seed_nodes():
+    # Node IDs with tensor.
+    item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
+    assert item_set.names == ("seeds",)
+    # Iterating over ItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert i == item.item()
+        assert i == item_set[i]
+    # Indexing with a slice.
+    assert torch.equal(item_set[::2], torch.tensor([0, 2, 4]))
+    # Indexing with an Iterable.
+    assert torch.equal(item_set[torch.arange(0, 5)], torch.arange(0, 5))
+
+    # Node IDs with single integer.
+    item_set = gb.ItemSet(5, names="seeds")
+    assert item_set.names == ("seeds",)
+    # Iterating over ItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert i == item.item()
+        assert i == item_set[i]
+    # Indexing with a slice.
+    assert torch.equal(item_set[::2], torch.tensor([0, 2, 4]))
+    assert torch.equal(item_set[torch.arange(0, 5)], torch.arange(0, 5))
+    # Indexing with an integer.
+    assert item_set[0] == 0
+    assert item_set[-1] == 4
+    # Indexing that is out of range.
+    with pytest.raises(IndexError, match="ItemSet index out of range."):
+        _ = item_set[5]
+    with pytest.raises(IndexError, match="ItemSet index out of range."):
+        _ = item_set[-10]
+    # Indexing with invalid input type.
+    with pytest.raises(
+        TypeError,
+        match="ItemSet indices must be int, slice, or torch.Tensor, not <class 'float'>.",
+    ):
+        _ = item_set[1.5]
+
+
+def test_ItemSet_seed_nodes_labels():
+    # Node IDs and labels.
+    seed_nodes = torch.arange(0, 5)
+    labels = torch.randint(0, 3, (5,))
+    item_set = gb.ItemSet((seed_nodes, labels), names=("seeds", "labels"))
+    assert item_set.names == ("seeds", "labels")
+    # Iterating over ItemSet and indexing one by one.
+    for i, (seed_node, label) in enumerate(item_set):
+        assert seed_node == seed_nodes[i]
+        assert label == labels[i]
+        assert seed_node == item_set[i][0]
+        assert label == item_set[i][1]
+    # Indexing with a slice.
+    assert torch.equal(item_set[:][0], seed_nodes)
+    assert torch.equal(item_set[:][1], labels)
+    # Indexing with an Iterable.
+    assert torch.equal(item_set[torch.arange(0, 5)][0], seed_nodes)
+    assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
+
+
+def test_ItemSet_node_pairs():
+    # Node pairs.
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    item_set = gb.ItemSet(node_pairs, names="seeds")
+    assert item_set.names == ("seeds",)
+    # Iterating over ItemSet and indexing one by one.
+    for i, (src, dst) in enumerate(item_set):
+        assert node_pairs[i][0] == src
+        assert node_pairs[i][1] == dst
+        assert node_pairs[i][0] == item_set[i][0]
+        assert node_pairs[i][1] == item_set[i][1]
+    # Indexing with a slice.
+    assert torch.equal(item_set[:], node_pairs)
+    # Indexing with an Iterable.
+    assert torch.equal(item_set[torch.arange(0, 5)], node_pairs)
+
+
+def test_ItemSet_node_pairs_labels():
+    # Node pairs and labels
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    labels = torch.randint(0, 3, (5,))
+    item_set = gb.ItemSet((node_pairs, labels), names=("seeds", "labels"))
+    assert item_set.names == ("seeds", "labels")
+    # Iterating over ItemSet and indexing one by one.
+    for i, (node_pair, label) in enumerate(item_set):
+        assert torch.equal(node_pairs[i], node_pair)
+        assert labels[i] == label
+        assert torch.equal(node_pairs[i], item_set[i][0])
+        assert labels[i] == item_set[i][1]
+    # Indexing with a slice.
+    assert torch.equal(item_set[:][0], node_pairs)
+    assert torch.equal(item_set[:][1], labels)
+    # Indexing with an Iterable.
+    assert torch.equal(item_set[torch.arange(0, 5)][0], node_pairs)
+    assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
+
+
+def test_ItemSet_node_pairs_labels_indexes():
+    # Node pairs and negative destinations.
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    labels = torch.tensor([1, 1, 0, 0, 0])
+    indexes = torch.tensor([0, 1, 0, 0, 1])
+    item_set = gb.ItemSet(
+        (node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
+    )
+    assert item_set.names == ("seeds", "labels", "indexes")
+    # Iterating over ItemSet and indexing one by one.
+    for i, (node_pair, label, index) in enumerate(item_set):
+        assert torch.equal(node_pairs[i], node_pair)
+        assert torch.equal(labels[i], label)
+        assert torch.equal(indexes[i], index)
+        assert torch.equal(node_pairs[i], item_set[i][0])
+        assert torch.equal(labels[i], item_set[i][1])
+        assert torch.equal(indexes[i], item_set[i][2])
+    # Indexing with a slice.
+    assert torch.equal(item_set[:][0], node_pairs)
+    assert torch.equal(item_set[:][1], labels)
+    assert torch.equal(item_set[:][2], indexes)
+    # Indexing with an Iterable.
+    assert torch.equal(item_set[torch.arange(0, 5)][0], node_pairs)
+    assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
+    assert torch.equal(item_set[torch.arange(0, 5)][2], indexes)
+
+
+def test_ItemSet_graphs():
+    # Graphs.
+    graphs = [dgl.rand_graph(10, 20) for _ in range(5)]
+    item_set = gb.ItemSet(graphs)
+    assert item_set.names is None
+    # Iterating over ItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert graphs[i] == item
+        assert graphs[i] == item_set[i]
+    # Indexing with a slice.
+    assert item_set[:] == graphs
+
+
+def test_HeteroItemSet_names():
+    # HeteroItemSet with single name.
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(torch.arange(0, 5), names="seeds"),
+            "item": gb.ItemSet(torch.arange(5, 10), names="seeds"),
+        }
+    )
+    assert item_set.names == ("seeds",)
+
+    # HeteroItemSet with multiple names.
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(
+                (torch.arange(0, 5), torch.arange(5, 10)),
+                names=("seeds", "labels"),
+            ),
+            "item": gb.ItemSet(
+                (torch.arange(5, 10), torch.arange(10, 15)),
+                names=("seeds", "labels"),
+            ),
+        }
+    )
+    assert item_set.names == ("seeds", "labels")
+
+    # HeteroItemSet with no name.
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(torch.arange(0, 5)),
+            "item": gb.ItemSet(torch.arange(5, 10)),
+        }
+    )
+    assert item_set.names is None
+
+    # HeteroItemSet with mismatched items and names.
+    with pytest.raises(
+        AssertionError,
+        match=re.escape("All itemsets must have the same names."),
+    ):
+        _ = gb.HeteroItemSet(
+            {
+                "user": gb.ItemSet(
+                    (torch.arange(0, 5), torch.arange(5, 10)),
+                    names=("seeds", "labels"),
+                ),
+                "item": gb.ItemSet((torch.arange(5, 10),), names=("seeds",)),
+            }
+        )
+
+
+def test_HeteroItemSet_length():
+    # Single iterable with valid length.
+    user_ids = torch.arange(0, 5)
+    item_ids = torch.arange(0, 5)
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(user_ids),
+            "item": gb.ItemSet(item_ids),
+        }
+    )
+    assert len(item_set) == len(user_ids) + len(item_ids)
+
+    # Tuple of iterables with valid length.
+    node_pairs_like = torch.arange(0, 10).reshape(-1, 2)
+    neg_dsts_like = torch.arange(10, 20).reshape(-1, 2)
+    node_pairs_follow = torch.arange(0, 10).reshape(-1, 2)
+    neg_dsts_follow = torch.arange(10, 20).reshape(-1, 2)
+    item_set = gb.HeteroItemSet(
+        {
+            "user:like:item": gb.ItemSet((node_pairs_like, neg_dsts_like)),
+            "user:follow:user": gb.ItemSet(
+                (node_pairs_follow, neg_dsts_follow)
+            ),
+        }
+    )
+    assert len(item_set) == node_pairs_like.size(0) + node_pairs_follow.size(0)
+
+    class InvalidLength:
+        def __iter__(self):
+            return iter([0, 1, 2])
+
+    # Single iterable with invalid length.
+    with pytest.raises(
+        TypeError, match="object of type 'InvalidLength' has no len()"
+    ):
+        item_set = gb.HeteroItemSet(
+            {
+                "user": gb.ItemSet(InvalidLength()),
+                "item": gb.ItemSet(InvalidLength()),
+            }
+        )
+
+    # Tuple of iterables with invalid length.
+    with pytest.raises(
+        TypeError, match="object of type 'InvalidLength' has no len()"
+    ):
+        item_set = gb.HeteroItemSet(
+            {
+                "user:like:item": gb.ItemSet(
+                    (InvalidLength(), InvalidLength())
+                ),
+                "user:follow:user": gb.ItemSet(
+                    (InvalidLength(), InvalidLength())
+                ),
+            }
+        )
+
+
+def test_HeteroItemSet_iteration_seed_nodes():
+    # Node IDs.
+    user_ids = torch.arange(0, 5)
+    item_ids = torch.arange(5, 10)
+    ids = {
+        "user": gb.ItemSet(user_ids, names="seeds"),
+        "item": gb.ItemSet(item_ids, names="seeds"),
+    }
+    chained_ids = []
+    for key, value in ids.items():
+        chained_ids += [(key, v) for v in value]
+    item_set = gb.HeteroItemSet(ids)
+    assert item_set.names == ("seeds",)
+    # Iterating over HeteroItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert len(item) == 1
+        assert isinstance(item, dict)
+        assert chained_ids[i][0] in item
+        assert item[chained_ids[i][0]] == chained_ids[i][1]
+        assert item_set[i] == item
+        assert item_set[i - len(item_set)] == item
+    # Indexing all with a slice.
+    assert torch.equal(item_set[:]["user"], user_ids)
+    assert torch.equal(item_set[:]["item"], item_ids)
+    # Indexing partial with a slice.
+    partial_data = item_set[:3]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["user"], user_ids[:3])
+    partial_data = item_set[7:]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["item"], item_ids[2:])
+    partial_data = item_set[3:8:2]
+    assert len(list(partial_data.keys())) == 2
+    assert torch.equal(partial_data["user"], user_ids[3:-1:2])
+    assert torch.equal(partial_data["item"], item_ids[0:3:2])
+    # Indexing with an iterable of int.
+    partial_data = item_set[torch.tensor([1, 0, 4])]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["user"], torch.tensor([1, 0, 4]))
+    partial_data = item_set[torch.tensor([9, 8, 5])]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["item"], torch.tensor([9, 8, 5]))
+    partial_data = item_set[torch.tensor([8, 1, 0, 9, 7, 5])]
+    assert len(list(partial_data.keys())) == 2
+    assert torch.equal(partial_data["user"], torch.tensor([1, 0]))
+    assert torch.equal(partial_data["item"], torch.tensor([8, 9, 7, 5]))
+
+    # Exception cases.
+    with pytest.raises(
+        AssertionError, match="Start must be smaller than stop."
+    ):
+        _ = item_set[5:3]
+    with pytest.raises(
+        AssertionError, match="Start must be smaller than stop."
+    ):
+        _ = item_set[-1:3]
+    with pytest.raises(IndexError, match="HeteroItemSet index out of range."):
+        _ = item_set[20]
+    with pytest.raises(IndexError, match="HeteroItemSet index out of range."):
+        _ = item_set[-20]
+    with pytest.raises(
+        TypeError,
+        match="HeteroItemSet indices must be int, slice, or iterable of int, not <class 'float'>.",
+    ):
+        _ = item_set[1.5]
+
+
+def test_HeteroItemSet_iteration_seed_nodes_labels():
+    # Node IDs and labels.
+    user_ids = torch.arange(0, 5)
+    user_labels = torch.randint(0, 3, (5,))
+    item_ids = torch.arange(5, 10)
+    item_labels = torch.randint(0, 3, (5,))
+    ids_labels = {
+        "user": gb.ItemSet((user_ids, user_labels), names=("seeds", "labels")),
+        "item": gb.ItemSet((item_ids, item_labels), names=("seeds", "labels")),
+    }
+    chained_ids = []
+    for key, value in ids_labels.items():
+        chained_ids += [(key, v) for v in value]
+    item_set = gb.HeteroItemSet(ids_labels)
+    assert item_set.names == ("seeds", "labels")
+    # Iterating over HeteroItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert len(item) == 1
+        assert isinstance(item, dict)
+        assert chained_ids[i][0] in item
+        assert item[chained_ids[i][0]] == chained_ids[i][1]
+        assert item_set[i] == item
+    # Indexing with a slice.
+    assert torch.equal(item_set[:]["user"][0], user_ids)
+    assert torch.equal(item_set[:]["user"][1], user_labels)
+    assert torch.equal(item_set[:]["item"][0], item_ids)
+    assert torch.equal(item_set[:]["item"][1], item_labels)
+
+
+def test_HeteroItemSet_iteration_node_pairs():
+    # Node pairs.
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    node_pairs_dict = {
+        "user:like:item": gb.ItemSet(node_pairs, names="seeds"),
+        "user:follow:user": gb.ItemSet(node_pairs, names="seeds"),
+    }
+    expected_data = []
+    for key, value in node_pairs_dict.items():
+        expected_data += [(key, v) for v in value]
+    item_set = gb.HeteroItemSet(node_pairs_dict)
+    assert item_set.names == ("seeds",)
+    # Iterating over HeteroItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert len(item) == 1
+        assert isinstance(item, dict)
+        assert expected_data[i][0] in item
+        assert torch.equal(item[expected_data[i][0]], expected_data[i][1])
+        assert item_set[i].keys() == item.keys()
+        key = list(item.keys())[0]
+        assert torch.equal(item_set[i][key], item[key])
+    # Indexing with a slice.
+    assert torch.equal(item_set[:]["user:like:item"], node_pairs)
+    assert torch.equal(item_set[:]["user:follow:user"], node_pairs)
+
+
+def test_HeteroItemSet_iteration_node_pairs_labels():
+    # Node pairs and labels
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    labels = torch.randint(0, 3, (5,))
+    node_pairs_labels = {
+        "user:like:item": gb.ItemSet(
+            (node_pairs, labels), names=("seeds", "labels")
+        ),
+        "user:follow:user": gb.ItemSet(
+            (node_pairs, labels), names=("seeds", "labels")
+        ),
+    }
+    expected_data = []
+    for key, value in node_pairs_labels.items():
+        expected_data += [(key, v) for v in value]
+    item_set = gb.HeteroItemSet(node_pairs_labels)
+    assert item_set.names == ("seeds", "labels")
+    # Iterating over HeteroItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert len(item) == 1
+        assert isinstance(item, dict)
+        key, value = expected_data[i]
+        assert key in item
+        assert torch.equal(item[key][0], value[0])
+        assert item[key][1] == value[1]
+        assert item_set[i].keys() == item.keys()
+        key = list(item.keys())[0]
+        assert torch.equal(item_set[i][key][0], item[key][0])
+        assert torch.equal(item_set[i][key][1], item[key][1])
+    # Indexing with a slice.
+    assert torch.equal(item_set[:]["user:like:item"][0], node_pairs)
+    assert torch.equal(item_set[:]["user:like:item"][1], labels)
+    assert torch.equal(item_set[:]["user:follow:user"][0], node_pairs)
+    assert torch.equal(item_set[:]["user:follow:user"][1], labels)
+
+
+def test_HeteroItemSet_iteration_node_pairs_labels_indexes():
+    # Node pairs and negative destinations.
+    node_pairs = torch.arange(0, 10).reshape(-1, 2)
+    labels = torch.tensor([1, 1, 0, 0, 0])
+    indexes = torch.tensor([0, 1, 0, 0, 1])
+    node_pairs_neg_dsts = {
+        "user:like:item": gb.ItemSet(
+            (node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
+        ),
+        "user:follow:user": gb.ItemSet(
+            (node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
+        ),
+    }
+    expected_data = []
+    for key, value in node_pairs_neg_dsts.items():
+        expected_data += [(key, v) for v in value]
+    item_set = gb.HeteroItemSet(node_pairs_neg_dsts)
+    assert item_set.names == ("seeds", "labels", "indexes")
+    # Iterating over HeteroItemSet and indexing one by one.
+    for i, item in enumerate(item_set):
+        assert len(item) == 1
+        assert isinstance(item, dict)
+        key, value = expected_data[i]
+        assert key in item
+        assert torch.equal(item[key][0], value[0])
+        assert torch.equal(item[key][1], value[1])
+        assert torch.equal(item[key][2], value[2])
+        assert item_set[i].keys() == item.keys()
+        key = list(item.keys())[0]
+        assert torch.equal(item_set[i][key][0], item[key][0])
+        assert torch.equal(item_set[i][key][1], item[key][1])
+        assert torch.equal(item_set[i][key][2], item[key][2])
+    # Indexing with a slice.
+    assert torch.equal(item_set[:]["user:like:item"][0], node_pairs)
+    assert torch.equal(item_set[:]["user:like:item"][1], labels)
+    assert torch.equal(item_set[:]["user:like:item"][2], indexes)
+    assert torch.equal(item_set[:]["user:follow:user"][0], node_pairs)
+    assert torch.equal(item_set[:]["user:follow:user"][1], labels)
+    assert torch.equal(item_set[:]["user:follow:user"][2], indexes)
+
+
+def test_ItemSet_repr():
+    # ItemSet with single name.
+    item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
+    expected_str = (
+        "ItemSet(\n"
+        "    items=(tensor([0, 1, 2, 3, 4]),),\n"
+        "    names=('seeds',),\n"
+        ")"
+    )
+
+    assert str(item_set) == expected_str, item_set
+
+    # ItemSet with multiple names.
+    item_set = gb.ItemSet(
+        (torch.arange(0, 5), torch.arange(5, 10)),
+        names=("seeds", "labels"),
+    )
+    expected_str = (
+        "ItemSet(\n"
+        "    items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n"
+        "    names=('seeds', 'labels'),\n"
+        ")"
+    )
+    assert str(item_set) == expected_str, item_set
+
+
+def test_HeteroItemSet_repr():
+    # HeteroItemSet with single name.
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(torch.arange(0, 5), names="seeds"),
+            "item": gb.ItemSet(torch.arange(5, 10), names="seeds"),
+        }
+    )
+    expected_str = (
+        "HeteroItemSet(\n"
+        "    itemsets={'user': ItemSet(\n"
+        "                 items=(tensor([0, 1, 2, 3, 4]),),\n"
+        "                 names=('seeds',),\n"
+        "             ), 'item': ItemSet(\n"
+        "                 items=(tensor([5, 6, 7, 8, 9]),),\n"
+        "                 names=('seeds',),\n"
+        "             )},\n"
+        "    names=('seeds',),\n"
+        ")"
+    )
+    assert str(item_set) == expected_str, item_set
+
+    # HeteroItemSet with multiple names.
+    item_set = gb.HeteroItemSet(
+        {
+            "user": gb.ItemSet(
+                (torch.arange(0, 5), torch.arange(5, 10)),
+                names=("seeds", "labels"),
+            ),
+            "item": gb.ItemSet(
+                (torch.arange(5, 10), torch.arange(10, 15)),
+                names=("seeds", "labels"),
+            ),
+        }
+    )
+    expected_str = (
+        "HeteroItemSet(\n"
+        "    itemsets={'user': ItemSet(\n"
+        "                 items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n"
+        "                 names=('seeds', 'labels'),\n"
+        "             ), 'item': ItemSet(\n"
+        "                 items=(tensor([5, 6, 7, 8, 9]), tensor([10, 11, 12, 13, 14])),\n"
+        "                 names=('seeds', 'labels'),\n"
+        "             )},\n"
+        "    names=('seeds', 'labels'),\n"
+        ")"
+    )
+    assert str(item_set) == expected_str, item_set
+
+
+def test_deprecation_alias():
+    """Test `ItemSetDict` as the alias for `HeteroItemSet`."""
+
+    user_ids = torch.arange(0, 5)
+    item_ids = torch.arange(5, 10)
+    ids = {
+        "user": gb.ItemSet(user_ids, names="seeds"),
+        "item": gb.ItemSet(item_ids, names="seeds"),
+    }
+    with pytest.warns(
+        DeprecationWarning,
+        match="ItemSetDict is deprecated and will be removed in the future. Please use HeteroItemSet instead.",
+    ):
+        item_set_dict = gb.ItemSetDict(ids)
+    hetero_item_set = gb.HeteroItemSet(ids)
+    assert len(item_set_dict) == len(hetero_item_set)
+    assert item_set_dict.names == hetero_item_set.names
+    assert item_set_dict._keys == hetero_item_set._keys
+    assert torch.equal(item_set_dict._offsets, hetero_item_set._offsets)
+    assert (
+        repr(item_set_dict)[len("ItemSetDict") :]
+        == repr(hetero_item_set)[len("HeteroItemSet") :]
+    )
+    # Indexing all with a slice.
+    assert torch.equal(item_set_dict[:]["user"], hetero_item_set[:]["user"])
+    assert torch.equal(item_set_dict[:]["item"], hetero_item_set[:]["item"])
+    # Indexing partial with a slice.
+    partial_data = item_set_dict[:3]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["user"], hetero_item_set[:3]["user"])
+    partial_data = item_set_dict[7:]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["item"], hetero_item_set[7:]["item"])
+    partial_data = item_set_dict[3:8:2]
+    assert len(list(partial_data.keys())) == 2
+    assert torch.equal(partial_data["user"], hetero_item_set[3:8:2]["user"])
+    assert torch.equal(partial_data["item"], hetero_item_set[3:8:2]["item"])
+    # Indexing with an iterable of int.
+    partial_data = item_set_dict[torch.tensor([1, 0, 4])]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["user"], hetero_item_set[1, 0, 4]["user"])
+    partial_data = item_set_dict[torch.tensor([9, 8, 5])]
+    assert len(list(partial_data.keys())) == 1
+    assert torch.equal(partial_data["item"], hetero_item_set[9, 8, 5]["item"])
+    partial_data = item_set_dict[torch.tensor([8, 1, 0, 9, 7, 5])]
+    assert len(list(partial_data.keys())) == 2
+    assert torch.equal(partial_data["user"], hetero_item_set[1, 0]["user"])
+    assert torch.equal(
+        partial_data["item"], hetero_item_set[8, 9, 7, 5]["item"]
+    )

tests/python/pytorch/graphbolt/test_minibatch.py

@@ -0,0 +1,755 @@
+import dgl
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+
+relation = "A:r:B"
+reverse_relation = "B:rr:A"
+
+
+@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
+def test_minibatch_representation_homo(indptr_dtype, indices_dtype):
+    seeds = torch.tensor([10, 11])
+    csc_formats = [
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 3, 5, 6], dtype=indptr_dtype),
+            indices=torch.tensor([0, 1, 2, 2, 1, 2], dtype=indices_dtype),
+        ),
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 3], dtype=indptr_dtype),
+            indices=torch.tensor([1, 2, 0], dtype=indices_dtype),
+        ),
+    ]
+    original_column_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11]),
+    ]
+    original_row_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11, 12]),
+    ]
+    original_edge_ids = [
+        torch.tensor([19, 20, 21, 22, 25, 30]),
+        torch.tensor([10, 15, 17]),
+    ]
+    node_features = {"x": torch.tensor([5, 0, 2, 1])}
+    edge_features = [
+        {"x": torch.tensor([9, 0, 1, 1, 7, 4])},
+        {"x": torch.tensor([0, 2, 2])},
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=csc_formats[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    input_nodes = torch.tensor([8, 1, 6, 5, 9, 0, 2, 4])
+    compacted_seeds = torch.tensor([0, 1])
+    labels = torch.tensor([1.0, 2.0])
+    # Test minibatch without data.
+    minibatch = gb.MiniBatch()
+    expect_result = str(
+        """MiniBatch(seeds=None,
+          sampled_subgraphs=None,
+          node_features=None,
+          labels=None,
+          input_nodes=None,
+          indexes=None,
+          edge_features=None,
+          compacted_seeds=None,
+          blocks=None,
+       )"""
+    )
+    result = str(minibatch)
+    assert result == expect_result, print(expect_result, result)
+    # Test minibatch with all attributes.
+    minibatch = gb.MiniBatch(
+        seeds=seeds,
+        sampled_subgraphs=subgraphs,
+        labels=labels,
+        node_features=node_features,
+        edge_features=edge_features,
+        compacted_seeds=compacted_seeds,
+        input_nodes=input_nodes,
+    )
+    expect_result = str(
+        """MiniBatch(seeds=tensor([10, 11]),
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 3, 5, 6], dtype=torch.int32),
+                                                                         indices=tensor([0, 1, 2, 2, 1, 2], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([10, 11, 12, 13]),
+                                               original_edge_ids=tensor([19, 20, 21, 22, 25, 30]),
+                                               original_column_node_ids=tensor([10, 11, 12, 13]),
+                            ),
+                            SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 2, 3], dtype=torch.int32),
+                                                                         indices=tensor([1, 2, 0], dtype=torch.int32),
+                                                           ),
+                                               original_row_node_ids=tensor([10, 11, 12]),
+                                               original_edge_ids=tensor([10, 15, 17]),
+                                               original_column_node_ids=tensor([10, 11]),
+                            )],
+          node_features={'x': tensor([5, 0, 2, 1])},
+          labels=tensor([1., 2.]),
+          input_nodes=tensor([8, 1, 6, 5, 9, 0, 2, 4]),
+          indexes=None,
+          edge_features=[{'x': tensor([9, 0, 1, 1, 7, 4])},
+                        {'x': tensor([0, 2, 2])}],
+          compacted_seeds=tensor([0, 1]),
+          blocks=[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=6),
+                 Block(num_src_nodes=3, num_dst_nodes=2, num_edges=3)],
+       )"""
+    )
+    result = str(minibatch)
+    assert result == expect_result, print(expect_result, result)
+
+
+@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
+def test_minibatch_representation_hetero(indptr_dtype, indices_dtype):
+    seeds = {relation: torch.tensor([10, 11])}
+    csc_formats = [
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2, 3], dtype=indptr_dtype),
+                indices=torch.tensor([0, 1, 1], dtype=indices_dtype),
+            ),
+            reverse_relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 0, 0, 1, 2], dtype=indptr_dtype),
+                indices=torch.tensor([1, 0], dtype=indices_dtype),
+            ),
+        },
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2], dtype=indptr_dtype),
+                indices=torch.tensor([1, 0], dtype=indices_dtype),
+            ),
+            reverse_relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 2], dtype=indptr_dtype),
+                indices=torch.tensor([1, 0], dtype=indices_dtype),
+            ),
+        },
+    ]
+    original_column_node_ids = [
+        {"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
+        {"B": torch.tensor([10, 11]), "A": torch.tensor([5])},
+    ]
+    original_row_node_ids = [
+        {
+            "A": torch.tensor([5, 7, 9, 11]),
+            "B": torch.tensor([10, 11, 12]),
+        },
+        {
+            "A": torch.tensor([5, 7]),
+            "B": torch.tensor([10, 11]),
+        },
+    ]
+    original_edge_ids = [
+        {
+            relation: torch.tensor([19, 20, 21]),
+            reverse_relation: torch.tensor([23, 26]),
+        },
+        {relation: torch.tensor([10, 12])},
+    ]
+    node_features = {
+        ("A", "x"): torch.tensor([6, 4, 0, 1]),
+    }
+    edge_features = [
+        {(relation, "x"): torch.tensor([4, 2, 4])},
+        {(relation, "x"): torch.tensor([0, 6])},
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=csc_formats[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    compacted_seeds = {relation: torch.tensor([0, 1])}
+    # Test minibatch with all attributes.
+    minibatch = gb.MiniBatch(
+        seeds=seeds,
+        sampled_subgraphs=subgraphs,
+        node_features=node_features,
+        edge_features=edge_features,
+        labels={"B": torch.tensor([2, 5])},
+        compacted_seeds=compacted_seeds,
+        input_nodes={
+            "A": torch.tensor([5, 7, 9, 11]),
+            "B": torch.tensor([10, 11, 12]),
+        },
+    )
+    expect_result = str(
+        """MiniBatch(seeds={'A:r:B': tensor([10, 11])},
+          sampled_subgraphs=[SampledSubgraphImpl(sampled_csc={'A:r:B': CSCFormatBase(indptr=tensor([0, 1, 2, 3], dtype=torch.int32),
+                                                                         indices=tensor([0, 1, 1], dtype=torch.int32),
+                                                           ), 'B:rr:A': CSCFormatBase(indptr=tensor([0, 0, 0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([1, 0], dtype=torch.int32),
+                                                           )},
+                                               original_row_node_ids={'A': tensor([ 5,  7,  9, 11]), 'B': tensor([10, 11, 12])},
+                                               original_edge_ids={'A:r:B': tensor([19, 20, 21]), 'B:rr:A': tensor([23, 26])},
+                                               original_column_node_ids={'B': tensor([10, 11, 12]), 'A': tensor([ 5,  7,  9, 11])},
+                            ),
+                            SampledSubgraphImpl(sampled_csc={'A:r:B': CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
+                                                                         indices=tensor([1, 0], dtype=torch.int32),
+                                                           ), 'B:rr:A': CSCFormatBase(indptr=tensor([0, 2], dtype=torch.int32),
+                                                                         indices=tensor([1, 0], dtype=torch.int32),
+                                                           )},
+                                               original_row_node_ids={'A': tensor([5, 7]), 'B': tensor([10, 11])},
+                                               original_edge_ids={'A:r:B': tensor([10, 12])},
+                                               original_column_node_ids={'B': tensor([10, 11]), 'A': tensor([5])},
+                            )],
+          node_features={('A', 'x'): tensor([6, 4, 0, 1])},
+          labels={'B': tensor([2, 5])},
+          input_nodes={'A': tensor([ 5,  7,  9, 11]), 'B': tensor([10, 11, 12])},
+          indexes=None,
+          edge_features=[{('A:r:B', 'x'): tensor([4, 2, 4])},
+                        {('A:r:B', 'x'): tensor([0, 6])}],
+          compacted_seeds={'A:r:B': tensor([0, 1])},
+          blocks=[Block(num_src_nodes={'A': 4, 'B': 3},
+                       num_dst_nodes={'A': 4, 'B': 3},
+                       num_edges={('A', 'r', 'B'): 3, ('B', 'rr', 'A'): 2},
+                       metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')]),
+                 Block(num_src_nodes={'A': 2, 'B': 2},
+                       num_dst_nodes={'A': 1, 'B': 2},
+                       num_edges={('A', 'r', 'B'): 2, ('B', 'rr', 'A'): 2},
+                       metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')])],
+       )"""
+    )
+    result = str(minibatch)
+    assert result == expect_result, print(result)
+
+
+@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
+def test_get_dgl_blocks_homo(indptr_dtype, indices_dtype):
+    csc_formats = [
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 3, 5, 6], dtype=indptr_dtype),
+            indices=torch.tensor([0, 1, 2, 2, 1, 2], dtype=indices_dtype),
+        ),
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 3], dtype=indptr_dtype),
+            indices=torch.tensor([0, 1, 2], dtype=indices_dtype),
+        ),
+    ]
+    original_column_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11]),
+    ]
+    original_row_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11, 12]),
+    ]
+    original_edge_ids = [
+        torch.tensor([19, 20, 21, 22, 25, 30]),
+        torch.tensor([10, 15, 17]),
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=csc_formats[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    # Test minibatch with all attributes.
+    minibatch = gb.MiniBatch(
+        sampled_subgraphs=subgraphs,
+    )
+    dgl_blocks = minibatch.blocks
+    expect_result = str(
+        """[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=6), Block(num_src_nodes=3, num_dst_nodes=2, num_edges=3)]"""
+    )
+    result = str(dgl_blocks)
+    assert result == expect_result
+
+
+def test_get_dgl_blocks_hetero():
+    csc_formats = [
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2, 3]),
+                indices=torch.tensor([0, 1, 1]),
+            ),
+            reverse_relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 0, 0, 1, 2]),
+                indices=torch.tensor([1, 0]),
+            ),
+        },
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2]), indices=torch.tensor([1, 0])
+            ),
+            reverse_relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1]),
+                indices=torch.tensor([1]),
+            ),
+        },
+    ]
+    original_column_node_ids = [
+        {"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
+        {"B": torch.tensor([10, 11]), "A": torch.tensor([5])},
+    ]
+    original_row_node_ids = [
+        {
+            "A": torch.tensor([5, 7, 9, 11]),
+            "B": torch.tensor([10, 11, 12]),
+        },
+        {
+            "A": torch.tensor([5, 7]),
+            "B": torch.tensor([10, 11]),
+        },
+    ]
+    original_edge_ids = [
+        {
+            relation: torch.tensor([19, 20, 21]),
+            reverse_relation: torch.tensor([23, 26]),
+        },
+        {relation: torch.tensor([10, 12])},
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=csc_formats[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    # Test minibatch with all attributes.
+    minibatch = gb.MiniBatch(
+        sampled_subgraphs=subgraphs,
+    )
+    dgl_blocks = minibatch.blocks
+    expect_result = str(
+        """[Block(num_src_nodes={'A': 4, 'B': 3},
+      num_dst_nodes={'A': 4, 'B': 3},
+      num_edges={('A', 'r', 'B'): 3, ('B', 'rr', 'A'): 2},
+      metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')]), Block(num_src_nodes={'A': 2, 'B': 2},
+      num_dst_nodes={'A': 1, 'B': 2},
+      num_edges={('A', 'r', 'B'): 2, ('B', 'rr', 'A'): 1},
+      metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')])]"""
+    )
+    result = str(dgl_blocks)
+    assert result == expect_result
+
+
+def test_get_dgl_blocks_hetero_partial_empty_edges():
+    hg = dgl.heterograph(
+        {
+            ("n1", "e1", "n1"): ([0, 1, 1], [1, 2, 0]),
+            ("n1", "e2", "n2"): ([0, 1, 2], [1, 0, 2]),
+        }
+    )
+
+    gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
+
+    train_set = gb.HeteroItemSet(
+        {"n1:e2:n2": gb.ItemSet(torch.LongTensor([[0, 1]]), names="seeds")}
+    )
+    datapipe = gb.ItemSampler(train_set, batch_size=1)
+    datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
+    dataloader = gb.DataLoader(datapipe)
+    blocks_str = str(next(iter(dataloader)).blocks)
+    expected_str = """[Block(num_src_nodes={'n1': 2, 'n2': 0},
+      num_dst_nodes={'n1': 2, 'n2': 0},
+      num_edges={('n1', 'e1', 'n1'): 2, ('n1', 'e2', 'n2'): 0},
+      metagraph=[('n1', 'n1', 'e1'), ('n1', 'n2', 'e2')]), Block(num_src_nodes={'n1': 2, 'n2': 0},
+      num_dst_nodes={'n1': 1, 'n2': 1},
+      num_edges={('n1', 'e1', 'n1'): 1, ('n1', 'e2', 'n2'): 1},
+      metagraph=[('n1', 'n1', 'e1'), ('n1', 'n2', 'e2')])]"""
+    assert expected_str == blocks_str
+
+
+def test_get_dgl_blocks_hetero_empty_edges():
+    hg = dgl.heterograph(
+        {
+            ("n3", "e1", "n1"): ([0, 1, 1], [1, 2, 0]),
+            ("n3", "e2", "n2"): ([0, 1, 2], [1, 0, 2]),
+        }
+    )
+
+    gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
+
+    train_set = gb.HeteroItemSet(
+        {"n3:e1:n1": gb.ItemSet(torch.LongTensor([[2, 1]]), names="seeds")}
+    )
+    datapipe = gb.ItemSampler(train_set, batch_size=1)
+    datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
+    dataloader = gb.DataLoader(datapipe)
+    blocks_str = str(next(iter(dataloader)).blocks)
+    expected_str = """[Block(num_src_nodes={'n1': 0, 'n2': 0, 'n3': 2},
+      num_dst_nodes={'n1': 0, 'n2': 0, 'n3': 2},
+      num_edges={('n3', 'e1', 'n1'): 0, ('n3', 'e2', 'n2'): 0},
+      metagraph=[('n3', 'n1', 'e1'), ('n3', 'n2', 'e2')]), Block(num_src_nodes={'n1': 0, 'n2': 0, 'n3': 2},
+      num_dst_nodes={'n1': 1, 'n2': 0, 'n3': 1},
+      num_edges={('n3', 'e1', 'n1'): 1, ('n3', 'e2', 'n2'): 0},
+      metagraph=[('n3', 'n1', 'e1'), ('n3', 'n2', 'e2')])]"""
+    assert expected_str == blocks_str
+
+
+def test_get_dgl_blocks_homo_empty_edges():
+    g = dgl.graph(([2, 3, 4], [3, 4, 5]))
+
+    gb_g = gb.from_dglgraph(g, is_homogeneous=True)
+    train_set = gb.ItemSet(torch.LongTensor([[0, 1]]), names="seeds")
+    datapipe = gb.ItemSampler(train_set, batch_size=1)
+    datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
+    dataloader = gb.DataLoader(datapipe)
+    blocks_str = str(next(iter(dataloader)).blocks)
+    expected_str = "[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=0), Block(num_src_nodes=2, num_dst_nodes=2, num_edges=0)]"
+    assert expected_str == blocks_str
+
+
+def test_seeds_ntype_being_passed():
+    hg = dgl.heterograph({("n1", "e1", "n2"): ([0, 1, 2], [2, 0, 1])})
+
+    gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
+    train_set = gb.HeteroItemSet(
+        {"n2": gb.ItemSet(torch.LongTensor([0, 1]), names="seeds")}
+    )
+    datapipe = gb.ItemSampler(train_set, batch_size=2)
+    datapipe = datapipe.sample_neighbor(gb_g, [-1, -1, -1])
+    dataloader = gb.DataLoader(datapipe)
+    blocks = next(iter(dataloader)).blocks
+    for block in blocks:
+        assert "n2" in block.srctypes
+
+
+def create_homo_minibatch():
+    csc_formats = [
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 1, 3, 5, 6]),
+            indices=torch.tensor([0, 1, 2, 2, 1, 2]),
+        ),
+        gb.CSCFormatBase(
+            indptr=torch.tensor([0, 2, 3]),
+            indices=torch.tensor([1, 2, 0]),
+        ),
+    ]
+    original_column_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11]),
+    ]
+    original_row_node_ids = [
+        torch.tensor([10, 11, 12, 13]),
+        torch.tensor([10, 11, 12]),
+    ]
+    original_edge_ids = [
+        torch.tensor([19, 20, 21, 22, 25, 30]),
+        torch.tensor([10, 15, 17]),
+    ]
+    node_features = {"x": torch.randint(0, 10, (4,))}
+    edge_features = [
+        {"x": torch.randint(0, 10, (6,))},
+        {"x": torch.randint(0, 10, (3,))},
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=csc_formats[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    return gb.MiniBatch(
+        sampled_subgraphs=subgraphs,
+        node_features=node_features,
+        edge_features=edge_features,
+        input_nodes=torch.tensor([10, 11, 12, 13]),
+    )
+
+
+def create_hetero_minibatch():
+    sampled_csc = [
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2, 3]),
+                indices=torch.tensor([0, 1, 1]),
+            ),
+            reverse_relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 0, 0, 1, 2]),
+                indices=torch.tensor([1, 0]),
+            ),
+        },
+        {
+            relation: gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 2]), indices=torch.tensor([1, 0])
+            )
+        },
+    ]
+    original_column_node_ids = [
+        {"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
+        {"B": torch.tensor([10, 11])},
+    ]
+    original_row_node_ids = [
+        {
+            "A": torch.tensor([5, 7, 9, 11]),
+            "B": torch.tensor([10, 11, 12]),
+        },
+        {
+            "A": torch.tensor([5, 7]),
+            "B": torch.tensor([10, 11]),
+        },
+    ]
+    original_edge_ids = [
+        {
+            relation: torch.tensor([19, 20, 21]),
+            reverse_relation: torch.tensor([23, 26]),
+        },
+        {relation: torch.tensor([10, 12])},
+    ]
+    node_features = {
+        ("A", "x"): torch.randint(0, 10, (4,)),
+    }
+    edge_features = [
+        {(relation, "x"): torch.randint(0, 10, (3,))},
+        {(relation, "x"): torch.randint(0, 10, (2,))},
+    ]
+    subgraphs = []
+    for i in range(2):
+        subgraphs.append(
+            gb.SampledSubgraphImpl(
+                sampled_csc=sampled_csc[i],
+                original_column_node_ids=original_column_node_ids[i],
+                original_row_node_ids=original_row_node_ids[i],
+                original_edge_ids=original_edge_ids[i],
+            )
+        )
+    return gb.MiniBatch(
+        sampled_subgraphs=subgraphs,
+        node_features=node_features,
+        edge_features=edge_features,
+        input_nodes={
+            "A": torch.tensor([5, 7, 9, 11]),
+            "B": torch.tensor([10, 11, 12]),
+        },
+    )
+
+
+def check_dgl_blocks_hetero(minibatch, blocks):
+    etype = gb.etype_str_to_tuple(relation)
+    sampled_csc = [
+        subgraph.sampled_csc for subgraph in minibatch.sampled_subgraphs
+    ]
+    original_edge_ids = [
+        subgraph.original_edge_ids for subgraph in minibatch.sampled_subgraphs
+    ]
+    original_row_node_ids = [
+        subgraph.original_row_node_ids
+        for subgraph in minibatch.sampled_subgraphs
+    ]
+
+    for i, block in enumerate(blocks):
+        edges = block.edges(etype=etype)
+        dst_ndoes = torch.arange(
+            0, len(sampled_csc[i][relation].indptr) - 1
+        ).repeat_interleave(sampled_csc[i][relation].indptr.diff())
+        assert torch.equal(edges[0], sampled_csc[i][relation].indices)
+        assert torch.equal(edges[1], dst_ndoes)
+        assert torch.equal(
+            block.edges[etype].data[dgl.EID], original_edge_ids[i][relation]
+        )
+    edges = blocks[0].edges(etype=gb.etype_str_to_tuple(reverse_relation))
+    dst_ndoes = torch.arange(
+        0, len(sampled_csc[0][reverse_relation].indptr) - 1
+    ).repeat_interleave(sampled_csc[0][reverse_relation].indptr.diff())
+    assert torch.equal(edges[0], sampled_csc[0][reverse_relation].indices)
+    assert torch.equal(edges[1], dst_ndoes)
+    assert torch.equal(
+        blocks[0].srcdata[dgl.NID]["A"], original_row_node_ids[0]["A"]
+    )
+    assert torch.equal(
+        blocks[0].srcdata[dgl.NID]["B"], original_row_node_ids[0]["B"]
+    )
+
+
+def check_dgl_blocks_homo(minibatch, blocks):
+    sampled_csc = [
+        subgraph.sampled_csc for subgraph in minibatch.sampled_subgraphs
+    ]
+    original_edge_ids = [
+        subgraph.original_edge_ids for subgraph in minibatch.sampled_subgraphs
+    ]
+    original_row_node_ids = [
+        subgraph.original_row_node_ids
+        for subgraph in minibatch.sampled_subgraphs
+    ]
+    for i, block in enumerate(blocks):
+        dst_ndoes = torch.arange(
+            0, len(sampled_csc[i].indptr) - 1
+        ).repeat_interleave(sampled_csc[i].indptr.diff())
+        assert torch.equal(block.edges()[0], sampled_csc[i].indices)
+        assert torch.equal(block.edges()[1], dst_ndoes)
+        assert torch.equal(block.edata[dgl.EID], original_edge_ids[i])
+    assert torch.equal(blocks[0].srcdata[dgl.NID], original_row_node_ids[0])
+
+
+def test_dgl_node_classification_without_feature():
+    # Arrange
+    minibatch = create_homo_minibatch()
+    minibatch.node_features = None
+    minibatch.labels = None
+    minibatch.seeds = torch.tensor([10, 15])
+    # Act
+    dgl_blocks = minibatch.blocks
+
+    # Assert
+    assert len(dgl_blocks) == 2
+    assert minibatch.node_features is None
+    assert minibatch.labels is None
+    check_dgl_blocks_homo(minibatch, dgl_blocks)
+
+
+def test_dgl_node_classification_homo():
+    # Arrange
+    minibatch = create_homo_minibatch()
+    minibatch.seeds = torch.tensor([10, 15])
+    minibatch.labels = torch.tensor([2, 5])
+    # Act
+    dgl_blocks = minibatch.blocks
+
+    # Assert
+    assert len(dgl_blocks) == 2
+    check_dgl_blocks_homo(minibatch, dgl_blocks)
+
+
+def test_dgl_node_classification_hetero():
+    minibatch = create_hetero_minibatch()
+    minibatch.labels = {"B": torch.tensor([2, 5])}
+    minibatch.seeds = {"B": torch.tensor([10, 15])}
+    # Act
+    dgl_blocks = minibatch.blocks
+
+    # Assert
+    assert len(dgl_blocks) == 2
+    check_dgl_blocks_hetero(minibatch, dgl_blocks)
+
+
+def test_dgl_link_predication_homo():
+    # Arrange
+    minibatch = create_homo_minibatch()
+    minibatch.compacted_seeds = (
+        torch.tensor([[0, 1, 0, 0, 1, 1], [1, 0, 1, 1, 0, 0]]).T,
+    )
+    minibatch.labels = torch.tensor([1, 1, 0, 0, 0, 0])
+    # Act
+    dgl_blocks = minibatch.blocks
+
+    # Assert
+    assert len(dgl_blocks) == 2
+    check_dgl_blocks_homo(minibatch, dgl_blocks)
+
+
+def test_dgl_link_predication_hetero():
+    # Arrange
+    minibatch = create_hetero_minibatch()
+    minibatch.compacted_seeds = {
+        relation: (torch.tensor([[1, 1, 2, 0, 1, 2], [1, 0, 1, 1, 0, 0]]).T,),
+        reverse_relation: (
+            torch.tensor([[0, 1, 1, 2, 0, 2], [1, 0, 1, 1, 0, 0]]).T,
+        ),
+    }
+    minibatch.labels = {
+        relation: (torch.tensor([1, 1, 0, 0, 0, 0]),),
+        reverse_relation: (torch.tensor([1, 1, 0, 0, 0, 0]),),
+    }
+    # Act
+    dgl_blocks = minibatch.blocks
+
+    # Assert
+    assert len(dgl_blocks) == 2
+    check_dgl_blocks_hetero(minibatch, dgl_blocks)
+
+
+def test_to_pyg_data():
+    test_minibatch = create_homo_minibatch()
+    test_minibatch.seeds = torch.tensor([0, 1])
+    test_minibatch.labels = torch.tensor([7, 8])
+
+    expected_edge_index = torch.tensor(
+        [[0, 0, 1, 1, 1, 2, 2, 2, 2], [0, 1, 0, 1, 2, 0, 1, 2, 3]]
+    )
+    expected_node_features = next(iter(test_minibatch.node_features.values()))
+    expected_labels = torch.tensor([7, 8])
+    expected_batch_size = 2
+    expected_n_id = torch.tensor([10, 11, 12, 13])
+
+    pyg_data = test_minibatch.to_pyg_data()
+    pyg_data.validate()
+    assert torch.equal(pyg_data.edge_index, expected_edge_index)
+    assert torch.equal(pyg_data.x, expected_node_features)
+    assert torch.equal(pyg_data.y, expected_labels)
+    assert pyg_data.batch_size == expected_batch_size
+    assert torch.equal(pyg_data.n_id, expected_n_id)
+
+    test_minibatch.seeds = torch.tensor([[0, 1], [2, 3]])
+    assert pyg_data.batch_size == expected_batch_size
+
+    test_minibatch.seeds = {"A": torch.tensor([0, 1])}
+    assert pyg_data.batch_size == expected_batch_size
+
+    test_minibatch.seeds = {"A": torch.tensor([[0, 1], [2, 3]])}
+    assert pyg_data.batch_size == expected_batch_size
+
+    subgraph = test_minibatch.sampled_subgraphs[0]
+    # Test with sampled_csc as None.
+    test_minibatch = gb.MiniBatch(
+        sampled_subgraphs=None,
+        node_features={"feat": expected_node_features},
+        labels=expected_labels,
+    )
+    pyg_data = test_minibatch.to_pyg_data()
+    assert pyg_data.edge_index is None, "Edge index should be none."
+
+    # Test with node_features as None.
+    test_minibatch = gb.MiniBatch(
+        sampled_subgraphs=[subgraph],
+        node_features=None,
+        labels=expected_labels,
+    )
+    pyg_data = test_minibatch.to_pyg_data()
+    assert pyg_data.x is None, "Node features should be None."
+
+    # Test with labels as None.
+    test_minibatch = gb.MiniBatch(
+        sampled_subgraphs=[subgraph],
+        node_features={"feat": expected_node_features},
+        labels=None,
+    )
+    pyg_data = test_minibatch.to_pyg_data()
+    assert pyg_data.y is None, "Labels should be None."
+
+    # Test with multiple features.
+    test_minibatch = gb.MiniBatch(
+        sampled_subgraphs=[subgraph],
+        node_features={
+            "feat": expected_node_features,
+            "extra_feat": torch.tensor([[3], [4]]),
+        },
+        labels=expected_labels,
+    )
+    try:
+        pyg_data = test_minibatch.to_pyg_data()
+        assert (
+            pyg_data.x is None
+        ), "Multiple features case should raise an error."
+    except AssertionError as e:
+        assert (
+            str(e)
+            == "`to_pyg_data` only supports single feature homogeneous graph."
+        )

tests/python/pytorch/graphbolt/test_utils.py

@@ -0,0 +1,308 @@
+import re
+import unittest
+
+from functools import partial
+
+import backend as F
+
+import dgl
+import dgl.graphbolt as gb
+import pytest
+import torch
+
+
+def test_add_reverse_edges_homo():
+    edges = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]]).T
+    combined_edges = gb.add_reverse_edges(edges)
+    assert torch.equal(
+        combined_edges,
+        torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7], [4, 5, 6, 7, 0, 1, 2, 3]]).T,
+    )
+    # Tensor with uncorrect dimensions.
+    edges = torch.tensor([0, 1, 2, 3])
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Only tensor with shape N*2 is supported now, but got torch.Size([4])."
+        ),
+    ):
+        gb.add_reverse_edges(edges)
+
+
+def test_add_reverse_edges_hetero():
+    # reverse_etype doesn't exist in original etypes.
+    edges = {"n1:e1:n2": torch.tensor([[0, 1, 2], [4, 5, 6]]).T}
+    reverse_etype_mapping = {"n1:e1:n2": "n2:e2:n1"}
+    combined_edges = gb.add_reverse_edges(edges, reverse_etype_mapping)
+    assert torch.equal(
+        combined_edges["n1:e1:n2"], torch.tensor([[0, 1, 2], [4, 5, 6]]).T
+    )
+    assert torch.equal(
+        combined_edges["n2:e2:n1"], torch.tensor([[4, 5, 6], [0, 1, 2]]).T
+    )
+    # reverse_etype exists in original etypes.
+    edges = {
+        "n1:e1:n2": torch.tensor([[0, 1, 2], [4, 5, 6]]).T,
+        "n2:e2:n1": torch.tensor([[7, 8, 9], [10, 11, 12]]).T,
+    }
+    reverse_etype_mapping = {"n1:e1:n2": "n2:e2:n1"}
+    combined_edges = gb.add_reverse_edges(edges, reverse_etype_mapping)
+    assert torch.equal(
+        combined_edges["n1:e1:n2"], torch.tensor([[0, 1, 2], [4, 5, 6]]).T
+    )
+    assert torch.equal(
+        combined_edges["n2:e2:n1"],
+        torch.tensor([[7, 8, 9, 4, 5, 6], [10, 11, 12, 0, 1, 2]]).T,
+    )
+    # Tensor with uncorrect dimensions.
+    edges = {
+        "n1:e1:n2": torch.tensor([0, 1, 2]),
+        "n2:e2:n1": torch.tensor([7, 8, 9]),
+    }
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            "Only tensor with shape N*2 is supported now, but got torch.Size([3])."
+        ),
+    ):
+        gb.add_reverse_edges(edges, reverse_etype_mapping)
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Fails due to different result on the GPU.",
+)
+@pytest.mark.parametrize("use_datapipe", [False, True])
+def test_exclude_seed_edges_homo_cpu(use_datapipe):
+    graph = dgl.graph(([5, 0, 6, 7, 2, 2, 4], [0, 1, 2, 2, 3, 4, 4]))
+    graph = gb.from_dglgraph(graph, True).to(F.ctx())
+    items = torch.LongTensor([[0, 3], [4, 4]])
+    names = "seeds"
+    itemset = gb.ItemSet(items, names=names)
+    datapipe = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    sampler = gb.NeighborSampler
+    datapipe = sampler(datapipe, graph, fanouts)
+    if use_datapipe:
+        datapipe = datapipe.exclude_seed_edges()
+    else:
+        datapipe = datapipe.transform(partial(gb.exclude_seed_edges))
+    original_row_node_ids = [
+        torch.tensor([0, 3, 4, 5, 2, 6, 7]).to(F.ctx()),
+        torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
+    ]
+    compacted_indices = [
+        torch.tensor([3, 4, 4, 5, 6]).to(F.ctx()),
+        torch.tensor([3, 4, 4]).to(F.ctx()),
+    ]
+    indptr = [
+        torch.tensor([0, 1, 2, 3, 3, 5]).to(F.ctx()),
+        torch.tensor([0, 1, 2, 3]).to(F.ctx()),
+    ]
+    seeds = [
+        torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
+        torch.tensor([0, 3, 4]).to(F.ctx()),
+    ]
+    for data in datapipe:
+        for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
+            assert torch.equal(
+                sampled_subgraph.original_row_node_ids,
+                original_row_node_ids[step],
+            )
+            assert torch.equal(
+                sampled_subgraph.sampled_csc.indices, compacted_indices[step]
+            )
+            assert torch.equal(
+                sampled_subgraph.sampled_csc.indptr, indptr[step]
+            )
+            assert torch.equal(
+                sampled_subgraph.original_column_node_ids, seeds[step]
+            )
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Fails due to different result on the CPU.",
+)
+@pytest.mark.parametrize("use_datapipe", [False, True])
+@pytest.mark.parametrize("async_op", [False, True])
+def test_exclude_seed_edges_gpu(use_datapipe, async_op):
+    graph = dgl.graph(([5, 0, 7, 7, 2, 4], [0, 1, 2, 2, 3, 4]))
+    graph = gb.from_dglgraph(graph, is_homogeneous=True).to(F.ctx())
+    items = torch.LongTensor([[0, 3], [4, 4]])
+    names = "seeds"
+    itemset = gb.ItemSet(items, names=names)
+    datapipe = gb.ItemSampler(itemset, batch_size=4).copy_to(F.ctx())
+    num_layer = 2
+    fanouts = [torch.LongTensor([-1]) for _ in range(num_layer)]
+    sampler = gb.NeighborSampler
+    datapipe = sampler(
+        datapipe,
+        graph,
+        fanouts,
+        deduplicate=True,
+    )
+    if use_datapipe:
+        datapipe = datapipe.exclude_seed_edges(asynchronous=async_op)
+    else:
+        datapipe = datapipe.transform(
+            partial(gb.exclude_seed_edges, async_op=async_op)
+        )
+    if torch.cuda.get_device_capability()[0] < 7:
+        original_row_node_ids = [
+            torch.tensor([0, 3, 4, 2, 5, 7]).to(F.ctx()),
+            torch.tensor([0, 3, 4, 2, 5]).to(F.ctx()),
+        ]
+        compacted_indices = [
+            torch.tensor([4, 3, 5, 5]).to(F.ctx()),
+            torch.tensor([4, 3]).to(F.ctx()),
+        ]
+        indptr = [
+            torch.tensor([0, 1, 2, 2, 5, 5]).to(F.ctx()),
+            torch.tensor([0, 1, 2, 2]).to(F.ctx()),
+        ]
+        seeds = [
+            torch.tensor([0, 3, 4, 2, 5]).to(F.ctx()),
+            torch.tensor([0, 3, 4]).to(F.ctx()),
+        ]
+    else:
+        original_row_node_ids = [
+            torch.tensor([0, 3, 4, 5, 2, 7]).to(F.ctx()),
+            torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
+        ]
+        compacted_indices = [
+            torch.tensor([3, 4, 5, 5]).to(F.ctx()),
+            torch.tensor([3, 4]).to(F.ctx()),
+        ]
+        indptr = [
+            torch.tensor([0, 1, 2, 2, 2, 4]).to(F.ctx()),
+            torch.tensor([0, 1, 2, 2]).to(F.ctx()),
+        ]
+        seeds = [
+            torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
+            torch.tensor([0, 3, 4]).to(F.ctx()),
+        ]
+    for data in datapipe:
+        for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
+            if async_op and not use_datapipe:
+                sampled_subgraph = sampled_subgraph.wait()
+            assert torch.equal(
+                sampled_subgraph.original_row_node_ids,
+                original_row_node_ids[step],
+            )
+            assert torch.equal(
+                (sampled_subgraph.sampled_csc.indices), compacted_indices[step]
+            )
+            assert torch.equal(
+                sampled_subgraph.sampled_csc.indptr, indptr[step]
+            )
+            assert torch.equal(
+                sampled_subgraph.original_column_node_ids, seeds[step]
+            )
+
+
+def get_hetero_graph():
+    # COO graph:
+    # [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
+    # [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
+    # [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
+    # num_nodes = 5, num_n1 = 2, num_n2 = 3
+    ntypes = {"n1": 0, "n2": 1}
+    etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
+    indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
+    indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
+    type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
+    node_type_offset = torch.LongTensor([0, 2, 5])
+    return gb.fused_csc_sampling_graph(
+        indptr,
+        indices,
+        node_type_offset=node_type_offset,
+        type_per_edge=type_per_edge,
+        node_type_to_id=ntypes,
+        edge_type_to_id=etypes,
+    )
+
+
+def test_exclude_seed_edges_hetero():
+    graph = get_hetero_graph().to(F.ctx())
+    itemset = gb.HeteroItemSet(
+        {"n1:e1:n2": gb.ItemSet(torch.tensor([[0, 1]]), names="seeds")}
+    )
+    item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
+    num_layer = 2
+    fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
+    Sampler = gb.NeighborSampler
+    datapipe = Sampler(
+        item_sampler,
+        graph,
+        fanouts,
+        deduplicate=True,
+    )
+    datapipe = datapipe.transform(partial(gb.exclude_seed_edges))
+    csc_formats = [
+        {
+            "n1:e1:n2": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1, 3, 5]),
+                indices=torch.tensor([1, 0, 1, 0, 1]),
+            ),
+            "n2:e2:n1": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 2, 4]),
+                indices=torch.tensor([1, 2, 1, 0]),
+            ),
+        },
+        {
+            "n1:e1:n2": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 1]),
+                indices=torch.tensor([1]),
+            ),
+            "n2:e2:n1": gb.CSCFormatBase(
+                indptr=torch.tensor([0, 2]),
+                indices=torch.tensor([1, 2], dtype=torch.int64),
+            ),
+        },
+    ]
+    original_column_node_ids = [
+        {
+            "n1": torch.tensor([0, 1]),
+            "n2": torch.tensor([0, 1, 2]),
+        },
+        {
+            "n1": torch.tensor([0]),
+            "n2": torch.tensor([1]),
+        },
+    ]
+    original_row_node_ids = [
+        {
+            "n1": torch.tensor([0, 1]),
+            "n2": torch.tensor([0, 1, 2]),
+        },
+        {
+            "n1": torch.tensor([0, 1]),
+            "n2": torch.tensor([0, 1, 2]),
+        },
+    ]
+    for data in datapipe:
+        for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
+            for ntype in ["n1", "n2"]:
+                assert torch.equal(
+                    torch.sort(sampled_subgraph.original_row_node_ids[ntype])[
+                        0
+                    ],
+                    original_row_node_ids[step][ntype].to(F.ctx()),
+                )
+                assert torch.equal(
+                    torch.sort(
+                        sampled_subgraph.original_column_node_ids[ntype]
+                    )[0],
+                    original_column_node_ids[step][ntype].to(F.ctx()),
+                )
+            for etype in ["n1:e1:n2", "n2:e2:n1"]:
+                assert torch.equal(
+                    sampled_subgraph.sampled_csc[etype].indices,
+                    csc_formats[step][etype].indices.to(F.ctx()),
+                )
+                assert torch.equal(
+                    sampled_subgraph.sampled_csc[etype].indptr,
+                    csc_formats[step][etype].indptr.to(F.ctx()),
+                )

tests/python/pytorch/ip_config.txt

@@ -0,0 +1,4 @@
+0 127.0.0.1 40050
+1 127.0.0.1 40051
+2 127.0.0.1 40052
+3 127.0.0.1 40053

tests/python/pytorch/mpops/test_edgewise.py

@@ -0,0 +1,209 @@
+import random
+
+import backend as F
+
+import dgl
+import numpy as np
+import pytest
+import torch
+from utils import parametrize_idtype
+
+random.seed(42)
+np.random.seed(42)
+dgl.seed(42)
+torch.random.manual_seed(42)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_u(idtype, feat_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn(
+        (g.num_nodes(),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_u(g, x)
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = g.edges()
+    y_true = x[u.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_u_hetero(idtype, feat_size):
+    ctx = F.ctx()
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): ([0, 1, 2], [2, 3, 4]),
+            ("user", "like", "movie"): ([3, 3, 1, 2], [0, 0, 1, 1]),
+        }
+    )
+
+    hg = hg.astype(idtype).to(ctx)
+    x = torch.randn(
+        (hg.num_nodes("user"),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_u(hg, x, etype="like")
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = hg.edges(etype="like")
+    y_true = x[u.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_v(idtype, feat_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn(
+        (g.num_nodes(),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_v(g, x)
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = g.edges()
+    y_true = x[v.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_v_hetero(idtype, feat_size):
+    ctx = F.ctx()
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): ([0, 1, 2], [2, 3, 4]),
+            ("user", "like", "movie"): ([3, 3, 1, 2], [0, 0, 1, 1]),
+        }
+    )
+
+    hg = hg.astype(idtype).to(ctx)
+    x = torch.randn(
+        (hg.num_nodes("movie"),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_v(hg, x, etype="like")
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = hg.edges(etype="like")
+    y_true = x[v.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+binary_arg_sizes = [
+    ((5,), (5,)),
+    ((5,), ()),
+    ((), (5,)),
+    ((1, 3, 3), (4, 1, 3)),
+    ((3, 3), (4, 1, 3)),
+    ((4, 1, 3), (3, 3)),
+]
+
+dot_arg_sizes = [
+    ((5,), (5,)),
+    ((1, 3, 3), (4, 1, 3)),
+    ((3, 3), (4, 1, 3)),
+    ((4, 1, 3), (3, 3)),
+]
+
+ops = ["add", "sub", "mul", "div"]
+
+
+def pad_shape(x, y, x_size, y_size):
+    xy_size = torch.broadcast_shapes(x_size, y_size)
+    new_x_size = (1,) * (len(xy_size) - len(x_size)) + x_size
+    new_y_size = (1,) * (len(xy_size) - len(y_size)) + y_size
+    new_x = x.view(-1, *new_x_size)
+    new_y = y.view(-1, *new_y_size)
+    return new_x, new_y
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("op", ops)
+@pytest.mark.parametrize("x_size,y_size", binary_arg_sizes)
+def test_u_op_v(idtype, op, x_size, y_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn((g.num_nodes(),) + x_size, requires_grad=True, device=ctx)
+    y = torch.randn((g.num_nodes(),) + y_size, requires_grad=True, device=ctx)
+
+    f_dgl = getattr(dgl, f"u_{op}_v")
+    z = f_dgl(g, x, y)
+    z.sum().backward()
+    x_grad = x.grad
+    y_grad = y.grad
+
+    x_grad.zero_()
+    y_grad.zero_()
+    u, v = g.edges()
+    f_torch = getattr(torch, op)
+    x_u, y_v = pad_shape(x[u.long()], y[v.long()], x_size, y_size)
+    z_true = f_torch(x_u, y_v)
+    z_true.sum().backward()
+    x_grad_true = x.grad
+    y_grad_true = y.grad
+
+    assert torch.allclose(z, z_true)
+    assert torch.allclose(x_grad, x_grad_true)
+    assert torch.allclose(y_grad, y_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("x_size,y_size", dot_arg_sizes)
+def test_u_dot_v(idtype, x_size, y_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn((g.num_nodes(),) + x_size, requires_grad=True, device=ctx)
+    y = torch.randn((g.num_nodes(),) + y_size, requires_grad=True, device=ctx)
+
+    z = dgl.u_dot_v(g, x, y)
+    z.sum().backward()
+    x_grad = x.grad
+    y_grad = y.grad
+
+    x_grad.zero_()
+    y_grad.zero_()
+    u, v = g.edges()
+    x_u, y_v = pad_shape(x[u.long()], y[v.long()], x_size, y_size)
+    z_true = (x_u * y_v).sum(-1).unsqueeze(-1)
+    z_true.sum().backward()
+    x_grad_true = x.grad
+    y_grad_true = y.grad
+
+    assert torch.allclose(z, z_true, atol=1e-4, rtol=1e-4)
+    assert torch.allclose(x_grad, x_grad_true)
+    assert torch.allclose(y_grad, y_grad_true)

tests/python/pytorch/nn/conv/test_gatedgcnconv.py

@@ -0,0 +1,26 @@
+import io
+
+import backend as F
+
+import dgl.nn.pytorch as nn
+import pytest
+from utils import parametrize_idtype
+from utils.graph_cases import get_cases
+
+tmp_buffer = io.BytesIO()
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["zero-degree"]))
+def test_gatedgcn_conv(g, idtype):
+    ctx = F.ctx()
+    g = g.astype(idtype).to(ctx)
+    gatedgcnconv = nn.GatedGCNConv(10, 10, 5)
+    feat = F.randn((g.num_nodes(), 10))
+    efeat = F.randn((g.num_edges(), 10))
+    gatedgcnconv = gatedgcnconv.to(ctx)
+
+    h, edge_h = gatedgcnconv(g, feat, efeat)
+    # current we only do shape check
+    assert h.shape == (g.number_of_dst_nodes(), 5)
+    assert edge_h.shape == (g.number_of_edges(), 5)

tests/python/pytorch/nn/test_sparse_emb.py

@@ -0,0 +1,133 @@
+import multiprocessing as mp
+import os
+import unittest
+
+import backend as F
+import pytest
+import torch as th
+
+from dgl.nn import NodeEmbedding
+from dgl.optim import SparseAdam
+
+
+def initializer(emb):
+    th.manual_seed(0)
+    emb.uniform_(-1.0, 1.0)
+    return emb
+
+
+def check_all_set_all_get_emb(device, init_emb):
+    num_embs = init_emb.shape[0]
+    emb_dim = init_emb.shape[1]
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test", device=device)
+    dgl_emb.all_set_embedding(init_emb)
+
+    out_emb = dgl_emb.all_get_embedding()
+    assert F.allclose(init_emb, out_emb)
+
+
+def check_all_set_all_get_optm_state(
+    device, state_step, state_mem, state_power
+):
+    num_embs = state_mem.shape[0]
+    emb_dim = state_mem.shape[1]
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test", device=device)
+    optm = SparseAdam(params=[dgl_emb], lr=0.01)
+
+    dgl_emb._all_set_optm_state((state_step, state_mem, state_power))
+
+    out_step, out_mem, out_power = dgl_emb._all_get_optm_state()
+
+    assert F.allclose(state_step, out_step)
+    assert F.allclose(state_mem, out_mem)
+    assert F.allclose(state_power, out_power)
+
+
+def start_sparse_worker(rank, world_size, test, args):
+    print("start sparse worker {}".format(rank))
+    dist_init_method = "tcp://{master_ip}:{master_port}".format(
+        master_ip="127.0.0.1", master_port="12345"
+    )
+    backend = "gloo"
+    device = F.ctx()
+    if device.type == "cuda":
+        device = th.device(rank)
+        th.cuda.set_device(device)
+    th.distributed.init_process_group(
+        backend=backend,
+        init_method=dist_init_method,
+        world_size=world_size,
+        rank=rank,
+    )
+
+    test(device, *args)
+    th.distributed.barrier()
+    th.distributed.destroy_process_group()
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("num_workers", [1, 2, 3])
+def test_multiprocess_sparse_emb_get_set(num_workers):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    worker_list = []
+
+    init_emb = th.rand([1000, 8])
+
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_sparse_worker,
+            args=(i, num_workers, check_all_set_all_get_emb, (init_emb,)),
+        )
+        p.start()
+        worker_list.append(p)
+
+    for p in worker_list:
+        p.join()
+    for p in worker_list:
+        assert p.exitcode == 0
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("num_workers", [1, 2, 3])
+def test_multiprocess_sparse_emb_get_set_optm_state(num_workers):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    worker_list = []
+
+    num_embs, emb_dim = 1000, 8
+    state_step = th.randint(1000, (num_embs,))
+    state_mem = th.rand((num_embs, emb_dim))
+    state_power = th.rand((num_embs, emb_dim))
+
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_sparse_worker,
+            args=(
+                i,
+                num_workers,
+                check_all_set_all_get_optm_state,
+                (state_step, state_mem, state_power),
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+
+    for p in worker_list:
+        p.join()
+    for p in worker_list:
+        assert p.exitcode == 0
+
+
+if __name__ == "__main__":
+    # test_multiprocess_sparse_emb_get_set(1)
+    # test_multiprocess_sparse_emb_get_set(2)
+    # test_multiprocess_sparse_emb_get_set(3)
+
+    test_multiprocess_sparse_emb_get_set_optm_state(1)
+    # test_multiprocess_sparse_emb_get_set_optm_state(2)
+    # test_multiprocess_sparse_emb_get_set_optm_state(3)

tests/python/pytorch/optim/test_optim.py

@@ -0,0 +1,698 @@
+import os
+import unittest
+
+import backend as F
+import pytest
+import torch as th
+import torch.multiprocessing as mp
+
+from dgl.nn import NodeEmbedding
+from dgl.optim import SparseAdagrad, SparseAdam
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
+def test_sparse_adam(emb_dim):
+    num_embs = 10
+    device = F.ctx()
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test")
+    torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
+    th.manual_seed(0)
+    th.nn.init.uniform_(dgl_emb.weight, 0, 1.0)
+
+    dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01)
+    torch_adam = th.optim.SparseAdam(list(torch_emb.parameters()), lr=0.01)
+
+    # first step
+    idx = th.randint(0, num_embs, size=(4,))
+    dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
+    torch_value = torch_emb(idx)
+    labels = th.zeros((4,)).long()
+    print("dgl_value = {}".format(dgl_value))
+    print("labels = {}".format(labels))
+
+    dgl_adam.zero_grad()
+    torch_adam.zero_grad()
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    dgl_loss.backward()
+    torch_loss.backward()
+
+    dgl_adam.step()
+    torch_adam.step()
+    assert F.allclose(dgl_emb.weight, torch_emb.weight)
+
+    # Can not test second step
+    # Pytorch sparseAdam maintains a global step
+    # DGL sparseAdam use a per embedding step
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("use_uva", [False, True, None])
+@pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
+def test_sparse_adam_uva(use_uva, emb_dim):
+    if F.ctx().type == "cpu" and use_uva == True:
+        # we want to only test values of False and None when not using GPU
+        pytest.skip("UVA cannot be used without GPUs.")
+
+    num_embs = 10
+    device = F.ctx()
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test_uva{}".format(use_uva))
+    torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
+    th.manual_seed(0)
+    th.nn.init.uniform_(dgl_emb.weight, 0, 1.0)
+
+    dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01, use_uva=use_uva)
+    torch_adam = th.optim.SparseAdam(list(torch_emb.parameters()), lr=0.01)
+
+    # first step
+    idx = th.randint(0, num_embs, size=(4,))
+    dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
+    torch_value = torch_emb(idx)
+    labels = th.zeros((4,)).long()
+
+    dgl_adam.zero_grad()
+    torch_adam.zero_grad()
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    dgl_loss.backward()
+    torch_loss.backward()
+
+    dgl_adam.step()
+    torch_adam.step()
+    assert F.allclose(dgl_emb.weight, torch_emb.weight)
+
+    # Can not test second step
+    # Pytorch sparseAdam maintains a global step
+    # DGL sparseAdam use a per embedding step
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@pytest.mark.parametrize("dtype", [th.float32, th.float16])
+@pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
+def test_sparse_adam_dtype(dtype, emb_dim):
+    num_embs = 10
+    device = F.ctx()
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test_dtype{}".format(dtype))
+    torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
+    th.manual_seed(0)
+    th.nn.init.uniform_(dgl_emb.weight, 0, 1.0)
+
+    dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01, dtype=dtype)
+    torch_adam = th.optim.SparseAdam(list(torch_emb.parameters()), lr=0.01)
+
+    # first step
+    idx = th.randint(0, num_embs, size=(4,))
+    dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
+    torch_value = torch_emb(idx)
+    labels = th.zeros((4,)).long()
+
+    dgl_adam.zero_grad()
+    torch_adam.zero_grad()
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    dgl_loss.backward()
+    torch_loss.backward()
+
+    dgl_adam.step()
+    torch_adam.step()
+    assert F.allclose(dgl_emb.weight, torch_emb.weight)
+
+    # Can not test second step
+    # Pytorch sparseAdam maintains a global step
+    # DGL sparseAdam use a per embedding step
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+def test_sparse_adam_zero_step():
+    num_embs = 10
+    emb_dim = 4
+    device = F.ctx()
+    dgl_emb = NodeEmbedding(num_embs, emb_dim, "test")
+    torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    dgl_emb_zero = NodeEmbedding(num_embs, emb_dim, "test2")
+    torch_emb_zero = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
+    th.nn.init.uniform_(torch_emb_zero.weight, 0, 1.0)
+    th.manual_seed(0)
+    th.nn.init.uniform_(dgl_emb.weight, 0, 1.0)
+    th.nn.init.uniform_(dgl_emb_zero.weight, 0, 1.0)
+
+    dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
+    torch_adam = th.optim.SparseAdam(
+        list(torch_emb.parameters()) + list(torch_emb_zero.parameters()),
+        lr=0.01,
+    )
+
+    # first step
+    idx = th.randint(0, num_embs, size=(4,))
+    dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
+    torch_value = torch_emb(idx)
+    labels = th.ones((4,)).long()
+
+    dgl_adam.zero_grad()
+    torch_adam.zero_grad()
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    dgl_loss.backward()
+    torch_loss.backward()
+
+    dgl_adam.step()
+    torch_adam.step()
+    assert F.allclose(dgl_emb.weight, torch_emb.weight)
+
+
+def initializer(emb):
+    th.manual_seed(0)
+    emb.uniform_(-1.0, 1.0)
+    return emb
+
+
+def start_sparse_adam_worker(
+    rank,
+    device,
+    world_size,
+    weight,
+    tensor_dev="cpu",
+    has_zero_grad=False,
+    backend="gloo",
+    num_embs=128,
+    emb_dim=10,
+    zero_comm=True,
+):
+    print("start sparse worker for adam {}".format(rank))
+    dist_init_method = "tcp://{master_ip}:{master_port}".format(
+        master_ip="127.0.0.1", master_port="12345"
+    )
+
+    if device.type == "cuda":
+        th.cuda.set_device(device)
+
+    th.distributed.init_process_group(
+        backend=backend,
+        init_method=dist_init_method,
+        world_size=world_size,
+        rank=rank,
+    )
+
+    init_weight = th.empty((num_embs, emb_dim))
+    th.manual_seed(0)
+    th.nn.init.uniform_(init_weight, -1.0, 1.0)
+    dgl_emb = NodeEmbedding(
+        num_embs, emb_dim, "test", init_func=initializer, device=tensor_dev
+    )
+    dgl_emb.all_set_embedding(init_weight)
+
+    if has_zero_grad:
+        dgl_emb_zero = NodeEmbedding(
+            num_embs, emb_dim, "zero", init_func=initializer, device=tensor_dev
+        )
+        dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
+    else:
+        dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01)
+
+    th.manual_seed(rank)
+    if zero_comm:
+        start = (num_embs // world_size) * rank
+        end = (num_embs // world_size) * (rank + 1)
+        idx = th.randint(start, end, size=(4,)).to(tensor_dev)
+    else:
+        idx = th.randint(0, num_embs, size=(4,)).to(tensor_dev)
+    dgl_value = dgl_emb(idx, device)
+    labels = th.ones((4,)).long().to(device)
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    dgl_adam.zero_grad()
+    dgl_loss.backward()
+    dgl_adam.step()
+    th.distributed.barrier()
+    dgl_weight = dgl_emb.all_get_embedding().detach()
+    after_step = dgl_emb(idx, device).cpu()
+
+    if rank == 0:
+        dgl_value = dgl_value.detach().cpu()
+        assert F.allclose(dgl_value, after_step) is False
+        weight[:] = dgl_weight[:]
+    th.distributed.barrier()
+
+
+def start_torch_adam_worker(
+    rank,
+    world_size,
+    weight,
+    has_zero_grad=False,
+    num_embs=128,
+    emb_dim=10,
+    zero_comm=True,
+):
+    print("start sparse worker for adam {}".format(rank))
+    dist_init_method = "tcp://{master_ip}:{master_port}".format(
+        master_ip="127.0.0.1", master_port="12345"
+    )
+    backend = "gloo"
+
+    th.distributed.init_process_group(
+        backend=backend,
+        init_method=dist_init_method,
+        world_size=world_size,
+        rank=rank,
+    )
+
+    torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+    th.manual_seed(0)
+    th.nn.init.uniform_(torch_emb.weight, -1.0, 1.0)
+    torch_emb = th.nn.parallel.DistributedDataParallel(torch_emb)
+    if has_zero_grad:
+        torch_emb_zero = th.nn.Embedding(num_embs, emb_dim, sparse=True)
+        torch_emb_zero = torch_emb_zero.to(tensor_dev)
+        th.manual_seed(0)
+        th.nn.init.uniform_(torch_emb_zero.weight, -1.0, 1.0)
+        torch_emb_zero = th.nn.parallel.DistributedDataParallel(torch_emb_zero)
+        torch_adam = th.optim.SparseAdam(
+            list(torch_emb.module.parameters())
+            + list(torch_emb_zero.module.parameters()),
+            lr=0.01,
+        )
+    else:
+        torch_adam = th.optim.SparseAdam(
+            list(torch_emb.module.parameters()), lr=0.01
+        )
+
+    th.manual_seed(rank)
+    if zero_comm:
+        start = (num_embs // world_size) * rank
+        end = (num_embs // world_size) * (rank + 1)
+        idx = th.randint(start, end, size=(4,))
+    else:
+        idx = th.randint(0, num_embs, size=(4,))
+    labels = th.ones((4,)).long()
+    torch_value = torch_emb(idx)
+    torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
+    torch_adam.zero_grad()
+    torch_loss.backward()
+    torch_adam.step()
+    th.distributed.barrier()
+
+    if rank == 0:
+        weight[:] = torch_emb.module.weight.cpu()[:]
+    th.distributed.barrier()
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(F.ctx().type != "cpu", reason="cpu only test")
+@pytest.mark.parametrize("num_workers", [2, 4])
+def test_multiprocess_cpu_sparse_adam(num_workers):
+    backend = "gloo"
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = F.ctx()
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(
+                i,
+                device,
+                num_workers,
+                dgl_weight,
+                th.device("cpu"),
+                True,
+                backend,
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(i, num_workers, torch_weight, False),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(F.ctx().type == "cpu", reason="gpu only test")
+@pytest.mark.parametrize("num_workers", [2, 4, 8])
+@pytest.mark.parametrize("backend", ["nccl", "gloo"])
+@pytest.mark.parametrize("zero_comm", [True, False])
+def test_multiprocess_sparse_adam(num_workers, backend, zero_comm):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = F.ctx()
+        if device.type == "cuda":
+            # make sure each process has a unique GPU
+            device = th.device(i)
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(
+                i,
+                device,
+                num_workers,
+                dgl_weight,
+                th.device("cpu"),
+                True,
+                backend,
+                num_embs,
+                emb_dim,
+                zero_comm,
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(
+                i,
+                num_workers,
+                torch_weight,
+                False,
+                num_embs,
+                emb_dim,
+                zero_comm,
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(
+    F.ctx().type == "cpu", reason="cuda tensor is not supported for cpu"
+)
+@pytest.mark.parametrize("num_workers", [2, 4, 8])
+def test_multiprocess_sparse_adam_cuda_tensor(num_workers):
+    if F.ctx().type == "cpu":
+        pytest.skip("Do not test CPU")
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    backend = "nccl"
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = th.device(i)
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(i, device, num_workers, dgl_weight, device, False, backend),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(i, num_workers, torch_weight, False),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(F.ctx().type != "cpu", reason="cpu only test")
+@pytest.mark.parametrize("num_workers", [2, 4])
+def test_multiprocess_sparse_adam_cpu_zero_step(num_workers):
+    backend = "gloo"
+
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = F.ctx()
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(
+                i,
+                device,
+                num_workers,
+                dgl_weight,
+                th.device("cpu"),
+                True,
+                backend,
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(i, num_workers, torch_weight, False),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(F.ctx().type == "cpu", reason="gpu only test")
+@pytest.mark.parametrize("num_workers", [2, 4, 8])
+@pytest.mark.parametrize("backend", ["nccl", "gloo"])
+def test_multiprocess_sparse_adam_zero_step(num_workers, backend):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = F.ctx()
+        if device.type == "cuda":
+            # make sure each process has a unique GPU
+            device = th.device(i)
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(
+                i,
+                device,
+                num_workers,
+                dgl_weight,
+                th.device("cpu"),
+                True,
+                backend,
+            ),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(i, num_workers, torch_weight, False),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(
+    F.ctx().type == "cpu", reason="cuda tensor is not supported for cpu"
+)
+@pytest.mark.parametrize("num_workers", [2, 4, 8])
+def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    backend = "nccl"
+    worker_list = []
+    num_embs = 128
+    emb_dim = 10
+    dgl_weight = th.empty((num_embs, emb_dim))
+    ctx = mp.get_context("spawn")
+    for i in range(num_workers):
+        device = th.device(i)
+        p = ctx.Process(
+            target=start_sparse_adam_worker,
+            args=(i, device, num_workers, dgl_weight, device, True, backend),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    worker_list = []
+    torch_weight = th.empty((num_embs, emb_dim))
+    for i in range(num_workers):
+        p = ctx.Process(
+            target=start_torch_adam_worker,
+            args=(i, num_workers, torch_weight, False),
+        )
+        p.start()
+        worker_list.append(p)
+    for p in worker_list:
+        p.join()
+
+    assert F.allclose(dgl_weight, torch_weight)
+
+
+def start_sparse_adam_state_dict_worker(
+    rank,
+    world_size,
+    init_weight,
+    backend,
+    num_embs,
+    emb_dim,
+):
+    print("start sparse worker for adam {}".format(rank))
+    dist_init_method = "tcp://{master_ip}:{master_port}".format(
+        master_ip="127.0.0.1", master_port="12345"
+    )
+
+    device = th.device(f"cuda:{rank}")
+    th.cuda.set_device(device)
+    tensor_dev = device if backend == "nccl" else th.device("cpu")
+
+    th.distributed.init_process_group(
+        backend=backend,
+        init_method=dist_init_method,
+        world_size=world_size,
+        rank=rank,
+    )
+
+    th.manual_seed(0)
+    dgl_emb = NodeEmbedding(
+        num_embs, emb_dim, "test", init_func=initializer, device=tensor_dev
+    )
+    dgl_emb.all_set_embedding(init_weight)
+
+    dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01)
+
+    start = (num_embs // world_size) * rank
+    end = (num_embs // world_size) * (rank + 1)
+    th.manual_seed(rank)
+    idx = th.randint(start, end, size=(4,)).to(tensor_dev)
+    dgl_value = dgl_emb(idx, device)
+    labels = th.ones((4,)).long().to(device)
+    dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
+    dgl_adam.zero_grad()
+    dgl_loss.backward()
+    dgl_adam.step()
+    th.distributed.barrier()
+
+    worker_state_dict = [t.detach().clone() for t in dgl_emb.optm_state]
+    state_dict = dgl_adam.state_dict()
+    for t in dgl_emb.optm_state:
+        t.zero_()
+    dgl_adam.load_state_dict(state_dict)
+
+    for i, j in zip(worker_state_dict, dgl_emb.optm_state):
+        F.allclose(i, j)
+
+    th.distributed.barrier()
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+@unittest.skipIf(F.ctx().type == "cpu", reason="gpu only test")
+@pytest.mark.parametrize("num_workers", [1, 2, 4, 8])
+@pytest.mark.parametrize("backend", ["nccl", "gloo"])
+def test_multiprocess_sparse_adam_state_dict(num_workers, backend):
+    if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
+        pytest.skip("Not enough GPUs to run test.")
+
+    num_embs = 128
+    emb_dim = 10
+    init_weight = th.rand((num_embs, emb_dim))
+    mp.spawn(
+        start_sparse_adam_state_dict_worker,
+        (
+            num_workers,
+            init_weight,
+            backend,
+            num_embs,
+            emb_dim,
+        ),
+        nprocs=num_workers,
+    )
+
+
+if __name__ == "__main__":
+    test_sparse_adam(1)
+    test_sparse_adam(4)
+    test_sparse_adam(101)
+    test_sparse_adam(1024)
+    test_sparse_adam_zero_step()
+
+    test_multiprocess_cpu_sparse_adam(2)
+    test_multiprocess_cpu_sparse_adam(4)
+    test_multiprocess_cpu_sparse_adam(8)
+    test_multiprocess_sparse_adam_cpu_zero_step(2)
+
+    test_multiprocess_sparse_adam(2, backend="gloo")
+    test_multiprocess_sparse_adam(4, backend="gloo")
+    test_multiprocess_sparse_adam(8, backend="gloo")
+    test_multiprocess_sparse_adam(2, backend="nccl")
+    test_multiprocess_sparse_adam(4, backend="nccl")
+    test_multiprocess_sparse_adam(8, backend="nccl")
+
+    test_multiprocess_sparse_adam_zero_step(2, backend="gloo")
+    test_multiprocess_sparse_adam_zero_step(4, backend="nccl")
+
+    test_multiprocess_sparse_adam_cuda_tensor(2)
+    test_multiprocess_sparse_adam_zero_step_cuda_tensor(4)
+
+    test_multiprocess_sparse_adam_state_dict(2, "nccl")
+    test_multiprocess_sparse_adam_state_dict(2, "gloo")

tests/python/pytorch/sparse/__init__.py

@@ -0,0 +1 @@
+""" DGL sparse tests"""

tests/python/pytorch/sparse/test_broadcast.py

@@ -0,0 +1,45 @@
+import operator
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import sp_broadcast_v
+
+from .utils import rand_coo
+
+
+@pytest.mark.parametrize("shape", [(3, 4), (1, 5), (5, 1)])
+@pytest.mark.parametrize("nnz", [1, 4])
+@pytest.mark.parametrize("nz_dim", [None, 2])
+@pytest.mark.parametrize("op", ["add", "sub", "mul", "truediv"])
+def test_sp_broadcast_v(shape, nnz, nz_dim, op):
+    dev = F.ctx()
+    A = rand_coo(shape, nnz, dev, nz_dim)
+
+    v = torch.randn(A.shape[1], device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v[A.col]
+    else:
+        rhs = v[A.col].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)
+
+    v = torch.randn(1, A.shape[1], device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v.view(-1)[A.col]
+    else:
+        rhs = v.view(-1)[A.col].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)
+
+    v = torch.randn(A.shape[0], 1, device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v.view(-1)[A.row]
+    else:
+        rhs = v.view(-1)[A.row].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)

tests/python/pytorch/sparse/test_elementwise_op.py

@@ -0,0 +1,242 @@
+import operator
+
+import backend as F
+
+import dgl.sparse as dglsp
+import pytest
+import torch
+
+from dgl.sparse import diag, power
+
+
+@pytest.mark.parametrize("opname", ["add", "sub", "mul", "truediv"])
+def test_diag_op_diag(opname):
+    op = getattr(operator, opname)
+    ctx = F.ctx()
+    shape = (3, 4)
+    D1 = diag(torch.arange(1, 4).to(ctx), shape=shape)
+    D2 = diag(torch.arange(10, 13).to(ctx), shape=shape)
+    result = op(D1, D2)
+    assert torch.allclose(result.val, op(D1.val, D2.val), rtol=1e-4, atol=1e-4)
+    assert result.shape == D1.shape
+
+
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_diag_op_scalar(v_scalar):
+    ctx = F.ctx()
+    shape = (3, 4)
+    D1 = diag(torch.arange(1, 4).to(ctx), shape=shape)
+
+    # D * v
+    D2 = D1 * v_scalar
+    assert torch.allclose(D1.val * v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # v * D
+    D2 = v_scalar * D1
+    assert torch.allclose(v_scalar * D1.val, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # D / v
+    D2 = D1 / v_scalar
+    assert torch.allclose(D1.val / v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # D ^ v
+    D1 = diag(torch.arange(1, 4).to(ctx))
+    D2 = D1**v_scalar
+    assert torch.allclose(D1.val**v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # pow(D, v)
+    D2 = power(D1, v_scalar)
+    assert torch.allclose(D1.val**v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    with pytest.raises(TypeError):
+        D1 + v_scalar
+    with pytest.raises(TypeError):
+        v_scalar + D1
+
+    with pytest.raises(TypeError):
+        D1 - v_scalar
+    with pytest.raises(TypeError):
+        v_scalar - D1
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_coo(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    row = torch.tensor([1, 0]).to(ctx)
+    col = torch.tensor([0, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    B = dglsp.from_coo(row, col, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_csr(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    indptr = torch.tensor([0, 1, 2, 3]).to(ctx)
+    indices = torch.tensor([3, 0, 2]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    A = dglsp.from_csr(indptr, indices, val)
+
+    indptr = torch.tensor([0, 1, 2, 2]).to(ctx)
+    indices = torch.tensor([2, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    B = dglsp.from_csr(indptr, indices, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_csc(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    indptr = torch.tensor([0, 1, 1, 2, 3]).to(ctx)
+    indices = torch.tensor([1, 2, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    A = dglsp.from_csc(indptr, indices, val)
+
+    indptr = torch.tensor([0, 1, 1, 2, 2]).to(ctx)
+    indices = torch.tensor([1, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    B = dglsp.from_csc(indptr, indices, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_diag(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D1 = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+    D2 = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    C1 = op(D1, D2).to_dense()
+    C2 = func(D1, D2).to_dense()
+    dense_C = op(D1.to_dense(), D2.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(D1, 2)
+    with pytest.raises(TypeError):
+        op(2, D1)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+def test_add_sparse_diag(val_shape):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    sum1 = (A + D).to_dense()
+    sum2 = (D + A).to_dense()
+    sum3 = dglsp.add(A, D).to_dense()
+    sum4 = dglsp.add(D, A).to_dense()
+    dense_sum = A.to_dense() + D.to_dense()
+
+    assert torch.allclose(dense_sum, sum1)
+    assert torch.allclose(dense_sum, sum2)
+    assert torch.allclose(dense_sum, sum3)
+    assert torch.allclose(dense_sum, sum4)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+def test_sub_sparse_diag(val_shape):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    diff1 = (A - D).to_dense()
+    diff2 = (D - A).to_dense()
+    diff3 = dglsp.sub(A, D).to_dense()
+    diff4 = dglsp.sub(D, A).to_dense()
+    dense_diff = A.to_dense() - D.to_dense()
+
+    assert torch.allclose(dense_diff, diff1)
+    assert torch.allclose(dense_diff, -diff2)
+    assert torch.allclose(dense_diff, diff3)
+    assert torch.allclose(dense_diff, -diff4)
+
+
+@pytest.mark.parametrize("op", ["pow"])
+def test_error_op_sparse_diag(op):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    D = dglsp.diag(torch.randn(row.shape[0]).to(ctx), shape=shape)
+
+    with pytest.raises(TypeError):
+        getattr(operator, op)(A, D)
+    with pytest.raises(TypeError):
+        getattr(operator, op)(D, A)

tests/python/pytorch/sparse/test_elementwise_op_sp.py

@@ -0,0 +1,157 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import div, from_coo, mul, power, spmatrix, val_like
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+def all_close_sparse(A, row, col, val, shape):
+    rowA, colA = A.coo()
+    valA = A.val
+    assert torch.allclose(rowA, row)
+    assert torch.allclose(colA, col)
+    assert torch.allclose(valA, val)
+    assert A.shape == shape
+
+
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_muldiv_scalar(v_scalar):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(len(row)).to(ctx)
+    A1 = from_coo(row, col, val, shape=(3, 4))
+
+    # A * v
+    A2 = A1 * v_scalar
+    assert torch.allclose(A1.val * v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # v * A
+    A2 = v_scalar * A1
+    assert torch.allclose(A1.val * v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # A / v
+    A2 = A1 / v_scalar
+    assert torch.allclose(A1.val / v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # v / A
+    with pytest.raises(TypeError):
+        v_scalar / A1
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+def test_pow(val_shape):
+    # A ** v
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(val_shape).to(ctx)
+    A = from_coo(row, col, val, shape=(3, 4))
+    exponent = 2
+    A_new = A**exponent
+    assert torch.allclose(A_new.val, val**exponent)
+    assert A_new.shape == A.shape
+    new_row, new_col = A_new.coo()
+    assert torch.allclose(new_row, row)
+    assert torch.allclose(new_col, col)
+
+    # power(A, v)
+    A_new = power(A, exponent)
+    assert torch.allclose(A_new.val, val**exponent)
+    assert A_new.shape == A.shape
+    new_row, new_col = A_new.coo()
+    assert torch.allclose(new_row, row)
+    assert torch.allclose(new_col, col)
+
+
+@pytest.mark.parametrize("op", ["add", "sub"])
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_error_op_scalar(op, v_scalar):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(len(row)).to(ctx)
+    A = from_coo(row, col, val, shape=(3, 4))
+
+    with pytest.raises(TypeError):
+        A + v_scalar
+    with pytest.raises(TypeError):
+        v_scalar + A
+
+    with pytest.raises(TypeError):
+        A - v_scalar
+    with pytest.raises(TypeError):
+        v_scalar - A
+
+
+@pytest.mark.parametrize(
+    "create_func1", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize(
+    "create_func2", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (5, 3)])
+@pytest.mark.parametrize("nnz1", [5, 15])
+@pytest.mark.parametrize("nnz2", [1, 14])
+@pytest.mark.parametrize("nz_dim", [None, 3])
+def test_spspmul(create_func1, create_func2, shape, nnz1, nnz2, nz_dim):
+    dev = F.ctx()
+    A = create_func1(shape, nnz1, dev, nz_dim)
+    B = create_func2(shape, nnz2, dev, nz_dim)
+    C = mul(A, B)
+    assert not C.has_duplicate()
+
+    DA = sparse_matrix_to_dense(A)
+    DB = sparse_matrix_to_dense(B)
+    DC = DA * DB
+
+    grad = torch.rand_like(C.val)
+    C.val.backward(grad)
+    DC_grad = sparse_matrix_to_dense(val_like(C, grad))
+    DC.backward(DC_grad)
+
+    assert torch.allclose(sparse_matrix_to_dense(C), DC, atol=1e-05)
+    assert torch.allclose(
+        val_like(A, A.val.grad).to_dense(), DA.grad, atol=1e-05
+    )
+    assert torch.allclose(
+        val_like(B, B.val.grad).to_dense(), DB.grad, atol=1e-05
+    )
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (5, 3)])
+@pytest.mark.parametrize("nnz", [1, 14])
+@pytest.mark.parametrize("nz_dim", [None, 3])
+def test_spspdiv(create_func, nnz, shape, nz_dim):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev, nz_dim)
+
+    perm = torch.randperm(A.nnz, device=dev)
+    rperm = torch.argsort(perm)
+    B = spmatrix(A.indices()[:, perm], A.val[perm], A.shape)
+    C = div(A, B)
+    assert not C.has_duplicate()
+    assert torch.allclose(C.val, A.val / B.val[rperm], atol=1e-05)
+    assert torch.allclose(C.indices(), A.indices(), atol=1e-05)
+
+    # No need to test backward here, since it is handled by Pytorch

tests/python/pytorch/sparse/test_matmul.py

@@ -0,0 +1,218 @@
+import warnings
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import bspmm, diag, from_coo, val_like
+from dgl.sparse.matmul import matmul
+
+from .utils import (
+    clone_detach_and_grad,
+    dense_mask,
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_stride,
+    sparse_matrix_to_dense,
+    sparse_matrix_to_torch_sparse,
+)
+
+
+def _torch_sparse_mm(torch_A1, torch_A2):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return torch.sparse.mm(torch_A1, torch_A2)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 7), (5, 2)])
+@pytest.mark.parametrize("nnz", [1, 10])
+@pytest.mark.parametrize("out_dim", [None, 10])
+def test_spmm(create_func, shape, nnz, out_dim):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if out_dim is not None:
+        X = torch.randn(shape[1], out_dim, requires_grad=True, device=dev)
+    else:
+        X = torch.randn(shape[1], requires_grad=True, device=dev)
+
+    X = rand_stride(X)
+    sparse_result = matmul(A, X)
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = sparse_matrix_to_dense(A)
+    XX = clone_detach_and_grad(X)
+    dense_result = torch.matmul(adj, XX)
+    if out_dim is None:
+        dense_result = dense_result.view(-1)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result, atol=1e-05)
+    assert torch.allclose(X.grad, XX.grad, atol=1e-05)
+    assert torch.allclose(
+        dense_mask(adj.grad, A),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+        atol=1e-05,
+    )
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 7), (5, 2)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_bspmm(create_func, shape, nnz):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev, 2)
+    X = torch.randn(shape[1], 10, 2, requires_grad=True, device=dev)
+    X = rand_stride(X)
+
+    sparse_result = matmul(A, X)
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    XX = clone_detach_and_grad(X)
+    torch_A = A.to_dense().clone().detach().requires_grad_()
+    torch_result = torch_A.permute(2, 0, 1) @ XX.permute(2, 0, 1)
+
+    torch_result.backward(grad.permute(2, 0, 1))
+    assert torch.allclose(
+        sparse_result.permute(2, 0, 1), torch_result, atol=1e-05
+    )
+    assert torch.allclose(X.grad, XX.grad, atol=1e-05)
+    assert torch.allclose(
+        dense_mask(torch_A.grad, A),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+        atol=1e-05,
+    )
+
+
+@pytest.mark.parametrize("create_func1", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("create_func2", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape_n_m", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("shape_k", [3, 4])
+@pytest.mark.parametrize("nnz1", [1, 10])
+@pytest.mark.parametrize("nnz2", [1, 10])
+def test_spspmm(create_func1, create_func2, shape_n_m, shape_k, nnz1, nnz2):
+    dev = F.ctx()
+    shape1 = shape_n_m
+    shape2 = (shape_n_m[1], shape_k)
+    A1 = create_func1(shape1, nnz1, dev)
+    A2 = create_func2(shape2, nnz2, dev)
+    A3 = matmul(A1, A2)
+    grad = torch.randn_like(A3.val)
+    A3.val.backward(grad)
+
+    torch_A1 = sparse_matrix_to_torch_sparse(A1)
+    torch_A2 = sparse_matrix_to_torch_sparse(A2)
+    torch_A3 = _torch_sparse_mm(torch_A1, torch_A2)
+    torch_A3_grad = sparse_matrix_to_torch_sparse(A3, grad)
+    torch_A3.backward(torch_A3_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(A3.to_dense(), torch_A3.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A1, A1.val.grad).to_dense(),
+            torch_A1.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            val_like(A2, A2.val.grad).to_dense(),
+            torch_A2.grad.to_dense(),
+            atol=1e-05,
+        )
+
+
+def test_spspmm_duplicate():
+    dev = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A1 = from_coo(row, col, val, shape)
+
+    row = torch.tensor([1, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A2 = from_coo(row, col, val, shape)
+
+    try:
+        matmul(A1, A2)
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
+
+    try:
+        matmul(A2, A1)
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("sparse_shape", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_sparse_diag_mm(create_func, sparse_shape, nnz):
+    dev = F.ctx()
+    diag_shape = sparse_shape[1], sparse_shape[1]
+    A = create_func(sparse_shape, nnz, dev)
+    diag_val = torch.randn(sparse_shape[1], device=dev, requires_grad=True)
+    D = diag(diag_val, diag_shape)
+    B = matmul(A, D)
+    grad = torch.randn_like(B.val)
+    B.val.backward(grad)
+
+    torch_A = sparse_matrix_to_torch_sparse(A)
+    torch_D = sparse_matrix_to_torch_sparse(D)
+    torch_B = _torch_sparse_mm(torch_A, torch_D)
+    torch_B_grad = sparse_matrix_to_torch_sparse(B, grad)
+    torch_B.backward(torch_B_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(B.to_dense(), torch_B.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A, A.val.grad).to_dense(),
+            torch_A.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            diag(D.val.grad, D.shape).to_dense(),
+            torch_D.grad.to_dense(),
+            atol=1e-05,
+        )
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("sparse_shape", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_diag_sparse_mm(create_func, sparse_shape, nnz):
+    dev = F.ctx()
+    diag_shape = sparse_shape[0], sparse_shape[0]
+    A = create_func(sparse_shape, nnz, dev)
+    diag_val = torch.randn(sparse_shape[0], device=dev, requires_grad=True)
+    D = diag(diag_val, diag_shape)
+    B = matmul(D, A)
+    grad = torch.randn_like(B.val)
+    B.val.backward(grad)
+
+    torch_A = sparse_matrix_to_torch_sparse(A)
+    torch_D = sparse_matrix_to_torch_sparse(D)
+    torch_B = _torch_sparse_mm(torch_D, torch_A)
+    torch_B_grad = sparse_matrix_to_torch_sparse(B, grad)
+    torch_B.backward(torch_B_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(B.to_dense(), torch_B.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A, A.val.grad).to_dense(),
+            torch_A.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            diag(D.val.grad, D.shape).to_dense(),
+            torch_D.grad.to_dense(),
+            atol=1e-05,
+        )

tests/python/pytorch/sparse/test_matrix_op.py

@@ -0,0 +1,47 @@
+import backend as F
+import pytest
+import torch
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("dim", [0, 1])
+@pytest.mark.parametrize("index", [None, (1, 3), (4, 0, 2)])
+def test_compact(create_func, dim, index):
+    ctx = F.ctx()
+    shape = (5, 5)
+    ans_idx = []
+    if index is not None:
+        ans_idx = list(dict.fromkeys(index))
+        index = torch.tensor(index).to(ctx)
+
+    A = create_func(shape, 8, ctx)
+
+    A_compact, ret_id = A.compact(dim, index)
+    A_compact_dense = sparse_matrix_to_dense(A_compact)
+
+    A_dense = sparse_matrix_to_dense(A)
+
+    for i in range(shape[dim]):
+        if dim == 0:
+            row = list(A_dense[i, :].nonzero().reshape(-1))
+        else:
+            row = list(A_dense[:, i].nonzero().reshape(-1))
+        if (i not in list(ans_idx)) and len(row) > 0:
+            ans_idx.append(i)
+    if len(ans_idx):
+        ans_idx = torch.tensor(ans_idx).to(ctx)
+    A_dense_select = sparse_matrix_to_dense(A.index_select(dim, ans_idx))
+
+    assert A_compact_dense.shape == A_dense_select.shape
+    assert torch.allclose(A_compact_dense, A_dense_select)
+    assert torch.allclose(ans_idx, ret_id)

tests/python/pytorch/sparse/test_reduction.py

@@ -0,0 +1,160 @@
+import doctest
+import operator
+import sys
+
+import backend as F
+
+import dgl.sparse as dglsp
+import pytest
+import torch
+
+dgl_op_map = {
+    "sum": "sum",
+    "amin": "smin",
+    "amax": "smax",
+    "mean": "smean",
+    "prod": "sprod",
+}
+default_entry = {
+    "sum": 0,
+    "amin": float("inf"),
+    "amax": float("-inf"),
+    "mean": 0,
+    "prod": 1,
+}
+binary_op_map = {
+    "sum": operator.add,
+    "amin": torch.min,
+    "amax": torch.max,
+    "mean": operator.add,
+    "prod": operator.mul,
+}
+
+NUM_ROWS = 10
+NUM_COLS = 15
+
+
+def _coalesce_dense(row, col, val, nrows, ncols, op):
+    # Sparse matrix coalescing on a dense matrix.
+    #
+    # It is done by stacking every non-zero entry on an individual slice
+    # of an (nrows x ncols x nnz), that is, construct a tensor A with
+    # shape (nrows, ncols, len(val)) where
+    #
+    #     A[row[i], col[i], i] = val[i]
+    #
+    # and then reducing on the third "nnz" dimension.
+    #
+    # The mask matrix M has the same sparsity pattern as A with 1 being
+    # the non-zero entries.  This is used for division if the reduce
+    # operator is mean.
+    M = torch.zeros(NUM_ROWS, NUM_COLS, device=F.ctx())
+    A = torch.full(
+        (NUM_ROWS, NUM_COLS, 20) + val.shape[1:],
+        default_entry[op],
+        device=F.ctx(),
+        dtype=val.dtype,
+    )
+    A = torch.index_put(A, (row, col, torch.arange(20)), val)
+    for i in range(20):
+        M[row[i], col[i]] += 1
+    if op == "mean":
+        A = A.sum(2)
+    else:
+        A = getattr(A, op)(2)
+    M = M.view(NUM_ROWS, NUM_COLS, *([1] * (val.dim() - 1)))
+    return A, M
+
+
+# Add docstring tests of dglsp.reduction to unit tests
+@pytest.mark.parametrize(
+    "func", ["reduce", "sum", "smin", "smax", "sprod", "smean"]
+)
+def test_docstring(func):
+    globs = {"torch": torch, "dglsp": dglsp}
+    runner = doctest.DebugRunner()
+    finder = doctest.DocTestFinder()
+    obj = getattr(dglsp, func)
+    for test in finder.find(obj, func, globs=globs):
+        runner.run(test)
+
+
+@pytest.mark.parametrize("shape", [(20,), (20, 20)])
+@pytest.mark.parametrize("op", ["sum", "amin", "amax", "mean", "prod"])
+@pytest.mark.parametrize("use_reduce", [False, True])
+def test_reduce_all(shape, op, use_reduce):
+    row = torch.randint(0, NUM_ROWS, (20,), device=F.ctx())
+    col = torch.randint(0, NUM_COLS, (20,), device=F.ctx())
+    val = torch.randn(*shape, device=F.ctx())
+    val2 = val.clone()
+    val = val.requires_grad_()
+    val2 = val2.requires_grad_()
+    A = dglsp.from_coo(row, col, val, shape=(NUM_ROWS, NUM_COLS))
+
+    A2, M = _coalesce_dense(row, col, val2, NUM_ROWS, NUM_COLS, op)
+
+    if not use_reduce:
+        output = getattr(A, dgl_op_map[op])()
+    else:
+        output = A.reduce(rtype=dgl_op_map[op])
+
+    if op == "mean":
+        output2 = A2.sum((0, 1)) / M.sum()
+    elif op == "prod":
+        output2 = A2.prod(0).prod(0)  # prod() does not support tuple of dims
+    else:
+        output2 = getattr(A2, op)((0, 1))
+    assert (output - output2).abs().max() < 1e-4
+
+    head = torch.randn(*output.shape).to(val) if output.dim() > 0 else None
+    output.backward(head)
+    output2.backward(head)
+    assert (val.grad - val2.grad).abs().max() < 1e-4
+
+
+@pytest.mark.parametrize("shape", [(20,), (20, 20)])
+@pytest.mark.parametrize("dim", [0, 1])
+@pytest.mark.parametrize("empty_nnz", [False, True])
+@pytest.mark.parametrize("op", ["sum", "amin", "amax", "mean", "prod"])
+@pytest.mark.parametrize("use_reduce", [False, True])
+def test_reduce_along(shape, dim, empty_nnz, op, use_reduce):
+    row = torch.randint(0, NUM_ROWS, (20,), device=F.ctx())
+    col = torch.randint(0, NUM_COLS, (20,), device=F.ctx())
+    if dim == 0:
+        mask = torch.bincount(col, minlength=NUM_COLS) == 0
+    else:
+        mask = torch.bincount(row, minlength=NUM_ROWS) == 0
+    val = torch.randn(*shape, device=F.ctx())
+    val2 = val.clone()
+    val = val.requires_grad_()
+    val2 = val2.requires_grad_()
+
+    # empty_nnz controls whether at least one column or one row has no
+    # non-zero entry.
+    if empty_nnz:
+        row[row == 0] = 1
+        col[col == 0] = 1
+
+    A = dglsp.from_coo(row, col, val, shape=(NUM_ROWS, NUM_COLS))
+
+    A2, M = _coalesce_dense(row, col, val2, NUM_ROWS, NUM_COLS, op)
+
+    if not use_reduce:
+        output = getattr(A, dgl_op_map[op])(dim)
+    else:
+        output = A.reduce(dim=dim, rtype=dgl_op_map[op])
+
+    if op == "mean":
+        output2 = A2.sum(dim) / M.sum(dim)
+    else:
+        output2 = getattr(A2, op)(dim)
+    zero_entry_idx = (M.sum(dim) != 0).nonzero(as_tuple=True)[0]
+    output3 = torch.index_put(
+        torch.zeros_like(output2), (zero_entry_idx,), output2[zero_entry_idx]
+    )
+    assert (output - output3).abs().max() < 1e-4
+
+    head = torch.randn(*output.shape).to(val) if output.dim() > 0 else None
+    output.backward(head)
+    output3.backward(head)
+    assert (val.grad - val2.grad).abs().max() < 1e-4

tests/python/pytorch/sparse/test_sddmm.py

@@ -0,0 +1,92 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import bsddmm, sddmm
+
+from .utils import (
+    clone_detach_and_grad,
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_stride,
+)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("hidden", [1, 5])
+def test_sddmm(create_func, shape, nnz, hidden):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if hidden > 1:
+        B = torch.rand(shape[0], hidden, requires_grad=True, device=dev)
+        C = torch.rand(hidden, shape[1], requires_grad=True, device=dev)
+    else:
+        B = torch.rand(shape[0], requires_grad=True, device=dev)
+        C = torch.rand(shape[1], requires_grad=True, device=dev)
+
+    B = rand_stride(B)
+    C = rand_stride(C)
+
+    A_val_clone = clone_detach_and_grad(A.val)
+    dense_B = clone_detach_and_grad(B)
+    dense_C = clone_detach_and_grad(C)
+
+    sparse_result = sddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    if hidden == 1:
+        dense_result = dense_B.view(-1, 1) @ dense_C.view(1, -1)
+    else:
+        dense_result = dense_B @ dense_C
+
+    row, col = A.coo()
+    dense_val = dense_result[row, col] * A_val_clone
+    dense_val.backward(grad)
+
+    assert torch.allclose(dense_val, sparse_result.val, atol=1e-05)
+    assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
+    assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
+    assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("nz_dim", [2, 10])
+def test_bsddmm(create_func, shape, nnz, nz_dim):
+    dev = F.ctx()
+    hidden = 2
+    A = create_func(shape, nnz, dev, nz_dim)
+    B = torch.rand(shape[0], hidden, nz_dim, requires_grad=True, device=dev)
+    C = torch.rand(hidden, shape[1], nz_dim, requires_grad=True, device=dev)
+
+    B = rand_stride(B)
+    C = rand_stride(C)
+
+    A_val_clone = clone_detach_and_grad(A.val)
+    dense_B = clone_detach_and_grad(B)
+    dense_C = clone_detach_and_grad(C)
+
+    sparse_result = bsddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    dense_result = dense_B.permute(2, 0, 1) @ dense_C.permute(2, 0, 1)
+    dense_result = dense_result.permute(1, 2, 0)
+
+    row, col = A.coo()
+    dense_val = dense_result[row, col] * A_val_clone
+    dense_val.backward(grad)
+
+    assert torch.allclose(dense_val, sparse_result.val, atol=1e-05)
+    assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
+    assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
+    assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)

tests/python/pytorch/sparse/test_softmax.py

@@ -0,0 +1,43 @@
+import sys
+
+import backend as F
+
+import dgl
+import pytest
+import torch
+from dgl.sparse import from_coo, softmax
+
+
+@pytest.mark.parametrize("val_D", [None, 2])
+@pytest.mark.parametrize("csr", [True, False])
+@pytest.mark.parametrize("dim", [0, 1])
+def test_softmax(val_D, csr, dim):
+    dev = F.ctx()
+    row = torch.tensor([0, 0, 1, 1]).to(dev)
+    col = torch.tensor([0, 2, 1, 2]).to(dev)
+    nnz = len(row)
+    if val_D is None:
+        val = torch.randn(nnz).to(dev)
+    else:
+        val = torch.randn(nnz, val_D).to(dev)
+
+    val_sparse = val.clone().requires_grad_()
+    A = from_coo(row, col, val_sparse)
+
+    if csr:
+        # Test CSR
+        A.csr()
+
+    A_max = softmax(A, dim)
+    if dim == 1:
+        g = dgl.graph((col, row), num_nodes=max(A.shape))
+    else:
+        g = dgl.graph((row, col), num_nodes=max(A.shape))
+    val_g = val.clone().requires_grad_()
+    score = dgl.nn.functional.edge_softmax(g, val_g)
+    assert torch.allclose(A_max.val, score, atol=1e-05)
+
+    grad = torch.randn_like(score).to(dev)
+    A_max.val.backward(grad)
+    score.backward(grad)
+    assert torch.allclose(A.val.grad, val_g.grad, atol=1e-05)

tests/python/pytorch/sparse/test_sparse_matrix.py

@@ -0,0 +1,878 @@
+import unittest
+import warnings
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import (
+    diag,
+    from_coo,
+    from_csc,
+    from_csr,
+    from_torch_sparse,
+    identity,
+    to_torch_sparse_coo,
+    to_torch_sparse_csc,
+    to_torch_sparse_csr,
+    val_like,
+)
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+def _torch_sparse_csr_tensor(indptr, indices, val, torch_sparse_shape):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return torch.sparse_csr_tensor(indptr, indices, val, torch_sparse_shape)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_from_coo(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_from_csr(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_from_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)
+
+
+@pytest.mark.parametrize("val_shape", [(3), (3, 2)])
+def test_dense(val_shape):
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    val = torch.randn(val_shape).to(ctx)
+    A = from_coo(row, col, val)
+    A_dense = A.to_dense()
+
+    shape = A.shape + val.shape[1:]
+    mat = torch.zeros(shape, device=ctx)
+    mat[row, col] = val
+    assert torch.allclose(A_dense, mat)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    col = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+def _scatter_add(a, index, v=1):
+    index = index.tolist()
+    for i in index:
+        a[i] += v
+    return a
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csr(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_csr(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    row, sort_index = row.sort(stable=True)
+    col = col[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csc(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+
+    col = indices
+    col, sort_index = col.sort(stable=True)
+    row = row[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    indptr = _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("shape", [(3, 5), (5, 5), (5, 4)])
+def test_diag_conversions(shape):
+    n_rows, n_cols = shape
+    nnz = min(shape)
+    ctx = F.ctx()
+    val = torch.randn(nnz).to(ctx)
+    D = diag(val, shape)
+    row, col = D.coo()
+    assert torch.allclose(row, torch.arange(nnz).to(ctx))
+    assert torch.allclose(col, torch.arange(nnz).to(ctx))
+
+    indptr, indices, _ = D.csr()
+    exp_indptr = list(range(0, nnz + 1)) + [nnz] * (n_rows - nnz)
+    assert torch.allclose(indptr, torch.tensor(exp_indptr).to(ctx))
+    assert torch.allclose(indices, torch.arange(nnz).to(ctx))
+
+    indptr, indices, _ = D.csc()
+    exp_indptr = list(range(0, nnz + 1)) + [nnz] * (n_cols - nnz)
+    assert torch.allclose(indptr, torch.tensor(exp_indptr).to(ctx))
+    assert torch.allclose(indices, torch.arange(nnz).to(ctx))
+
+
+@pytest.mark.parametrize("val_shape", [(3), (3, 2)])
+@pytest.mark.parametrize("shape", [(3, 5), (5, 5)])
+def test_val_like(val_shape, shape):
+    def check_val_like(A, B):
+        assert A.shape == B.shape
+        assert A.nnz == B.nnz
+        assert torch.allclose(torch.stack(A.coo()), torch.stack(B.coo()))
+        assert A.val.device == B.val.device
+
+    ctx = F.ctx()
+
+    # COO
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    val = torch.randn(3).to(ctx)
+    coo_A = from_coo(row, col, val, shape)
+    new_val = torch.randn(val_shape).to(ctx)
+    coo_B = val_like(coo_A, new_val)
+    check_val_like(coo_A, coo_B)
+
+    # CSR
+    indptr, indices, _ = coo_A.csr()
+    csr_A = from_csr(indptr, indices, val, shape)
+    csr_B = val_like(csr_A, new_val)
+    check_val_like(csr_A, csr_B)
+
+    # CSC
+    indptr, indices, _ = coo_A.csc()
+    csc_A = from_csc(indptr, indices, val, shape)
+    csc_B = val_like(csc_A, new_val)
+    check_val_like(csc_A, csc_B)
+
+
+def test_coalesce():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    A = from_coo(row, col, val, (4, 4))
+
+    assert A.has_duplicate()
+
+    A_coalesced = A.coalesce()
+
+    assert A_coalesced.nnz == 4
+    assert A_coalesced.shape == (4, 4)
+    assert list(A_coalesced.row) == [0, 0, 1, 1]
+    assert list(A_coalesced.col) == [1, 2, 1, 2]
+    # Values of duplicate indices are added together.
+    assert list(A_coalesced.val) == [3, 3, 0, 4]
+    assert not A_coalesced.has_duplicate()
+
+
+def test_has_duplicate():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    shape = (4, 4)
+
+    # COO
+    coo_A = from_coo(row, col, val, shape)
+    assert coo_A.has_duplicate()
+
+    # CSR
+    indptr, indices, _ = coo_A.csr()
+    csr_A = from_csr(indptr, indices, val, shape)
+    assert csr_A.has_duplicate()
+
+    # CSC
+    indptr, indices, _ = coo_A.csc()
+    csc_A = from_csc(indptr, indices, val, shape)
+    assert csc_A.has_duplicate()
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (6, 4)])
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("select_dim", [0, 1])
+@pytest.mark.parametrize("index", [(0, 1, 3), (1, 2)])
+def test_index_select(create_func, shape, dense_dim, select_dim, index):
+    ctx = F.ctx()
+    A = create_func(shape, 20, ctx, dense_dim)
+    index = torch.tensor(index).to(ctx)
+    A_select = A.index_select(select_dim, index)
+
+    dense = sparse_matrix_to_dense(A)
+    dense_select = torch.index_select(dense, select_dim, index)
+
+    A_select_to_dense = sparse_matrix_to_dense(A_select)
+
+    assert A_select_to_dense.shape == dense_select.shape
+    assert torch.allclose(A_select_to_dense, dense_select)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (6, 4)])
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("select_dim", [0, 1])
+@pytest.mark.parametrize("rang", [slice(0, 2), slice(1, 3)])
+def test_range_select(create_func, shape, dense_dim, select_dim, rang):
+    ctx = F.ctx()
+    A = create_func(shape, 20, ctx, dense_dim)
+    A_select = A.range_select(select_dim, rang)
+
+    dense = sparse_matrix_to_dense(A)
+    if select_dim == 0:
+        dense_select = dense[rang, :]
+    else:
+        dense_select = dense[:, rang]
+
+    A_select_to_dense = sparse_matrix_to_dense(A_select)
+
+    assert A_select_to_dense.shape == dense_select.shape
+    assert torch.allclose(A_select_to_dense, dense_select)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("index", [(0, 1, 2, 3, 4), (0, 1, 3), (1, 1, 2)])
+@pytest.mark.parametrize("replace", [False, True])
+@pytest.mark.parametrize("bias", [False, True])
+def test_sample_rowwise(create_func, index, replace, bias):
+    ctx = F.ctx()
+    shape = (5, 5)
+    sample_dim = 0
+    sample_num = 3
+    A = create_func(shape, 10, ctx)
+    A = val_like(A, torch.abs(A.val))
+
+    index = torch.tensor(index).to(ctx)
+
+    A_sample = A.sample(sample_dim, sample_num, index, replace, bias)
+    A_dense = sparse_matrix_to_dense(A)
+    A_sample_to_dense = sparse_matrix_to_dense(A_sample)
+
+    ans_shape = (index.size(0), shape[1])
+    # Verify sample elements in origin rows
+    for i, row in enumerate(list(index)):
+        ans_ele = list(A_dense[row, :].nonzero().reshape(-1))
+        ret_ele = list(A_sample_to_dense[i, :].nonzero().reshape(-1))
+        for e in ret_ele:
+            assert e in ans_ele
+        if replace:
+            # The number of sample elements in one row should be equal to
+            # 'sample_num' if the row is not empty otherwise should be
+            # equal to 0.
+            assert list(A_sample.row).count(torch.tensor(i)) == (
+                sample_num if len(ans_ele) != 0 else 0
+            )
+        else:
+            assert len(ret_ele) == min(sample_num, len(ans_ele))
+
+    assert A_sample.shape == ans_shape
+    if not replace:
+        assert not A_sample.has_duplicate()
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("index", [(0, 1, 2, 3, 4), (0, 1, 3), (1, 1, 2)])
+@pytest.mark.parametrize("replace", [False, True])
+@pytest.mark.parametrize("bias", [False, True])
+def test_sample_columnwise(create_func, index, replace, bias):
+    ctx = F.ctx()
+    shape = (5, 5)
+    sample_dim = 1
+    sample_num = 3
+    A = create_func(shape, 10, ctx)
+    A = val_like(A, torch.abs(A.val))
+
+    index = torch.tensor(index).to(ctx)
+
+    A_sample = A.sample(sample_dim, sample_num, index, replace, bias)
+    A_dense = sparse_matrix_to_dense(A)
+    A_sample_to_dense = sparse_matrix_to_dense(A_sample)
+
+    ans_shape = (shape[0], index.size(0))
+    # Verify sample elements in origin columns
+    for i, col in enumerate(list(index)):
+        ans_ele = list(A_dense[:, col].nonzero().reshape(-1))
+        ret_ele = list(A_sample_to_dense[:, i].nonzero().reshape(-1))
+        for e in ret_ele:
+            assert e in ans_ele
+        if replace:
+            # The number of sample elements in one column should be equal to
+            # 'sample_num' if the column is not empty otherwise should be
+            # equal to 0.
+            assert list(A_sample.col).count(torch.tensor(i)) == (
+                sample_num if len(ans_ele) != 0 else 0
+            )
+        else:
+            assert len(ret_ele) == min(sample_num, len(ans_ele))
+
+    assert A_sample.shape == ans_shape
+    if not replace:
+        assert not A_sample.has_duplicate()
+
+
+def test_print():
+    ctx = F.ctx()
+
+    # basic
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([2, 1, 3]).to(ctx)
+    val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
+    A = from_coo(row, col, val)
+    expected = (
+        str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]]),
+             values=tensor([1., 1., 2.]),
+             shape=(4, 4), nnz=3)"""
+        )
+        if str(ctx) == "cpu"
+        else str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]], device='cuda:0'),
+             values=tensor([1., 1., 2.], device='cuda:0'),
+             shape=(4, 4), nnz=3)"""
+        )
+    )
+    assert str(A) == expected, print(A, expected)
+
+    # vector-shape non zero
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([2, 1, 3]).to(ctx)
+    val = torch.tensor(
+        [[1.3080, 1.5984], [-0.4126, 0.7250], [-0.5416, -0.7022]]
+    ).to(ctx)
+    A = from_coo(row, col, val)
+    expected = (
+        str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]]),
+             values=tensor([[ 1.3080,  1.5984],
+                            [-0.4126,  0.7250],
+                            [-0.5416, -0.7022]]),
+             shape=(4, 4), nnz=3, val_size=(2,))"""
+        )
+        if str(ctx) == "cpu"
+        else str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]], device='cuda:0'),
+             values=tensor([[ 1.3080,  1.5984],
+                            [-0.4126,  0.7250],
+                            [-0.5416, -0.7022]], device='cuda:0'),
+             shape=(4, 4), nnz=3, val_size=(2,))"""
+        )
+    )
+    assert str(A) == expected, print(A, expected)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Device conversions don't need to be tested on CPU.",
+)
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_to_device(device):
+    row = torch.tensor([1, 1, 2])
+    col = torch.tensor([1, 2, 0])
+    mat = from_coo(row, col, shape=(3, 4))
+
+    target_row = row.to(device)
+    target_col = col.to(device)
+    target_val = mat.val.to(device)
+
+    mat2 = mat.to(device=device)
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.row, target_row)
+    assert torch.allclose(mat2.col, target_col)
+    assert torch.allclose(mat2.val, target_val)
+
+    mat2 = getattr(mat, device)()
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.row, target_row)
+    assert torch.allclose(mat2.col, target_col)
+    assert torch.allclose(mat2.val, target_val)
+
+
+@pytest.mark.parametrize(
+    "dtype", [torch.float, torch.double, torch.int, torch.long]
+)
+def test_to_dtype(dtype):
+    row = torch.tensor([1, 1, 2])
+    col = torch.tensor([1, 2, 0])
+    mat = from_coo(row, col, shape=(3, 4))
+
+    target_val = mat.val.to(dtype=dtype)
+
+    mat2 = mat.to(dtype=dtype)
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.val, target_val)
+
+    func_name = {
+        torch.float: "float",
+        torch.double: "double",
+        torch.int: "int",
+        torch.long: "long",
+    }
+    mat2 = getattr(mat, func_name[dtype])()
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.val, target_val)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 2])
+@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("extra_shape", [(0, 1), (2, 1)])
+def test_sparse_matrix_transpose(dense_dim, row, col, extra_shape):
+    mat_shape = (max(row) + 1 + extra_shape[0], max(col) + 1 + extra_shape[1])
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, mat_shape).transpose()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == mat_shape[::-1]
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, col)
+    assert torch.allclose(mat_col, row)
+
+
+@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("nz_dim", [None, 2])
+@pytest.mark.parametrize("shape", [(5, 5), (6, 7)])
+def test_torch_sparse_coo_conversion(row, col, nz_dim, shape):
+    dev = F.ctx()
+    row = torch.tensor(row).to(dev)
+    col = torch.tensor(col).to(dev)
+    indices = torch.stack([row, col])
+    torch_sparse_shape = shape
+    val_shape = (row.shape[0],)
+    if nz_dim is not None:
+        torch_sparse_shape += (nz_dim,)
+        val_shape += (nz_dim,)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_coo = torch.sparse_coo_tensor(indices, val, torch_sparse_shape)
+    spmat = from_torch_sparse(torch_sparse_coo)
+
+    def _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo):
+        assert torch_sparse_coo.layout == torch.sparse_coo
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert (
+            spmat.indices().data_ptr() == torch_sparse_coo._indices().data_ptr()
+        )
+        assert spmat.val.data_ptr() == torch_sparse_coo._values().data_ptr()
+        assert spmat.shape == torch_sparse_coo.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+    torch_sparse_coo = to_torch_sparse_coo(spmat)
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+
+
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [(3, 5), (3, 7)])
+def test_torch_sparse_csr_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csr = _torch_sparse_csr_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csr)
+
+    def _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr):
+        indptr, indices, value_indices = spmat.csr()
+        assert torch_sparse_csr.layout == torch.sparse_csr
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csr.crow_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csr.col_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csr.values().data_ptr()
+        assert spmat.shape == torch_sparse_csr.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+    torch_sparse_csr = to_torch_sparse_csr(spmat)
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+
+
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [(8, 3), (5, 3)])
+def test_torch_sparse_csc_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csc = torch.sparse_csc_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csc)
+
+    def _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc):
+        indptr, indices, value_indices = spmat.csc()
+        assert torch_sparse_csc.layout == torch.sparse_csc
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csc.ccol_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csc.row_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csc.values().data_ptr()
+        assert spmat.shape == torch_sparse_csc.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)
+    torch_sparse_csc = to_torch_sparse_csc(spmat)
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)
+
+
+### Diag foramt related tests ###
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+@pytest.mark.parametrize("mat_shape", [None, (3, 5), (5, 3)])
+def test_diag(val_shape, mat_shape):
+    ctx = F.ctx()
+    # creation
+    val = torch.randn(val_shape).to(ctx)
+    mat = diag(val, mat_shape)
+
+    # val, shape attributes
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape
+
+    val = torch.randn(val_shape).to(ctx)
+
+    # nnz
+    assert mat.nnz == val.shape[0]
+    # dtype
+    assert mat.dtype == val.dtype
+    # device
+    assert mat.device == val.device
+
+    # row, col, val
+    edge_index = torch.arange(len(val)).to(mat.device)
+    row, col = mat.coo()
+    val = mat.val
+    assert torch.allclose(row, edge_index)
+    assert torch.allclose(col, edge_index)
+    assert torch.allclose(val, val)
+
+
+@pytest.mark.parametrize("shape", [(3, 3), (3, 5), (5, 3)])
+@pytest.mark.parametrize("d", [None, 2])
+def test_identity(shape, d):
+    ctx = F.ctx()
+    # creation
+    mat = identity(shape, d)
+    # shape
+    assert mat.shape == shape
+    # val
+    len_val = min(shape)
+    if d is None:
+        val_shape = len_val
+    else:
+        val_shape = (len_val, d)
+    val = torch.ones(val_shape)
+    assert torch.allclose(val, mat.val)
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+@pytest.mark.parametrize("mat_shape", [None, (3, 5), (5, 3)])
+def test_diag_matrix_transpose(val_shape, mat_shape):
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    mat = diag(val, mat_shape).transpose()
+
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape[::-1]

tests/python/pytorch/sparse/test_unary_op.py

@@ -0,0 +1,40 @@
+import sys
+
+import backend as F
+import torch
+
+from dgl.sparse import diag, spmatrix
+
+
+def test_neg():
+    ctx = F.ctx()
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([1, 2, 3]).to(ctx)
+    val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
+    A = spmatrix(torch.stack([row, col]), val)
+    neg_A = -A
+    assert A.shape == neg_A.shape
+    assert A.nnz == neg_A.nnz
+    assert torch.allclose(-A.val, neg_A.val)
+    assert torch.allclose(torch.stack(A.coo()), torch.stack(neg_A.coo()))
+    assert A.val.device == neg_A.val.device
+
+
+def test_diag_neg():
+    ctx = F.ctx()
+    val = torch.arange(3).float().to(ctx)
+    D = diag(val)
+    neg_D = -D
+    assert D.shape == neg_D.shape
+    assert torch.allclose(-D.val, neg_D.val)
+    assert D.val.device == neg_D.val.device
+
+
+def test_diag_inv():
+    ctx = F.ctx()
+    val = torch.arange(1, 4).float().to(ctx)
+    D = diag(val)
+    inv_D = D.inv()
+    assert D.shape == inv_D.shape
+    assert torch.allclose(1.0 / D.val, inv_D.val)
+    assert D.val.device == inv_D.val.device

tests/python/pytorch/sparse/utils.py

@@ -0,0 +1,163 @@
+import numpy as np
+import torch
+
+from dgl.sparse import diag, from_csc, from_csr, SparseMatrix, spmatrix
+
+np.random.seed(42)
+torch.random.manual_seed(42)
+
+
+def clone_detach_and_grad(t):
+    t = t.clone().detach()
+    t.requires_grad_()
+    return t
+
+
+def rand_stride(t):
+    """Add stride to the last dimension of a tensor."""
+    stride = np.random.randint(2, 4)
+    ret = torch.stack([t] * stride, dim=-1)[..., 0]
+    ret = ret.detach()
+    if torch.is_floating_point(t):
+        ret.requires_grad_()
+    return ret
+
+
+def rand_coo(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indices = torch.stack([row, col])
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return spmatrix(indices, val, shape)
+
+
+def rand_csr(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64)
+    for r in row.tolist():
+        indptr[r + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[1] + 1, device=dev, dtype=torch.int64)
+    for c in col.tolist():
+        indptr[c + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csc(indptr, indices, val, shape=shape)
+
+
+def rand_diag(shape, nnz, dev, nz_dim=None):
+    nnz = min(shape)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    return diag(val, shape)
+
+
+def rand_coo_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indices = torch.stack([row, col])
+    indices = rand_stride(indices)
+    return spmatrix(indices, val, shape)
+
+
+def rand_csr_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64)
+    for r in row.tolist():
+        indptr[r + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[1] + 1, device=dev, dtype=torch.int64)
+    for c in col.tolist():
+        indptr[c + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csc(indptr, indices, val, shape=shape)
+
+
+def sparse_matrix_to_dense(A: SparseMatrix):
+    dense = A.to_dense()
+    return clone_detach_and_grad(dense)
+
+
+def sparse_matrix_to_torch_sparse(A: SparseMatrix, val=None):
+    row, col = A.coo()
+    edge_index = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+    shape = A.shape
+    if val is None:
+        val = A.val
+    val = val.clone().detach()
+    if len(A.val.shape) > 1:
+        shape += (A.val.shape[-1],)
+    ret = torch.sparse_coo_tensor(edge_index, val, shape).coalesce()
+    ret.requires_grad_()
+    return ret
+
+
+def dense_mask(dense, sparse):
+    ret = torch.zeros_like(dense)
+    row, col = sparse.coo()
+    for r, c in zip(row, col):
+        ret[r, c] = dense[r, c]
+    return ret

tests/python/pytorch/test_ffi-stream.py

@@ -0,0 +1,201 @@
+import unittest
+from statistics import mean
+
+import backend as F
+
+import dgl
+import dgl.ndarray as nd
+import dgl.ops as OPS
+import numpy as np
+import torch
+from dgl import rand_graph
+from dgl._ffi.streams import _dgl_get_stream, to_dgl_stream_handle
+from dgl.utils import to_dgl_context
+
+
+# borrowed from PyTorch, torch/testing/_internal/common_utils.py
+def _get_cycles_per_ms() -> float:
+    """Measure and return approximate number of cycles per millisecond for torch.cuda._sleep"""
+
+    def measure() -> float:
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        torch.cuda._sleep(1000000)
+        end.record()
+        end.synchronize()
+        cycles_per_ms = 1000000 / start.elapsed_time(end)
+        return cycles_per_ms
+
+    # Get 10 values and remove the 2 max and 2 min and return the avg.
+    # This is to avoid system disturbance that skew the results, e.g.
+    # the very first cuda call likely does a bunch of init, which takes
+    # much longer than subsequent calls.
+    num = 10
+    vals = []
+    for _ in range(num):
+        vals.append(measure())
+    vals = sorted(vals)
+    return mean(vals[2 : num - 2])
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="stream only runs on GPU."
+)
+def test_basics():
+    g = rand_graph(10, 20, device=F.cpu())
+    x = torch.ones(g.num_nodes(), 10)
+    result = OPS.copy_u_sum(g, x).to(F.ctx())
+
+    # launch on default stream used in DGL
+    xx = x.to(device=F.ctx())
+    gg = g.to(device=F.ctx())
+    OPS.copy_u_sum(gg, xx)
+    assert torch.equal(OPS.copy_u_sum(gg, xx), result)
+
+    # launch on new stream created via torch.cuda
+    s = torch.cuda.Stream(device=F.ctx())
+    with torch.cuda.stream(s):
+        xx = x.to(device=F.ctx(), non_blocking=True)
+        gg = g.to(device=F.ctx())
+        OPS.copy_u_sum(gg, xx)
+    s.synchronize()
+    assert torch.equal(OPS.copy_u_sum(gg, xx), result)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="stream only runs on GPU."
+)
+def test_set_get_stream():
+    current_stream = torch.cuda.current_stream()
+    # test setting another stream
+    s = torch.cuda.Stream(device=F.ctx())
+    torch.cuda.set_stream(s)
+    assert (
+        to_dgl_stream_handle(s).value
+        == _dgl_get_stream(to_dgl_context(F.ctx())).value
+    )
+    # revert to default stream
+    torch.cuda.set_stream(current_stream)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="stream only runs on GPU."
+)
+# borrowed from PyTorch, test/test_cuda.py: test_record_stream()
+def test_record_stream_ndarray():
+    cycles_per_ms = _get_cycles_per_ms()
+
+    t = nd.array(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), ctx=nd.cpu())
+    t.pin_memory_()
+    result = nd.empty([4], ctx=nd.gpu(0))
+    stream = torch.cuda.Stream()
+    ptr = [None]
+
+    # Performs the CPU->GPU copy in a background stream
+    def perform_copy():
+        with torch.cuda.stream(stream):
+            tmp = t.copyto(nd.gpu(0))
+            ptr[0] = F.from_dgl_nd(tmp).data_ptr()
+        torch.cuda.current_stream().wait_stream(stream)
+        tmp.record_stream(to_dgl_stream_handle(torch.cuda.current_stream()))
+        torch.cuda._sleep(int(50 * cycles_per_ms))  # delay the copy
+        result.copyfrom(tmp)
+
+    perform_copy()
+    with torch.cuda.stream(stream):
+        tmp2 = nd.empty([4], ctx=nd.gpu(0))
+        assert (
+            F.from_dgl_nd(tmp2).data_ptr() != ptr[0]
+        ), "allocation re-used too soon"
+
+    assert torch.equal(
+        F.from_dgl_nd(result).cpu(), torch.tensor([1.0, 2.0, 3.0, 4.0])
+    )
+
+    # Check that the block will be re-used after the main stream finishes
+    torch.cuda.current_stream().synchronize()
+    with torch.cuda.stream(stream):
+        tmp3 = nd.empty([4], ctx=nd.gpu(0))
+        assert (
+            F.from_dgl_nd(tmp3).data_ptr() == ptr[0]
+        ), "allocation not re-used"
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="stream only runs on GPU."
+)
+def test_record_stream_graph_positive():
+    cycles_per_ms = _get_cycles_per_ms()
+
+    g = rand_graph(10, 20, device=F.cpu())
+    g.create_formats_()
+    x = torch.ones(g.num_nodes(), 10).to(F.ctx())
+    g1 = g.to(F.ctx())
+    # this is necessary to initialize the cusparse handle
+    result = OPS.copy_u_sum(g1, x)
+    torch.cuda.current_stream().synchronize()
+
+    stream = torch.cuda.Stream()
+    results2 = torch.zeros_like(result)
+
+    # Performs the computing in a background stream
+    def perform_computing():
+        with torch.cuda.stream(stream):
+            g2 = g.to(F.ctx())
+        torch.cuda.current_stream().wait_stream(stream)
+        g2.record_stream(torch.cuda.current_stream())
+        torch.cuda._sleep(int(50 * cycles_per_ms))  # delay the computing
+        results2.copy_(OPS.copy_u_sum(g2, x))
+
+    perform_computing()
+    with torch.cuda.stream(stream):
+        # since we have called record stream for g2, g3 won't reuse its memory
+        g3 = rand_graph(10, 20, device=F.ctx())
+        g3.create_formats_()
+    torch.cuda.current_stream().synchronize()
+    assert torch.equal(result, results2)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="stream only runs on GPU."
+)
+def test_record_stream_graph_negative():
+    cycles_per_ms = _get_cycles_per_ms()
+
+    g = rand_graph(10, 20, device=F.cpu())
+    g.create_formats_()
+    x = torch.ones(g.num_nodes(), 10).to(F.ctx())
+    g1 = g.to(F.ctx())
+    # this is necessary to initialize the cusparse handle
+    result = OPS.copy_u_sum(g1, x)
+    torch.cuda.current_stream().synchronize()
+
+    stream = torch.cuda.Stream()
+    results2 = torch.zeros_like(result)
+
+    # Performs the computing in a background stream
+    def perform_computing():
+        with torch.cuda.stream(stream):
+            g2 = g.to(F.ctx())
+        torch.cuda.current_stream().wait_stream(stream)
+        # omit record_stream will produce a wrong result
+        # g2.record_stream(torch.cuda.current_stream())
+        torch.cuda._sleep(int(50 * cycles_per_ms))  # delay the computing
+        results2.copy_(OPS.copy_u_sum(g2, x))
+
+    perform_computing()
+    with torch.cuda.stream(stream):
+        # g3 will reuse g2's memory block, resulting a wrong result
+        g3 = rand_graph(10, 20, device=F.ctx())
+        g3.create_formats_()
+    torch.cuda.current_stream().synchronize()
+    assert not torch.equal(result, results2)
+
+
+if __name__ == "__main__":
+    test_basics()
+    test_set_get_stream()
+    test_record_stream_ndarray()
+    test_record_stream_graph_positive()
+    test_record_stream_graph_negative()

tests/python/pytorch/test_heterograph-pickle.py

@@ -0,0 +1,31 @@
+import io
+import pickle
+
+import dgl
+
+import networkx as nx
+import torch
+
+
+def _reconstruct_pickle(obj):
+    f = io.BytesIO()
+    pickle.dump(obj, f)
+    f.seek(0)
+    obj = pickle.load(f)
+    f.close()
+    return obj
+
+
+def test_pickling_batched_graph():
+    # NOTE: this is a test for a wierd bug mentioned in
+    #   https://github.com/dmlc/dgl/issues/438
+    glist = [nx.path_graph(i + 5) for i in range(5)]
+    glist = [dgl.from_networkx(g) for g in glist]
+    bg = dgl.batch(glist)
+    bg.ndata["x"] = torch.randn((35, 5))
+    bg.edata["y"] = torch.randn((60, 3))
+    new_bg = _reconstruct_pickle(bg)
+
+
+if __name__ == "__main__":
+    test_pickling_batched_graph()

tests/python/pytorch/test_multiprocessing-ipc.py

@@ -0,0 +1,26 @@
+import os
+import unittest
+
+import dgl
+
+import torch as th
+import torch.multiprocessing as mp
+
+
+def sub_ipc(g):
+    print(g)
+    return g
+
+
+@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
+def test_torch_ipc():
+    g = dgl.graph(([0, 1, 2], [1, 2, 3]))
+    ctx = mp.get_context("spawn")
+    p = ctx.Process(target=sub_ipc, args=(g,))
+
+    p.start()
+    p.join()
+
+
+if __name__ == "__main__":
+    test_torch_ipc()

tests/python/pytorch/utils/test_pin_memory.py

@@ -0,0 +1,96 @@
+import backend as F
+
+import dgl
+import pytest
+import torch
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test."
+)
+def test_pin_noncontiguous():
+    t = torch.empty([10, 100]).transpose(0, 1)
+
+    assert not t.is_contiguous()
+    assert not F.is_pinned(t)
+
+    with pytest.raises(dgl.DGLError):
+        dgl.utils.pin_memory_inplace(t)
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test."
+)
+def test_pin_view():
+    t = torch.empty([100, 10])
+    v = t[10:20]
+
+    assert v.is_contiguous()
+    assert not F.is_pinned(t)
+
+    with pytest.raises(dgl.DGLError):
+        dgl.utils.pin_memory_inplace(v)
+
+    # make sure an empty view does not generate an error
+    u = t[10:10]
+    u = dgl.utils.pin_memory_inplace(u)
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test."
+)
+def test_unpin_automatically():
+    # run a sufficient number of iterations such that the memory pool should be
+    # re-used
+    for j in range(10):
+        t = torch.ones(10000, 10)
+        assert not F.is_pinned(t)
+        nd = dgl.utils.pin_memory_inplace(t)
+        assert F.is_pinned(t)
+        del nd
+        # dgl.ndarray will unpin its data upon destruction
+        assert not F.is_pinned(t)
+        del t
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test."
+)
+def test_pin_unpin_column():
+    g = dgl.graph(([1, 2, 3, 4], [0, 0, 0, 0]))
+
+    g.ndata["x"] = torch.randn(g.num_nodes())
+    g.pin_memory_()
+    assert g.is_pinned()
+    assert g.ndata["x"].is_pinned()
+    for col in g._node_frames[0].values():
+        assert col.pinned_by_dgl
+        assert col._data_nd is not None
+
+    g.ndata["x"] = torch.randn(g.num_nodes())  # unpin the old ndata['x']
+    assert g.is_pinned()
+    for col in g._node_frames[0].values():
+        assert not col.pinned_by_dgl
+        assert col._data_nd is None
+    assert not g.ndata["x"].is_pinned()
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test."
+)
+def test_pin_empty():
+    t = torch.tensor([])
+    assert not t.is_pinned()
+
+    # Empty tensors will not be pinned or unpinned. It's a no-op.
+    # This is also the default behavior in PyTorch.
+    # We just check that it won't raise an error.
+    nd = dgl.utils.pin_memory_inplace(t)
+    assert not t.is_pinned()
+
+
+if __name__ == "__main__":
+    test_pin_noncontiguous()
+    test_pin_view()
+    test_unpin_automatically()
+    test_pin_unpin_column()