chore: import upstream snapshot with attribution

tests/python/common/backend/test_set_default_backend.py

@@ -0,0 +1,12 @@
+import os
+import unittest
+
+import backend as F
+
+
+def test_set_default_backend():
+    default_dir = os.path.join(os.path.expanduser("~"), ".dgl_unit_test")
+    F.set_default_backend(default_dir, "pytorch")
+
+    # make sure the config file was created
+    assert os.path.exists(os.path.join(default_dir, "config.json"))

tests/python/common/backend/test_tensor.py

@@ -0,0 +1,54 @@
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.ndarray as nd
+import numpy as np
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="TF doesn't support inplace update",
+)
+def test_dlpack():
+    # test dlpack conversion.
+    def nd2th():
+        ans = np.array(
+            [[1.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]
+        )
+        x = nd.array(np.zeros((3, 4), dtype=np.float32))
+        dl = x.to_dlpack()
+        y = F.zerocopy_from_dlpack(dl)
+        y[0] = 1
+        print(x)
+        print(y)
+        assert np.allclose(x.asnumpy(), ans)
+
+    def th2nd():
+        ans = np.array(
+            [[1.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]
+        )
+        x = F.zeros((3, 4))
+        dl = F.zerocopy_to_dlpack(x)
+        y = nd.from_dlpack(dl)
+        x[0] = 1
+        print(x)
+        print(y)
+        assert np.allclose(y.asnumpy(), ans)
+
+    def th2nd_incontiguous():
+        x = F.astype(F.tensor([[0, 1], [2, 3]]), F.int64)
+        ans = np.array([0, 2])
+        y = x[:2, 0]
+        # Uncomment this line and comment the one below to observe error
+        # dl = dlpack.to_dlpack(y)
+        dl = F.zerocopy_to_dlpack(y)
+        z = nd.from_dlpack(dl)
+        print(x)
+        print(z)
+        assert np.allclose(z.asnumpy(), ans)
+
+    nd2th()
+    th2nd()
+    th2nd_incontiguous()

tests/python/common/cuda/test_gpu_cache.py

@@ -0,0 +1,77 @@
+#
+#   Copyright (c) 2022 by Contributors
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+
+import unittest
+
+import backend as F
+
+import dgl
+from utils import parametrize_idtype
+
+D = 5
+
+
+def generate_graph(idtype, grad=False, add_data=True):
+    g = dgl.graph([]).to(F.ctx(), dtype=idtype)
+    g.add_nodes(10)
+    u, v = [], []
+    # create a graph where 0 is the source and 9 is the sink
+    for i in range(1, 9):
+        u.append(0)
+        v.append(i)
+        u.append(i)
+        v.append(9)
+    # add a back flow from 9 to 0
+    u.append(9)
+    v.append(0)
+    g.add_edges(u, v)
+    if add_data:
+        ncol = F.randn((10, D))
+        ecol = F.randn((17, D))
+        if grad:
+            ncol = F.attach_grad(ncol)
+            ecol = F.attach_grad(ecol)
+        g.ndata["h"] = ncol
+        g.edata["l"] = ecol
+    return g
+
+
+@unittest.skipIf(not F.gpu_ctx(), reason="only necessary with GPU")
+@parametrize_idtype
+def test_gpu_cache(idtype):
+    g = generate_graph(idtype)
+    cache = dgl.cuda.GPUCache(5, D, idtype)
+    h = g.ndata["h"]
+
+    t = 5
+    keys = F.arange(0, t, dtype=idtype)
+    values, m_idx, m_keys = cache.query(keys)
+    m_values = h[F.tensor(m_keys, F.int64)]
+    values[F.tensor(m_idx, F.int64)] = m_values
+    cache.replace(m_keys, m_values)
+
+    keys = F.arange(3, 8, dtype=idtype)
+    values, m_idx, m_keys = cache.query(keys)
+    assert m_keys.shape[0] == 3 and m_idx.shape[0] == 3
+    m_values = h[F.tensor(m_keys, F.int64)]
+    values[F.tensor(m_idx, F.int64)] = m_values
+    assert (values != h[F.tensor(keys, F.int64)]).sum().item() == 0
+    cache.replace(m_keys, m_values)
+
+
+if __name__ == "__main__":
+    test_gpu_cache(F.int64)
+    test_gpu_cache(F.int32)

tests/python/common/data/data/test_heterophilous_graphs.py

@@ -0,0 +1,97 @@
+import unittest
+
+import backend as F
+
+import dgl
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only supports PyTorch backend.",
+)
+def test_roman_empire():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.RomanEmpireDataset(force_reload=True)[0]
+    assert g.num_nodes() == 22662
+    assert g.num_edges() == 65854
+    g2 = dgl.data.RomanEmpireDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only supports PyTorch backend.",
+)
+def test_amazon_ratings():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.AmazonRatingsDataset(force_reload=True)[0]
+    assert g.num_nodes() == 24492
+    assert g.num_edges() == 186100
+    g2 = dgl.data.AmazonRatingsDataset(force_reload=True, transform=transform)[
+        0
+    ]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only supports PyTorch backend.",
+)
+def test_minesweeper():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.MinesweeperDataset(force_reload=True)[0]
+    assert g.num_nodes() == 10000
+    assert g.num_edges() == 78804
+    g2 = dgl.data.MinesweeperDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only supports PyTorch backend.",
+)
+def test_tolokers():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.TolokersDataset(force_reload=True)[0]
+    assert g.num_nodes() == 11758
+    assert g.num_edges() == 1038000
+    g2 = dgl.data.TolokersDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only supports PyTorch backend.",
+)
+def test_questions():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.QuestionsDataset(force_reload=True)[0]
+    assert g.num_nodes() == 48921
+    assert g.num_edges() == 307080
+    g2 = dgl.data.QuestionsDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()

tests/python/common/data/test_actor.py

@@ -0,0 +1,22 @@
+import unittest
+
+import backend as F
+
+import dgl
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_actor():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.ActorDataset(force_reload=True)[0]
+    assert g.num_nodes() == 7600
+    assert g.num_edges() == 33391
+    g2 = dgl.data.ActorDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()

tests/python/common/data/test_geom_gcn.py

@@ -0,0 +1,90 @@
+import unittest
+
+import backend as F
+
+import dgl
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_chameleon():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.ChameleonDataset(force_reload=True)[0]
+    assert g.num_nodes() == 2277
+    assert g.num_edges() == 36101
+    g2 = dgl.data.ChameleonDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_squirrel():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.SquirrelDataset(force_reload=True)[0]
+    assert g.num_nodes() == 5201
+    assert g.num_edges() == 217073
+    g2 = dgl.data.SquirrelDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_cornell():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.CornellDataset(force_reload=True)[0]
+    assert g.num_nodes() == 183
+    assert g.num_edges() == 298
+    g2 = dgl.data.CornellDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_texas():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.TexasDataset(force_reload=True)[0]
+    assert g.num_nodes() == 183
+    assert g.num_edges() == 325
+    g2 = dgl.data.TexasDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_wisconsin():
+    transform = dgl.AddSelfLoop(allow_duplicate=True)
+
+    g = dgl.data.WisconsinDataset(force_reload=True)[0]
+    assert g.num_nodes() == 251
+    assert g.num_edges() == 515
+    g2 = dgl.data.WisconsinDataset(force_reload=True, transform=transform)[0]
+    assert g2.num_edges() - g.num_edges() == g.num_nodes()

tests/python/common/data/test_movielens.py

@@ -0,0 +1,53 @@
+import unittest
+
+import backend as F
+
+import dgl
+from dgl.data.movielens import MovieLensDataset
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_movielens():
+    transform = dgl.AddSelfLoop(new_etypes=True)
+
+    movielens = MovieLensDataset(name="ml-100k", valid_ratio=0.2, verbose=True)
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 100000
+    assert (
+        g.nodes["user"].data["feat"].shape[1]
+        == g.nodes["user"].data["feat"].shape[1]
+        == g.nodes["user"].data["feat"].shape[1]
+        == 23
+    )
+    assert (
+        g.nodes["movie"].data["feat"].shape[1]
+        == g.nodes["movie"].data["feat"].shape[1]
+        == g.nodes["movie"].data["feat"].shape[1]
+        == 320
+    )
+
+    movielens = MovieLensDataset(
+        name="ml-100k", valid_ratio=0.2, transform=transform, verbose=True
+    )
+    g1 = movielens[0]
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+
+    movielens = MovieLensDataset(
+        name="ml-1m", valid_ratio=0.2, test_ratio=0.1, verbose=True
+    )
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 1000209
+
+    movielens = MovieLensDataset(
+        name="ml-10m", valid_ratio=0.2, test_ratio=0.1, verbose=True
+    )
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 10000054

tests/python/common/data/test_serialize.py

@@ -0,0 +1,442 @@
+import os
+import tempfile
+import time
+import unittest
+import warnings
+
+import backend as F
+
+import dgl
+import dgl.ndarray as nd
+import numpy as np
+import pytest
+import scipy as sp
+from dgl.data.utils import load_labels, load_tensors, save_tensors
+
+np.random.seed(44)
+
+
+def generate_rand_graph(n):
+    arr = (sp.sparse.random(n, n, density=0.1, format="coo") != 0).astype(
+        np.int64
+    )
+    return dgl.from_scipy(arr)
+
+
+def construct_graph(n):
+    g_list = []
+    for _ in range(n):
+        g = generate_rand_graph(30)
+        g.edata["e1"] = F.randn((g.num_edges(), 32))
+        g.edata["e2"] = F.ones((g.num_edges(), 32))
+        g.ndata["n1"] = F.randn((g.num_nodes(), 64))
+        g_list.append(g)
+    return g_list
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_graph_serialize_with_feature():
+    num_graphs = 100
+
+    t0 = time.time()
+
+    g_list = construct_graph(num_graphs)
+
+    t1 = time.time()
+
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    dgl.save_graphs(path, g_list)
+
+    t2 = time.time()
+    idx_list = np.random.permutation(np.arange(num_graphs)).tolist()
+    loadg_list, _ = dgl.load_graphs(path, idx_list)
+
+    t3 = time.time()
+    idx = idx_list[0]
+    load_g = loadg_list[0]
+    print("Save time: {} s".format(t2 - t1))
+    print("Load time: {} s".format(t3 - t2))
+    print("Graph Construction time: {} s".format(t1 - t0))
+
+    assert F.allclose(load_g.nodes(), g_list[idx].nodes())
+
+    load_edges = load_g.all_edges("uv", "eid")
+    g_edges = g_list[idx].all_edges("uv", "eid")
+    assert F.allclose(load_edges[0], g_edges[0])
+    assert F.allclose(load_edges[1], g_edges[1])
+    assert F.allclose(load_g.edata["e1"], g_list[idx].edata["e1"])
+    assert F.allclose(load_g.edata["e2"], g_list[idx].edata["e2"])
+    assert F.allclose(load_g.ndata["n1"], g_list[idx].ndata["n1"])
+
+    os.unlink(path)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_graph_serialize_without_feature():
+    num_graphs = 100
+    g_list = [generate_rand_graph(30) for _ in range(num_graphs)]
+
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    dgl.save_graphs(path, g_list)
+
+    idx_list = np.random.permutation(np.arange(num_graphs)).tolist()
+    loadg_list, _ = dgl.load_graphs(path, idx_list)
+
+    idx = idx_list[0]
+    load_g = loadg_list[0]
+
+    assert F.allclose(load_g.nodes(), g_list[idx].nodes())
+
+    load_edges = load_g.all_edges("uv", "eid")
+    g_edges = g_list[idx].all_edges("uv", "eid")
+    assert F.allclose(load_edges[0], g_edges[0])
+    assert F.allclose(load_edges[1], g_edges[1])
+
+    os.unlink(path)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_graph_serialize_with_labels():
+    num_graphs = 100
+    g_list = [generate_rand_graph(30) for _ in range(num_graphs)]
+    labels = {"label": F.zeros((num_graphs, 1))}
+
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    dgl.save_graphs(path, g_list, labels)
+
+    idx_list = np.random.permutation(np.arange(num_graphs)).tolist()
+    loadg_list, l_labels0 = dgl.load_graphs(path, idx_list)
+    l_labels = load_labels(path)
+    assert F.allclose(l_labels["label"], labels["label"])
+    assert F.allclose(l_labels0["label"], labels["label"])
+
+    idx = idx_list[0]
+    load_g = loadg_list[0]
+
+    assert F.allclose(load_g.nodes(), g_list[idx].nodes())
+
+    load_edges = load_g.all_edges("uv", "eid")
+    g_edges = g_list[idx].all_edges("uv", "eid")
+    assert F.allclose(load_edges[0], g_edges[0])
+    assert F.allclose(load_edges[1], g_edges[1])
+
+    os.unlink(path)
+
+
+def test_serialize_tensors():
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    tensor_dict = {
+        "a": F.tensor([1, 3, -1, 0], dtype=F.int64),
+        "1@1": F.tensor([1.5, 2], dtype=F.float32),
+    }
+
+    save_tensors(path, tensor_dict)
+
+    load_tensor_dict = load_tensors(path)
+
+    for key in tensor_dict:
+        assert key in load_tensor_dict
+        assert np.array_equal(
+            F.asnumpy(load_tensor_dict[key]), F.asnumpy(tensor_dict[key])
+        )
+
+    load_nd_dict = load_tensors(path, return_dgl_ndarray=True)
+
+    for key in tensor_dict:
+        assert key in load_nd_dict
+        assert isinstance(load_nd_dict[key], nd.NDArray)
+        assert np.array_equal(
+            load_nd_dict[key].asnumpy(), F.asnumpy(tensor_dict[key])
+        )
+
+    os.unlink(path)
+
+
+def test_serialize_empty_dict():
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    tensor_dict = {}
+
+    save_tensors(path, tensor_dict)
+
+    load_tensor_dict = load_tensors(path)
+    assert isinstance(load_tensor_dict, dict)
+    assert len(load_tensor_dict) == 0
+
+    os.unlink(path)
+
+
+def load_old_files(files):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return dgl.load_graphs(os.path.join(os.path.dirname(__file__), files))
+
+
+def test_load_old_files1():
+    loadg_list, _ = load_old_files("data/1.bin")
+    idx, num_nodes, edge0, edge1, edata_e1, edata_e2, ndata_n1 = np.load(
+        os.path.join(os.path.dirname(__file__), "data/1.npy"), allow_pickle=True
+    )
+
+    load_g = loadg_list[idx]
+    load_edges = load_g.all_edges("uv", "eid")
+
+    assert np.allclose(F.asnumpy(load_edges[0]), edge0)
+    assert np.allclose(F.asnumpy(load_edges[1]), edge1)
+    assert np.allclose(F.asnumpy(load_g.edata["e1"]), edata_e1)
+    assert np.allclose(F.asnumpy(load_g.edata["e2"]), edata_e2)
+    assert np.allclose(F.asnumpy(load_g.ndata["n1"]), ndata_n1)
+
+
+def test_load_old_files2():
+    loadg_list, labels0 = load_old_files("data/2.bin")
+    labels1 = load_labels(os.path.join(os.path.dirname(__file__), "data/2.bin"))
+    idx, edges0, edges1, np_labels = np.load(
+        os.path.join(os.path.dirname(__file__), "data/2.npy"), allow_pickle=True
+    )
+    assert np.allclose(F.asnumpy(labels0["label"]), np_labels)
+    assert np.allclose(F.asnumpy(labels1["label"]), np_labels)
+
+    load_g = loadg_list[idx]
+    print(load_g)
+    load_edges = load_g.all_edges("uv", "eid")
+    assert np.allclose(F.asnumpy(load_edges[0]), edges0)
+    assert np.allclose(F.asnumpy(load_edges[1]), edges1)
+
+
+def create_heterographs(idtype):
+    g_x = dgl.heterograph(
+        {("user", "follows", "user"): ([0, 1, 2], [1, 2, 3])}, idtype=idtype
+    )
+    g_y = dgl.heterograph(
+        {("user", "knows", "user"): ([0, 2], [2, 3])}, idtype=idtype
+    ).formats("csr")
+    g_x.ndata["h"] = F.randn((4, 3))
+    g_x.edata["w"] = F.randn((3, 2))
+    g_y.ndata["hh"] = F.ones((4, 5))
+    g_y.edata["ww"] = F.randn((2, 10))
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1, 2], [1, 2, 3]),
+            ("user", "knows", "user"): ([0, 2], [2, 3]),
+        },
+        idtype=idtype,
+    )
+    g.nodes["user"].data["h"] = g_x.ndata["h"]
+    g.nodes["user"].data["hh"] = g_y.ndata["hh"]
+    g.edges["follows"].data["w"] = g_x.edata["w"]
+    g.edges["knows"].data["ww"] = g_y.edata["ww"]
+    return [g, g_x, g_y]
+
+
+def create_heterographs2(idtype):
+    g_x = dgl.heterograph(
+        {("user", "follows", "user"): ([0, 1, 2], [1, 2, 3])}, idtype=idtype
+    )
+    g_y = dgl.heterograph(
+        {("user", "knows", "user"): ([0, 2], [2, 3])}, idtype=idtype
+    ).formats("csr")
+    g_z = dgl.heterograph(
+        {("user", "knows", "knowledge"): ([0, 1, 3], [2, 3, 4])}, idtype=idtype
+    )
+    g_x.ndata["h"] = F.randn((4, 3))
+    g_x.edata["w"] = F.randn((3, 2))
+    g_y.ndata["hh"] = F.ones((4, 5))
+    g_y.edata["ww"] = F.randn((2, 10))
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1, 2], [1, 2, 3]),
+            ("user", "knows", "user"): ([0, 2], [2, 3]),
+            ("user", "knows", "knowledge"): ([0, 1, 3], [2, 3, 4]),
+        },
+        idtype=idtype,
+    )
+    g.nodes["user"].data["h"] = g_x.ndata["h"]
+    g.edges["follows"].data["w"] = g_x.edata["w"]
+    g.nodes["user"].data["hh"] = g_y.ndata["hh"]
+    g.edges[("user", "knows", "user")].data["ww"] = g_y.edata["ww"]
+    return [g, g_x, g_y, g_z]
+
+
+def test_deserialize_old_heterograph_file():
+    path = os.path.join(os.path.dirname(__file__), "data/hetero1.bin")
+    g_list, label_dict = dgl.load_graphs(path)
+    assert g_list[0].idtype == F.int64
+    assert g_list[3].idtype == F.int32
+    assert np.allclose(
+        F.asnumpy(g_list[2].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    assert np.allclose(
+        F.asnumpy(g_list[5].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    edges = g_list[0]["follows"].edges()
+    assert np.allclose(F.asnumpy(edges[0]), np.array([0, 1, 2]))
+    assert np.allclose(F.asnumpy(edges[1]), np.array([1, 2, 3]))
+    assert F.allclose(label_dict["graph_label"], F.ones(54))
+
+
+def create_old_heterograph_files():
+    path = os.path.join(os.path.dirname(__file__), "data/hetero1.bin")
+    g_list0 = create_heterographs(F.int64) + create_heterographs(F.int32)
+    labels_dict = {"graph_label": F.ones(54)}
+    dgl.save_graphs(path, g_list0, labels_dict)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_serialize_heterograph():
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+    g_list0 = create_heterographs2(F.int64) + create_heterographs2(F.int32)
+    dgl.save_graphs(path, g_list0)
+
+    g_list, _ = dgl.load_graphs(path)
+    assert g_list[0].idtype == F.int64
+    assert len(g_list[0].canonical_etypes) == 3
+    for i in range(len(g_list0)):
+        for j, etypes in enumerate(g_list0[i].canonical_etypes):
+            assert g_list[i].canonical_etypes[j] == etypes
+    # assert g_list[1].restrict_format() == 'any'
+    # assert g_list[2].restrict_format() == 'csr'
+
+    assert g_list[4].idtype == F.int32
+    assert np.allclose(
+        F.asnumpy(g_list[2].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    assert np.allclose(
+        F.asnumpy(g_list[6].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    edges = g_list[0]["follows"].edges()
+    assert np.allclose(F.asnumpy(edges[0]), np.array([0, 1, 2]))
+    assert np.allclose(F.asnumpy(edges[1]), np.array([1, 2, 3]))
+    for i in range(len(g_list)):
+        assert g_list[i].ntypes == g_list0[i].ntypes
+        assert g_list[i].etypes == g_list0[i].etypes
+
+    # test set feature after load_graph
+    g_list[3].nodes["user"].data["test"] = F.tensor([0, 1, 2, 4])
+    g_list[3].edata["test"] = F.tensor([0, 1, 2])
+
+    os.unlink(path)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@pytest.mark.skip(reason="lack of permission on CI")
+def test_serialize_heterograph_s3():
+    path = "s3://dglci-data-test/graph2.bin"
+    g_list0 = create_heterographs(F.int64) + create_heterographs(F.int32)
+    dgl.save_graphs(path, g_list0)
+
+    g_list = dgl.load_graphs(path, [0, 2, 5])
+    assert g_list[0].idtype == F.int64
+    # assert g_list[1].restrict_format() == 'csr'
+    assert np.allclose(
+        F.asnumpy(g_list[1].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    assert np.allclose(
+        F.asnumpy(g_list[2].nodes["user"].data["hh"]), np.ones((4, 5))
+    )
+    edges = g_list[0]["follows"].edges()
+    assert np.allclose(F.asnumpy(edges[0]), np.array([0, 1, 2]))
+    assert np.allclose(F.asnumpy(edges[1]), np.array([1, 2, 3]))
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@pytest.mark.parametrize(
+    "formats",
+    [
+        "coo",
+        "csr",
+        "csc",
+        ["coo", "csc"],
+        ["coo", "csr"],
+        ["csc", "csr"],
+        ["coo", "csr", "csc"],
+    ],
+)
+def test_graph_serialize_with_formats(formats):
+    num_graphs = 100
+    g_list = [generate_rand_graph(30) for _ in range(num_graphs)]
+
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    dgl.save_graphs(path, g_list, formats=formats)
+
+    idx_list = np.random.permutation(np.arange(num_graphs)).tolist()
+    loadg_list, _ = dgl.load_graphs(path, idx_list)
+
+    idx = idx_list[0]
+    load_g = loadg_list[0]
+    g_formats = load_g.formats()
+
+    # verify formats
+    if not isinstance(formats, list):
+        formats = [formats]
+    for fmt in formats:
+        assert fmt in g_formats["created"]
+
+    assert F.allclose(load_g.nodes(), g_list[idx].nodes())
+
+    load_edges = load_g.all_edges("uv", "eid")
+    g_edges = g_list[idx].all_edges("uv", "eid")
+    assert F.allclose(load_edges[0], g_edges[0])
+    assert F.allclose(load_edges[1], g_edges[1])
+
+    os.unlink(path)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_graph_serialize_with_restricted_formats():
+    g = dgl.rand_graph(100, 200)
+    g = g.formats(["coo"])
+    g_list = [g]
+
+    # create a temporary file and immediately release it so DGL can open it.
+    f = tempfile.NamedTemporaryFile(delete=False)
+    path = f.name
+    f.close()
+
+    expect_except = False
+    try:
+        dgl.save_graphs(path, g_list, formats=["csr"])
+    except:
+        expect_except = True
+    assert expect_except
+
+    os.unlink(path)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_deserialize_old_graph():
+    num_nodes = 100
+    num_edges = 200
+    path = os.path.join(os.path.dirname(__file__), "data/graph_0.9a220622.dgl")
+    g_list, _ = dgl.load_graphs(path)
+    g = g_list[0]
+    assert "coo" in g.formats()["created"]
+    assert "csr" in g.formats()["not created"]
+    assert "csc" in g.formats()["not created"]
+    assert num_nodes == g.num_nodes()
+    assert num_edges == g.num_edges()

tests/python/common/data/test_utils.py

@@ -0,0 +1,102 @@
+import gzip
+import io
+import os
+import tarfile
+import tempfile
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.data as data
+import numpy as np
+import pandas as pd
+import pytest
+import yaml
+from dgl import DGLError
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Skip MXNet")
+def test_add_nodepred_split():
+    dataset = data.AmazonCoBuyComputerDataset()
+    print("train_mask" in dataset[0].ndata)
+    data.utils.add_nodepred_split(dataset, [0.8, 0.1, 0.1])
+    assert "train_mask" in dataset[0].ndata
+
+    dataset = data.AIFBDataset()
+    print("train_mask" in dataset[0].nodes["Publikationen"].data)
+    data.utils.add_nodepred_split(
+        dataset, [0.8, 0.1, 0.1], ntype="Publikationen"
+    )
+    assert "train_mask" in dataset[0].nodes["Publikationen"].data
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Skip MXNet")
+def test_extract_archive():
+    # gzip
+    with tempfile.TemporaryDirectory() as src_dir:
+        gz_file = "gz_archive"
+        gz_path = os.path.join(src_dir, gz_file + ".gz")
+        content = b"test extract archive gzip"
+        with gzip.open(gz_path, "wb") as f:
+            f.write(content)
+        with tempfile.TemporaryDirectory() as dst_dir:
+            data.utils.extract_archive(gz_path, dst_dir, overwrite=True)
+            assert os.path.exists(os.path.join(dst_dir, gz_file))
+
+    # tar
+    with tempfile.TemporaryDirectory() as src_dir:
+        tar_file = "tar_archive"
+        tar_path = os.path.join(src_dir, tar_file + ".tar")
+        # default encode to utf8
+        content = "test extract archive tar\n".encode()
+        info = tarfile.TarInfo(name="tar_archive")
+        info.size = len(content)
+        with tarfile.open(tar_path, "w") as f:
+            f.addfile(info, io.BytesIO(content))
+        with tempfile.TemporaryDirectory() as dst_dir:
+            data.utils.extract_archive(tar_path, dst_dir, overwrite=True)
+            assert os.path.exists(os.path.join(dst_dir, tar_file))
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Skip MXNet")
+def test_mask_nodes_by_property():
+    num_nodes = 1000
+    property_values = np.random.uniform(size=num_nodes)
+    part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    split_masks = data.utils.mask_nodes_by_property(
+        property_values, part_ratios
+    )
+    assert "in_valid_mask" in split_masks
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="Skip MXNet")
+def test_add_node_property_split():
+    dataset = data.AmazonCoBuyComputerDataset()
+    part_ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+    for property_name in ["popularity", "locality", "density"]:
+        data.utils.add_node_property_split(dataset, part_ratios, property_name)
+        assert "in_valid_mask" in dataset[0].ndata
+
+
+if __name__ == "__main__":
+    test_extract_archive()
+    test_add_nodepred_split()
+    test_mask_nodes_by_property()
+    test_add_node_property_split()

tests/python/common/dataloading/test_dataloader.py

@@ -0,0 +1,52 @@
+import unittest
+
+import backend as F
+
+import dgl
+from dgl.dataloading import (
+    as_edge_prediction_sampler,
+    negative_sampler,
+    NeighborSampler,
+)
+from utils import parametrize_idtype
+
+
+def create_test_graph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+@parametrize_idtype
+def test_edge_prediction_sampler(idtype):
+    g = create_test_graph(idtype)
+    sampler = NeighborSampler([10, 10])
+    sampler = as_edge_prediction_sampler(
+        sampler, negative_sampler=negative_sampler.Uniform(1)
+    )
+
+    seeds = F.copy_to(F.arange(0, 2, dtype=idtype), ctx=F.ctx())
+    # just a smoke test to make sure we don't fail internal assertions
+    result = sampler.sample(g, {"follows": seeds})
+
+
+if __name__ == "__main__":
+    test_edge_prediction_sampler()

tests/python/common/function/test_basics.py

@@ -0,0 +1,798 @@
+import warnings
+from collections import defaultdict as ddict
+
+import backend as F
+
+import dgl
+import networkx as nx
+import numpy as np
+from utils import parametrize_idtype
+
+D = 5
+reduce_msg_shapes = set()
+
+
+def message_func(edges):
+    assert F.ndim(edges.src["h"]) == 2
+    assert F.shape(edges.src["h"])[1] == D
+    return {"m": edges.src["h"]}
+
+
+def reduce_func(nodes):
+    msgs = nodes.mailbox["m"]
+    reduce_msg_shapes.add(tuple(msgs.shape))
+    assert F.ndim(msgs) == 3
+    assert F.shape(msgs)[2] == D
+    return {"accum": F.sum(msgs, 1)}
+
+
+def apply_node_func(nodes):
+    return {"h": nodes.data["h"] + nodes.data["accum"]}
+
+
+def generate_graph_old(grad=False):
+    g = dgl.graph([])
+    g.add_nodes(10)  # 10 nodes
+    # create a graph where 0 is the source and 9 is the sink
+    # 17 edges
+    for i in range(1, 9):
+        g.add_edges(0, i)
+        g.add_edges(i, 9)
+    # add a back flow from 9 to 0
+    g.add_edges(9, 0)
+    g = g.to(F.ctx())
+    ncol = F.randn((10, D))
+    ecol = F.randn((17, D))
+    if grad:
+        ncol = F.attach_grad(ncol)
+        ecol = F.attach_grad(ecol)
+
+    g.ndata["h"] = ncol
+    g.edata["w"] = ecol
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
+    return g
+
+
+def generate_graph(idtype, grad=False):
+    """
+    s, d, eid
+    0, 1, 0
+    1, 9, 1
+    0, 2, 2
+    2, 9, 3
+    0, 3, 4
+    3, 9, 5
+    0, 4, 6
+    4, 9, 7
+    0, 5, 8
+    5, 9, 9
+    0, 6, 10
+    6, 9, 11
+    0, 7, 12
+    7, 9, 13
+    0, 8, 14
+    8, 9, 15
+    9, 0, 16
+    """
+    u = F.tensor([0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 9])
+    v = F.tensor([1, 9, 2, 9, 3, 9, 4, 9, 5, 9, 6, 9, 7, 9, 8, 9, 0])
+    g = dgl.graph((u, v), idtype=idtype)
+    assert g.device == F.ctx()
+    ncol = F.randn((10, D))
+    ecol = F.randn((17, D))
+    if grad:
+        ncol = F.attach_grad(ncol)
+        ecol = F.attach_grad(ecol)
+
+    g.ndata["h"] = ncol
+    g.edata["w"] = ecol
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
+    return g
+
+
+def test_compatible():
+    g = generate_graph_old()
+
+
+@parametrize_idtype
+def test_batch_setter_getter(idtype):
+    def _pfc(x):
+        return list(F.zerocopy_to_numpy(x)[:, 0])
+
+    g = generate_graph(idtype)
+    # set all nodes
+    g.ndata["h"] = F.zeros((10, D))
+    assert F.allclose(g.ndata["h"], F.zeros((10, D)))
+    # pop nodes
+    old_len = len(g.ndata)
+    g.ndata.pop("h")
+    assert len(g.ndata) == old_len - 1
+    g.ndata["h"] = F.zeros((10, D))
+    # set partial nodes
+    u = F.tensor([1, 3, 5], g.idtype)
+    g.nodes[u].data["h"] = F.ones((3, D))
+    assert _pfc(g.ndata["h"]) == [
+        0.0,
+        1.0,
+        0.0,
+        1.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+    ]
+    # get partial nodes
+    u = F.tensor([1, 2, 3], g.idtype)
+    assert _pfc(g.nodes[u].data["h"]) == [1.0, 0.0, 1.0]
+
+    """
+    s, d, eid
+    0, 1, 0
+    1, 9, 1
+    0, 2, 2
+    2, 9, 3
+    0, 3, 4
+    3, 9, 5
+    0, 4, 6
+    4, 9, 7
+    0, 5, 8
+    5, 9, 9
+    0, 6, 10
+    6, 9, 11
+    0, 7, 12
+    7, 9, 13
+    0, 8, 14
+    8, 9, 15
+    9, 0, 16
+    """
+    # set all edges
+    g.edata["l"] = F.zeros((17, D))
+    assert _pfc(g.edata["l"]) == [0.0] * 17
+    # pop edges
+    old_len = len(g.edata)
+    g.edata.pop("l")
+    assert len(g.edata) == old_len - 1
+    g.edata["l"] = F.zeros((17, D))
+    # set partial edges (many-many)
+    u = F.tensor([0, 0, 2, 5, 9], g.idtype)
+    v = F.tensor([1, 3, 9, 9, 0], g.idtype)
+    g.edges[u, v].data["l"] = F.ones((5, D))
+    truth = [0.0] * 17
+    truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.0
+    assert _pfc(g.edata["l"]) == truth
+    u = F.tensor([3, 4, 6], g.idtype)
+    v = F.tensor([9, 9, 9], g.idtype)
+    g.edges[u, v].data["l"] = F.ones((3, D))
+    truth[5] = truth[7] = truth[11] = 1.0
+    assert _pfc(g.edata["l"]) == truth
+    u = F.tensor([0, 0, 0], g.idtype)
+    v = F.tensor([4, 5, 6], g.idtype)
+    g.edges[u, v].data["l"] = F.ones((3, D))
+    truth[6] = truth[8] = truth[10] = 1.0
+    assert _pfc(g.edata["l"]) == truth
+    u = F.tensor([0, 6, 0], g.idtype)
+    v = F.tensor([6, 9, 7], g.idtype)
+    assert _pfc(g.edges[u, v].data["l"]) == [1.0, 1.0, 0.0]
+
+
+@parametrize_idtype
+def test_batch_setter_autograd(idtype):
+    g = generate_graph(idtype, grad=True)
+    h1 = g.ndata["h"]
+    # partial set
+    v = F.tensor([1, 2, 8], g.idtype)
+    hh = F.attach_grad(F.zeros((len(v), D)))
+    with F.record_grad():
+        g.nodes[v].data["h"] = hh
+        h2 = g.ndata["h"]
+        F.backward(h2, F.ones((10, D)) * 2)
+    assert F.array_equal(
+        F.grad(h1)[:, 0],
+        F.tensor([2.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 2.0]),
+    )
+    assert F.array_equal(F.grad(hh)[:, 0], F.tensor([2.0, 2.0, 2.0]))
+
+
+def _test_nx_conversion():
+    # check conversion between networkx and DGLGraph
+
+    def _check_nx_feature(nxg, nf, ef):
+        # check node and edge feature of nxg
+        # this is used to check to_networkx
+        num_nodes = len(nxg)
+        num_edges = nxg.size()
+        if num_nodes > 0:
+            node_feat = ddict(list)
+            for nid, attr in nxg.nodes(data=True):
+                assert len(attr) == len(nf)
+                for k in nxg.nodes[nid]:
+                    node_feat[k].append(F.unsqueeze(attr[k], 0))
+            for k in node_feat:
+                feat = F.cat(node_feat[k], 0)
+                assert F.allclose(feat, nf[k])
+        else:
+            assert len(nf) == 0
+        if num_edges > 0:
+            edge_feat = ddict(lambda: [0] * num_edges)
+            for u, v, attr in nxg.edges(data=True):
+                assert len(attr) == len(ef) + 1  # extra id
+                eid = attr["id"]
+                for k in ef:
+                    edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
+            for k in edge_feat:
+                feat = F.cat(edge_feat[k], 0)
+                assert F.allclose(feat, ef[k])
+        else:
+            assert len(ef) == 0
+
+    n1 = F.randn((5, 3))
+    n2 = F.randn((5, 10))
+    n3 = F.randn((5, 4))
+    e1 = F.randn((4, 5))
+    e2 = F.randn((4, 7))
+    g = dgl.graph(([0, 1, 3, 4], [2, 4, 0, 3]))
+    g.ndata.update({"n1": n1, "n2": n2, "n3": n3})
+    g.edata.update({"e1": e1, "e2": e2})
+
+    # convert to networkx
+    nxg = g.to_networkx(node_attrs=["n1", "n3"], edge_attrs=["e1", "e2"])
+    assert len(nxg) == 5
+    assert nxg.size() == 4
+    _check_nx_feature(nxg, {"n1": n1, "n3": n3}, {"e1": e1, "e2": e2})
+
+    # convert to DGLGraph, nx graph has id in edge feature
+    # use id feature to test non-tensor copy
+    g = dgl.from_networkx(nxg, node_attrs=["n1"], edge_attrs=["e1", "id"])
+    # check graph size
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 4
+    # check number of features
+    # test with existing dglgraph (so existing features should be cleared)
+    assert len(g.ndata) == 1
+    assert len(g.edata) == 2
+    # check feature values
+    assert F.allclose(g.ndata["n1"], n1)
+    # with id in nx edge feature, e1 should follow original order
+    assert F.allclose(g.edata["e1"], e1)
+    assert F.array_equal(
+        F.astype(g.edata["id"], F.int64), F.copy_to(F.arange(0, 4), F.cpu())
+    )
+
+    # test conversion after modifying DGLGraph
+    g.edata.pop("id")  # pop id so we don't need to provide id when adding edges
+    new_n = F.randn((2, 3))
+    new_e = F.randn((3, 5))
+    g.add_nodes(2, data={"n1": new_n})
+    # add three edges, one is a multi-edge
+    g.add_edges([3, 6, 0], [4, 5, 2], data={"e1": new_e})
+    n1 = F.cat((n1, new_n), 0)
+    e1 = F.cat((e1, new_e), 0)
+    # convert to networkx again
+    nxg = g.to_networkx(node_attrs=["n1"], edge_attrs=["e1"])
+    assert len(nxg) == 7
+    assert nxg.size() == 7
+    _check_nx_feature(nxg, {"n1": n1}, {"e1": e1})
+
+    # now test convert from networkx without id in edge feature
+    # first pop id in edge feature
+    for _, _, attr in nxg.edges(data=True):
+        attr.pop("id")
+    # test with a new graph
+    g = dgl.from_networkx(nxg, node_attrs=["n1"], edge_attrs=["e1"])
+    # check graph size
+    assert g.num_nodes() == 7
+    assert g.num_edges() == 7
+    # check number of features
+    assert len(g.ndata) == 1
+    assert len(g.edata) == 1
+    # check feature values
+    assert F.allclose(g.ndata["n1"], n1)
+    # edge feature order follows nxg.edges()
+    edge_feat = []
+    for _, _, attr in nxg.edges(data=True):
+        edge_feat.append(F.unsqueeze(attr["e1"], 0))
+    edge_feat = F.cat(edge_feat, 0)
+    assert F.allclose(g.edata["e1"], edge_feat)
+
+    # Test converting from a networkx graph whose nodes are
+    # not labeled with consecutive-integers.
+    nxg = nx.cycle_graph(5)
+    nxg.remove_nodes_from([0, 4])
+    for u in nxg.nodes():
+        nxg.nodes[u]["h"] = F.tensor([u])
+    for u, v, d in nxg.edges(data=True):
+        d["h"] = F.tensor([u, v])
+
+    g = dgl.from_networkx(nxg, node_attrs=["h"], edge_attrs=["h"])
+    assert g.num_nodes() == 3
+    assert g.num_edges() == 4
+    assert g.has_edge_between(0, 1)
+    assert g.has_edge_between(1, 2)
+    assert F.allclose(g.ndata["h"], F.tensor([[1.0], [2.0], [3.0]]))
+    assert F.allclose(
+        g.edata["h"], F.tensor([[1.0, 2.0], [1.0, 2.0], [2.0, 3.0], [2.0, 3.0]])
+    )
+
+
+@parametrize_idtype
+def test_apply_nodes(idtype):
+    def _upd(nodes):
+        return {"h": nodes.data["h"] * 2}
+
+    g = generate_graph(idtype)
+    old = g.ndata["h"]
+    g.apply_nodes(_upd)
+    assert F.allclose(old * 2, g.ndata["h"])
+    u = F.tensor([0, 3, 4, 6], g.idtype)
+    g.apply_nodes(lambda nodes: {"h": nodes.data["h"] * 0.0}, u)
+    assert F.allclose(F.gather_row(g.ndata["h"], u), F.zeros((4, D)))
+
+
+@parametrize_idtype
+def test_apply_edges(idtype):
+    def _upd(edges):
+        return {"w": edges.data["w"] * 2}
+
+    g = generate_graph(idtype)
+    old = g.edata["w"]
+    g.apply_edges(_upd)
+    assert F.allclose(old * 2, g.edata["w"])
+    u = F.tensor([0, 0, 0, 4, 5, 6], g.idtype)
+    v = F.tensor([1, 2, 3, 9, 9, 9], g.idtype)
+    g.apply_edges(lambda edges: {"w": edges.data["w"] * 0.0}, (u, v))
+    eid = F.tensor(g.edge_ids(u, v))
+    assert F.allclose(F.gather_row(g.edata["w"], eid), F.zeros((6, D)))
+
+
+@parametrize_idtype
+def test_update_routines(idtype):
+    g = generate_graph(idtype)
+
+    # send_and_recv
+    reduce_msg_shapes.clear()
+    u = [0, 0, 0, 4, 5, 6]
+    v = [1, 2, 3, 9, 9, 9]
+    g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
+    assert reduce_msg_shapes == {(1, 3, D), (3, 1, D)}
+    reduce_msg_shapes.clear()
+    try:
+        g.send_and_recv([u, v])
+        assert False
+    except:
+        pass
+
+    # pull
+    v = F.tensor([1, 2, 3, 9], g.idtype)
+    reduce_msg_shapes.clear()
+    g.pull(v, message_func, reduce_func, apply_node_func)
+    assert reduce_msg_shapes == {(1, 8, D), (3, 1, D)}
+    reduce_msg_shapes.clear()
+
+    # push
+    v = F.tensor([0, 1, 2, 3], g.idtype)
+    reduce_msg_shapes.clear()
+    g.push(v, message_func, reduce_func, apply_node_func)
+    assert reduce_msg_shapes == {(1, 3, D), (8, 1, D)}
+    reduce_msg_shapes.clear()
+
+    # update_all
+    reduce_msg_shapes.clear()
+    g.update_all(message_func, reduce_func, apply_node_func)
+    assert reduce_msg_shapes == {(1, 8, D), (9, 1, D)}
+    reduce_msg_shapes.clear()
+
+
+@parametrize_idtype
+def test_update_all_0deg(idtype):
+    # test#1
+    g = dgl.graph(([1, 2, 3, 4], [0, 0, 0, 0]), idtype=idtype, device=F.ctx())
+
+    def _message(edges):
+        return {"m": edges.src["h"]}
+
+    def _reduce(nodes):
+        return {"x": nodes.data["h"] + F.sum(nodes.mailbox["m"], 1)}
+
+    def _apply(nodes):
+        return {"x": nodes.data["x"] * 2}
+
+    def _init2(shape, dtype, ctx, ids):
+        return 2 + F.zeros(shape, dtype, ctx)
+
+    g.set_n_initializer(_init2, "x")
+    old_repr = F.randn((5, 5))
+    g.ndata["h"] = old_repr
+    g.update_all(_message, _reduce, _apply)
+    new_repr = g.ndata["x"]
+    # the first row of the new_repr should be the sum of all the node
+    # features; while the 0-deg nodes should be initialized by the
+    # initializer and applied with UDF.
+    assert F.allclose(new_repr[1:], 2 * (2 + F.zeros((4, 5))))
+    assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
+
+    # test#2: graph with no edge
+    g = dgl.graph(([], []), num_nodes=5, idtype=idtype, device=F.ctx())
+    g.ndata["h"] = old_repr
+    # Intercepting the warning: The input graph for the user-defined edge
+    # function does not contain valid edges.
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        g.update_all(
+            _message, _reduce, lambda nodes: {"h": nodes.data["h"] * 2}
+        )
+
+    new_repr = g.ndata["h"]
+    # should fallback to apply
+    assert F.allclose(new_repr, 2 * old_repr)
+
+
+@parametrize_idtype
+def test_pull_0deg(idtype):
+    g = dgl.graph(([0], [1]), idtype=idtype, device=F.ctx())
+
+    def _message(edges):
+        return {"m": edges.src["h"]}
+
+    def _reduce(nodes):
+        return {"x": nodes.data["h"] + F.sum(nodes.mailbox["m"], 1)}
+
+    def _apply(nodes):
+        return {"x": nodes.data["x"] * 2}
+
+    def _init2(shape, dtype, ctx, ids):
+        return 2 + F.zeros(shape, dtype, ctx)
+
+    g.set_n_initializer(_init2, "x")
+    # test#1: pull both 0deg and non-0deg nodes
+    old = F.randn((2, 5))
+    g.ndata["h"] = old
+    g.pull([0, 1], _message, _reduce, _apply)
+    new = g.ndata["x"]
+    # 0deg check: initialized with the func and got applied
+    assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
+    # non-0deg check
+    assert F.allclose(new[1], F.sum(old, 0) * 2)
+
+    # test#2: pull only 0deg node
+    old = F.randn((2, 5))
+    g.ndata["h"] = old
+    # Intercepting the warning: The input graph for the user-defined edge
+    # function does not contain valid edges
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        g.pull(0, _message, _reduce, lambda nodes: {"h": nodes.data["h"] * 2})
+
+    new = g.ndata["h"]
+    # 0deg check: fallback to apply
+    assert F.allclose(new[0], 2 * old[0])
+    # non-0deg check: not touched
+    assert F.allclose(new[1], old[1])
+
+
+def test_dynamic_addition():
+    N = 3
+    D = 1
+
+    g = dgl.graph([]).to(F.ctx())
+
+    # Test node addition
+    g.add_nodes(N)
+    g.ndata.update({"h1": F.randn((N, D)), "h2": F.randn((N, D))})
+    g.add_nodes(3)
+    assert g.ndata["h1"].shape[0] == g.ndata["h2"].shape[0] == N + 3
+
+    # Test edge addition
+    g.add_edges(0, 1)
+    g.add_edges(1, 0)
+    g.edata.update({"h1": F.randn((2, D)), "h2": F.randn((2, D))})
+    assert g.edata["h1"].shape[0] == g.edata["h2"].shape[0] == 2
+
+    g.add_edges([0, 2], [2, 0])
+    g.edata["h1"] = F.randn((4, D))
+    assert g.edata["h1"].shape[0] == g.edata["h2"].shape[0] == 4
+
+    g.add_edges(1, 2)
+    g.edges[4].data["h1"] = F.randn((1, D))
+    assert g.edata["h1"].shape[0] == g.edata["h2"].shape[0] == 5
+
+    # test add edge with part of the features
+    g.add_edges(2, 1, {"h1": F.randn((1, D))})
+    assert len(g.edata["h1"]) == len(g.edata["h2"])
+
+
+@parametrize_idtype
+def test_repr(idtype):
+    g = dgl.graph(
+        ([0, 0, 1], [1, 2, 2]), num_nodes=10, idtype=idtype, device=F.ctx()
+    )
+    repr_string = g.__repr__()
+    print(repr_string)
+    g.ndata["x"] = F.zeros((10, 5))
+    g.edata["y"] = F.zeros((3, 4))
+    repr_string = g.__repr__()
+    print(repr_string)
+
+
+@parametrize_idtype
+def test_local_var(idtype):
+    g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]), idtype=idtype, device=F.ctx())
+    g.ndata["h"] = F.zeros((g.num_nodes(), 3))
+    g.edata["w"] = F.zeros((g.num_edges(), 4))
+
+    # test override
+    def foo(g):
+        g = g.local_var()
+        g.ndata["h"] = F.ones((g.num_nodes(), 3))
+        g.edata["w"] = F.ones((g.num_edges(), 4))
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test out-place update
+    def foo(g):
+        g = g.local_var()
+        g.nodes[[2, 3]].data["h"] = F.ones((2, 3))
+        g.edges[[2, 3]].data["w"] = F.ones((2, 4))
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test out-place update 2
+    def foo(g):
+        g = g.local_var()
+        g.apply_nodes(lambda nodes: {"h": nodes.data["h"] + 10}, [2, 3])
+        g.apply_edges(lambda edges: {"w": edges.data["w"] + 10}, [2, 3])
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test auto-pop
+    def foo(g):
+        g = g.local_var()
+        g.ndata["hh"] = F.ones((g.num_nodes(), 3))
+        g.edata["ww"] = F.ones((g.num_edges(), 4))
+
+    foo(g)
+    assert "hh" not in g.ndata
+    assert "ww" not in g.edata
+
+    # test initializer1
+    g = dgl.graph(([0, 1], [1, 1]), idtype=idtype, device=F.ctx())
+    g.set_n_initializer(dgl.init.zero_initializer)
+
+    def foo(g):
+        g = g.local_var()
+        g.nodes[0].data["h"] = F.ones((1, 1))
+        assert F.allclose(g.ndata["h"], F.tensor([[1.0], [0.0]]))
+
+    foo(g)
+
+    # test initializer2
+    def foo_e_initializer(shape, dtype, ctx, id_range):
+        return F.ones(shape)
+
+    g.set_e_initializer(foo_e_initializer, field="h")
+
+    def foo(g):
+        g = g.local_var()
+        g.edges[0, 1].data["h"] = F.ones((1, 1))
+        assert F.allclose(g.edata["h"], F.ones((2, 1)))
+        g.edges[0, 1].data["w"] = F.ones((1, 1))
+        assert F.allclose(g.edata["w"], F.tensor([[1.0], [0.0]]))
+
+    foo(g)
+
+
+@parametrize_idtype
+def test_local_scope(idtype):
+    g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]), idtype=idtype, device=F.ctx())
+    g.ndata["h"] = F.zeros((g.num_nodes(), 3))
+    g.edata["w"] = F.zeros((g.num_edges(), 4))
+
+    # test override
+    def foo(g):
+        with g.local_scope():
+            g.ndata["h"] = F.ones((g.num_nodes(), 3))
+            g.edata["w"] = F.ones((g.num_edges(), 4))
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test out-place update
+    def foo(g):
+        with g.local_scope():
+            g.nodes[[2, 3]].data["h"] = F.ones((2, 3))
+            g.edges[[2, 3]].data["w"] = F.ones((2, 4))
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test out-place update 2
+    def foo(g):
+        with g.local_scope():
+            g.apply_nodes(lambda nodes: {"h": nodes.data["h"] + 10}, [2, 3])
+            g.apply_edges(lambda edges: {"w": edges.data["w"] + 10}, [2, 3])
+
+    foo(g)
+    assert F.allclose(g.ndata["h"], F.zeros((g.num_nodes(), 3)))
+    assert F.allclose(g.edata["w"], F.zeros((g.num_edges(), 4)))
+
+    # test auto-pop
+    def foo(g):
+        with g.local_scope():
+            g.ndata["hh"] = F.ones((g.num_nodes(), 3))
+            g.edata["ww"] = F.ones((g.num_edges(), 4))
+
+    foo(g)
+    assert "hh" not in g.ndata
+    assert "ww" not in g.edata
+
+    # test nested scope
+    def foo(g):
+        with g.local_scope():
+            g.ndata["hh"] = F.ones((g.num_nodes(), 3))
+            g.edata["ww"] = F.ones((g.num_edges(), 4))
+            with g.local_scope():
+                g.ndata["hhh"] = F.ones((g.num_nodes(), 3))
+                g.edata["www"] = F.ones((g.num_edges(), 4))
+            assert "hhh" not in g.ndata
+            assert "www" not in g.edata
+
+    foo(g)
+    assert "hh" not in g.ndata
+    assert "ww" not in g.edata
+
+    # test initializer1
+    g = dgl.graph(([0, 1], [1, 1]), idtype=idtype, device=F.ctx())
+    g.set_n_initializer(dgl.init.zero_initializer)
+
+    def foo(g):
+        with g.local_scope():
+            g.nodes[0].data["h"] = F.ones((1, 1))
+            assert F.allclose(g.ndata["h"], F.tensor([[1.0], [0.0]]))
+
+    foo(g)
+
+    # test initializer2
+    def foo_e_initializer(shape, dtype, ctx, id_range):
+        return F.ones(shape)
+
+    g.set_e_initializer(foo_e_initializer, field="h")
+
+    def foo(g):
+        with g.local_scope():
+            g.edges[0, 1].data["h"] = F.ones((1, 1))
+            assert F.allclose(g.edata["h"], F.ones((2, 1)))
+            g.edges[0, 1].data["w"] = F.ones((1, 1))
+            assert F.allclose(g.edata["w"], F.tensor([[1.0], [0.0]]))
+
+    foo(g)
+
+    # test exception handling
+    def foo(g):
+        try:
+            with g.local_scope():
+                g.ndata["hh"] = F.ones((g.num_nodes(), 1))
+                # throw TypeError
+                1 + "1"
+        except TypeError:
+            pass
+        assert "hh" not in g.ndata
+
+    foo(g)
+
+
+@parametrize_idtype
+def test_isolated_nodes(idtype):
+    g = dgl.graph(([0, 1], [1, 2]), num_nodes=5, idtype=idtype, device=F.ctx())
+    assert g.num_nodes() == 5
+
+    g = dgl.heterograph(
+        {("user", "plays", "game"): ([0, 0, 1], [2, 3, 2])},
+        {"user": 5, "game": 7},
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.num_nodes("user") == 5
+    assert g.num_nodes("game") == 7
+
+    # Test backward compatibility
+    g = dgl.heterograph(
+        {("user", "plays", "game"): ([0, 0, 1], [2, 3, 2])},
+        {"user": 5, "game": 7},
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.num_nodes("user") == 5
+    assert g.num_nodes("game") == 7
+
+
+@parametrize_idtype
+def test_send_multigraph(idtype):
+    g = dgl.graph(([0, 0, 0, 2], [1, 1, 1, 1]), idtype=idtype, device=F.ctx())
+
+    def _message_a(edges):
+        return {"a": edges.data["a"]}
+
+    def _message_b(edges):
+        return {"a": edges.data["a"] * 3}
+
+    def _reduce(nodes):
+        return {"a": F.max(nodes.mailbox["a"], 1)}
+
+    def answer(*args):
+        return F.max(F.stack(args, 0), 0)
+
+    assert g.is_multigraph
+
+    # send by eid
+    old_repr = F.randn((4, 5))
+    # send_and_recv_on
+    g.ndata["a"] = F.zeros((3, 5))
+    g.edata["a"] = old_repr
+    g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
+    new_repr = g.ndata["a"]
+    assert F.allclose(
+        new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3])
+    )
+    assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
+
+
+@parametrize_idtype
+def test_issue_1088(idtype):
+    # This test ensures that message passing on a heterograph with one edge type
+    # would not crash (GitHub issue #1088).
+    import dgl.function as fn
+
+    g = dgl.heterograph(
+        {("U", "E", "V"): ([0, 1, 2], [1, 2, 3])}, idtype=idtype, device=F.ctx()
+    )
+    g.nodes["U"].data["x"] = F.randn((3, 3))
+    g.update_all(fn.copy_u("x", "m"), fn.sum("m", "y"))
+
+
+@parametrize_idtype
+def test_degree_bucket_edge_ordering(idtype):
+    import dgl.function as fn
+
+    g = dgl.graph(
+        ([1, 3, 5, 0, 4, 2, 3, 3, 4, 5], [1, 1, 0, 0, 1, 2, 2, 0, 3, 3]),
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g.edata["eid"] = F.copy_to(F.arange(0, 10), F.ctx())
+
+    def reducer(nodes):
+        eid = F.asnumpy(F.copy_to(nodes.mailbox["eid"], F.cpu()))
+        assert np.array_equal(eid, np.sort(eid, 1))
+        return {"n": F.sum(nodes.mailbox["eid"], 1)}
+
+    g.update_all(fn.copy_e("eid", "eid"), reducer)
+
+
+@parametrize_idtype
+def test_issue_2484(idtype):
+    import dgl.function as fn
+
+    g = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
+    x = F.copy_to(F.randn((4,)), F.ctx())
+    g.ndata["x"] = x
+    g.pull([2, 1], fn.u_add_v("x", "x", "m"), fn.sum("m", "x"))
+    y1 = g.ndata["x"]
+
+    g.ndata["x"] = x
+    g.pull([1, 2], fn.u_add_v("x", "x", "m"), fn.sum("m", "x"))
+    y2 = g.ndata["x"]
+
+    assert F.allclose(y1, y2)

tests/python/common/ops/test_edge_softmax.py

@@ -0,0 +1,192 @@
+import itertools
+import math
+import unittest
+from collections import Counter
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as ssp
+from dgl import DGLError
+from dgl.ops import edge_softmax
+from scipy.sparse import rand
+from utils import get_cases, parametrize_idtype
+
+edge_softmax_shapes = [(1,), (1, 3), (3, 4, 5)]
+rfuncs = {"sum": fn.sum, "max": fn.max, "min": fn.min, "mean": fn.mean}
+fill_value = {"sum": 0, "max": float("-inf")}
+feat_size = 2
+
+
+@pytest.mark.parametrize("g", get_cases(["clique"]))
+@pytest.mark.parametrize("norm_by", ["src", "dst"])
+@pytest.mark.parametrize("shp", edge_softmax_shapes)
+@parametrize_idtype
+def test_edge_softmax(g, norm_by, shp, idtype):
+    g = g.astype(idtype).to(F.ctx())
+    edata = F.tensor(np.random.rand(g.num_edges(), *shp))
+    e1 = F.attach_grad(F.clone(edata))
+
+    with F.record_grad():
+        score1 = edge_softmax(g, e1, norm_by=norm_by)
+        F.backward(F.reduce_sum(score1))
+        grad_edata = F.grad(e1)
+
+    with F.record_grad():
+        e2 = F.attach_grad(F.clone(edata))
+        e2_2d = F.reshape(
+            e2,
+            (g.number_of_src_nodes(), g.number_of_dst_nodes(), *e2.shape[1:]),
+        )
+        if norm_by == "src":
+            score2 = F.softmax(e2_2d, 1)
+            score2 = F.reshape(score2, (-1, *e2.shape[1:]))
+        if norm_by == "dst":
+            score2 = F.softmax(e2_2d, 0)
+            score2 = F.reshape(score2, (-1, *e2.shape[1:]))
+        assert F.allclose(score1, score2)
+        print("forward passed")
+
+        F.backward(F.reduce_sum(score2))
+        assert F.allclose(F.grad(e2), grad_edata)
+        print("backward passed")
+
+
+def create_test_heterograph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1, 2, 1, 1], [0, 0, 1, 1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 1, 1], [0, 0, 1]),
+            ("developer", "develops", "game"): ([0, 1, 0], [0, 1, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+def test_edge_softmax_unidirectional():
+    g = dgl.heterograph(
+        {
+            ("A", "AB", "B"): (
+                [1, 2, 3, 1, 2, 3, 1, 2, 3],
+                [0, 0, 0, 1, 1, 1, 2, 2, 2],
+            ),
+            ("B", "BB", "B"): (
+                [0, 1, 2, 0, 1, 2, 0, 1, 2],
+                [0, 0, 0, 1, 1, 1, 2, 2, 2],
+            ),
+        }
+    )
+    g = g.to(F.ctx())
+    g.edges["AB"].data["x"] = F.ones(9) * 2
+    g.edges["BB"].data["x"] = F.ones(9)
+    result = dgl.ops.edge_softmax(
+        g, {"AB": g.edges["AB"].data["x"], "BB": g.edges["BB"].data["x"]}
+    )
+
+    ab = result["A", "AB", "B"]
+    bb = result["B", "BB", "B"]
+    e2 = F.zeros_like(ab) + math.exp(2) / ((math.exp(2) + math.exp(1)) * 3)
+    e1 = F.zeros_like(bb) + math.exp(1) / ((math.exp(2) + math.exp(1)) * 3)
+    assert F.allclose(ab, e2)
+    assert F.allclose(bb, e1)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@pytest.mark.parametrize("g", get_cases(["clique"]))
+@pytest.mark.parametrize("norm_by", ["src", "dst"])
+# @pytest.mark.parametrize('shp', edge_softmax_shapes)
+@parametrize_idtype
+def test_edge_softmax(g, norm_by, idtype):
+    print("params", norm_by, idtype)
+
+    g = create_test_heterograph(idtype)
+
+    x1 = F.randn((g.num_edges("plays"), feat_size))
+    x2 = F.randn((g.num_edges("follows"), feat_size))
+    x3 = F.randn((g.num_edges("develops"), feat_size))
+    x4 = F.randn((g.num_edges("wishes"), feat_size))
+
+    F.attach_grad(F.clone(x1))
+    F.attach_grad(F.clone(x2))
+    F.attach_grad(F.clone(x3))
+    F.attach_grad(F.clone(x4))
+
+    g["plays"].edata["eid"] = x1
+    g["follows"].edata["eid"] = x2
+    g["develops"].edata["eid"] = x3
+    g["wishes"].edata["eid"] = x4
+
+    #################################################################
+    #  edge_softmax() on homogeneous graph
+    #################################################################
+
+    with F.record_grad():
+        hm_g = dgl.to_homogeneous(g)
+        hm_x = F.cat((x3, x2, x1, x4), 0)
+        hm_e = F.attach_grad(F.clone(hm_x))
+        score_hm = edge_softmax(hm_g, hm_e, norm_by=norm_by)
+        hm_g.edata["score"] = score_hm
+        ht_g = dgl.to_heterogeneous(hm_g, g.ntypes, g.etypes)
+        r1 = ht_g.edata["score"][("user", "plays", "game")]
+        r2 = ht_g.edata["score"][("user", "follows", "user")]
+        r3 = ht_g.edata["score"][("developer", "develops", "game")]
+        r4 = ht_g.edata["score"][("user", "wishes", "game")]
+        F.backward(F.reduce_sum(r1) + F.reduce_sum(r2))
+        grad_edata_hm = F.grad(hm_e)
+
+    #################################################################
+    #  edge_softmax() on heterogeneous graph
+    #################################################################
+
+    e1 = F.attach_grad(F.clone(x1))
+    e2 = F.attach_grad(F.clone(x2))
+    e3 = F.attach_grad(F.clone(x3))
+    e4 = F.attach_grad(F.clone(x4))
+    e = {
+        ("user", "follows", "user"): e2,
+        ("user", "plays", "game"): e1,
+        ("user", "wishes", "game"): e4,
+        ("developer", "develops", "game"): e3,
+    }
+    with F.record_grad():
+        score = edge_softmax(g, e, norm_by=norm_by)
+        r5 = score[("user", "plays", "game")]
+        r6 = score[("user", "follows", "user")]
+        r7 = score[("developer", "develops", "game")]
+        r8 = score[("user", "wishes", "game")]
+        F.backward(F.reduce_sum(r5) + F.reduce_sum(r6))
+        grad_edata_ht = F.cat(
+            (F.grad(e3), F.grad(e2), F.grad(e1), F.grad(e4)), 0
+        )
+        # correctness check
+        assert F.allclose(r1, r5)
+        assert F.allclose(r2, r6)
+        assert F.allclose(r3, r7)
+        assert F.allclose(r4, r8)
+        assert F.allclose(grad_edata_hm, grad_edata_ht)
+
+
+if __name__ == "__main__":
+    test_edge_softmax_unidirectional()

tests/python/common/ops/test_ops.py

@@ -0,0 +1,510 @@
+import random
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+import pytest
+import torch
+from dgl.ops import gather_mm, gsddmm, gspmm, segment_reduce
+from utils import parametrize_idtype
+from utils.graph_cases import get_cases
+
+# Set seeds to make tests fully reproducible.
+SEED = 12345  # random.randint(1, 99999)
+random.seed(SEED)
+np.random.seed(SEED)
+dgl.seed(SEED)
+F.seed(SEED)
+
+udf_msg = {
+    "add": lambda edges: {"m": edges.src["x"] + edges.data["w"]},
+    "sub": lambda edges: {"m": edges.src["x"] - edges.data["w"]},
+    "mul": lambda edges: {"m": edges.src["x"] * edges.data["w"]},
+    "div": lambda edges: {"m": edges.src["x"] / edges.data["w"]},
+    "copy_lhs": lambda edges: {"m": edges.src["x"]},
+    "copy_rhs": lambda edges: {"m": edges.data["w"]},
+}
+
+
+def select(target, src, edge, dst):
+    if target == "u":
+        return src
+    elif target == "v":
+        return dst
+    elif target == "e":
+        return edge
+
+
+def binary_op(msg, x, y):
+    if msg == "add":
+        return x + y
+    elif msg == "sub":
+        return x - y
+    elif msg == "mul":
+        return x * y
+    elif msg == "div":
+        return x / y
+    elif msg == "dot":
+        return F.sum(x * y, -1, keepdims=True)
+    elif msg == "copy_lhs":
+        return x
+    elif msg == "copy_rhs":
+        return y
+
+
+def edge_func(lhs_target, rhs_target, msg):
+    def foo(edges):
+        return {
+            "m": binary_op(
+                msg,
+                select(lhs_target, edges.src, edges.data, edges.dst)["x"],
+                select(rhs_target, edges.src, edges.data, edges.dst)["y"],
+            )
+        }
+
+    return foo
+
+
+udf_apply_edges = {
+    lhs_target
+    + "_"
+    + msg
+    + "_"
+    + rhs_target: edge_func(lhs_target, rhs_target, msg)
+    for lhs_target in ["u", "v", "e"]
+    for rhs_target in ["u", "v", "e"]
+    for msg in ["add", "sub", "mul", "div", "dot", "copy_lhs", "copy_rhs"]
+}
+
+udf_reduce = {
+    "sum": lambda nodes: {"v": F.sum(nodes.mailbox["m"], 1)},
+    "min": lambda nodes: {"v": F.min(nodes.mailbox["m"], 1)},
+    "max": lambda nodes: {"v": F.max(nodes.mailbox["m"], 1)},
+}
+
+graphs = [
+    #    dgl.rand_graph(30, 0),
+    dgl.rand_graph(30, 100),
+    dgl.rand_bipartite("_U", "_E", "_V", 30, 40, 300),
+]
+
+spmm_shapes = [
+    ((1, 2, 1, 3, 1), (4, 1, 3, 1, 1)),
+    ((3, 3), (1, 3)),
+    ((1,), (3,)),
+    ((3,), (1,)),
+    ((1,), (1,)),
+    ((), ()),
+]
+
+sddmm_shapes = [
+    ((1, 2, 1, 3, 1), (4, 1, 3, 1, 1)),
+    ((5, 3, 1, 7), (1, 3, 7, 7)),
+    ((1, 3, 3), (4, 1, 3)),
+    ((3,), (3,)),
+    ((1,), (1,)),
+]
+
+
+@pytest.mark.parametrize("g", graphs)
+@pytest.mark.parametrize("shp", spmm_shapes)
+@pytest.mark.parametrize(
+    "msg", ["add", "sub", "mul", "div", "copy_lhs", "copy_rhs"]
+)
+@pytest.mark.parametrize("reducer", ["sum", "min", "max"])
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_spmm(idtype, dtype, g, shp, msg, reducer):
+    g = g.astype(idtype).to(F.ctx())
+    print(g)
+    print(g.idtype)
+
+    hu = F.tensor(
+        np.random.rand(*((g.number_of_src_nodes(),) + shp[0])).astype(dtype) + 1
+    )
+    he = F.tensor(
+        np.random.rand(*((g.num_edges(),) + shp[1])).astype(dtype) + 1
+    )
+    print("u shape: {}, e shape: {}".format(F.shape(hu), F.shape(he)))
+
+    g.srcdata["x"] = F.attach_grad(F.clone(hu))
+    g.edata["w"] = F.attach_grad(F.clone(he))
+    print("SpMM(message func: {}, reduce func: {})".format(msg, reducer))
+
+    u = F.attach_grad(F.clone(hu))
+    e = F.attach_grad(F.clone(he))
+    with F.record_grad():
+        v = gspmm(g, msg, reducer, u, e)
+        if reducer in ["max", "min"]:
+            v = F.replace_inf_with_zero(v)
+        if g.num_edges() > 0:
+            F.backward(F.reduce_sum(v))
+            if msg != "copy_rhs":
+                grad_u = F.grad(u)
+            if msg != "copy_lhs":
+                grad_e = F.grad(e)
+
+    with F.record_grad():
+        g.update_all(udf_msg[msg], udf_reduce[reducer])
+        if g.num_edges() > 0:
+            v1 = g.dstdata["v"]
+            assert F.allclose(v, v1)
+            print("forward passed")
+
+            F.backward(F.reduce_sum(v1))
+            if msg != "copy_rhs":
+                if reducer in [
+                    "min",
+                    "max",
+                ]:  # there might be some numerical errors
+                    rate = F.reduce_sum(
+                        F.abs(F.grad(g.srcdata["x"]) - grad_u)
+                    ) / F.reduce_sum(F.abs(grad_u))
+                    assert F.as_scalar(rate) < 1e-2, rate
+                else:
+                    assert F.allclose(F.grad(g.srcdata["x"]), grad_u)
+            if msg != "copy_lhs":
+                if reducer in ["min", "max"]:
+                    rate = F.reduce_sum(
+                        F.abs(F.grad(g.edata["w"]) - grad_e)
+                    ) / F.reduce_sum(F.abs(grad_e))
+                    assert F.as_scalar(rate) < 1e-2, rate
+                else:
+                    assert F.allclose(F.grad(g.edata["w"]), grad_e)
+            print("backward passed")
+
+    g.srcdata.pop("x")
+    g.edata.pop("w")
+    if "v" in g.dstdata:
+        g.dstdata.pop("v")
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Only support PyTorch for now.",
+)
+@parametrize_idtype
+@pytest.mark.parametrize(
+    "dtype, rtol, atol",
+    [(torch.float16, 1e-3, 0.5), (torch.bfloat16, 4e-3, 2.0)],
+)
+def test_half_spmm(idtype, dtype, rtol, atol):
+    if F._default_context_str == "cpu" and dtype == torch.float16:
+        pytest.skip("float16 is not supported on CPU.")
+    if (
+        F._default_context_str == "gpu"
+        and dtype == torch.bfloat16
+        and not torch.cuda.is_bf16_supported()
+    ):
+        pytest.skip("BF16 is not supported.")
+
+    # make sure the spmm result is < 512 to match the rtol/atol we set.
+    g = dgl.graph(
+        (torch.arange(900), torch.tensor([0] * 900)),
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    feat_fp32 = torch.rand((g.num_src_nodes(), 32)).to(F.ctx())
+    feat_half = feat_fp32.to(dtype)
+
+    # test SpMMCSR
+    g = g.formats(["csc"])
+    res_fp32 = dgl.ops.copy_u_sum(g, feat_fp32)[0]
+    res_half = dgl.ops.copy_u_sum(g, feat_half)[0].float()
+    assert torch.allclose(res_fp32, res_half, rtol=rtol, atol=atol)
+
+    # test SpMMCOO
+    # TODO(Xin): half-precision SpMMCoo is temporally disabled.
+    # g = g.formats(['coo'])
+    # res_fp32 = dgl.ops.copy_u_sum(g, feat_fp32)[0]
+    # res_half = dgl.ops.copy_u_sum(g, feat_half)[0].float()
+    # assert torch.allclose(res_fp32, res_half, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("g", graphs)
+@pytest.mark.parametrize("shp", sddmm_shapes)
+@pytest.mark.parametrize("lhs_target", ["u", "v", "e"])
+@pytest.mark.parametrize("rhs_target", ["u", "v", "e"])
+@pytest.mark.parametrize(
+    "msg", ["add", "sub", "mul", "div", "dot", "copy_lhs", "copy_rhs"]
+)
+@parametrize_idtype
+def test_sddmm(g, shp, lhs_target, rhs_target, msg, idtype):
+    if lhs_target == rhs_target:
+        return
+    g = g.astype(idtype).to(F.ctx())
+    if dgl.backend.backend_name == "mxnet" and g.num_edges() == 0:
+        pytest.skip()  # mxnet do not support zero shape tensor
+    print(g)
+    print(g.idtype)
+
+    len_lhs = select(
+        lhs_target,
+        g.number_of_src_nodes(),
+        g.num_edges(),
+        g.number_of_dst_nodes(),
+    )
+    lhs_shp = (len_lhs,) + shp[0]
+    len_rhs = select(
+        rhs_target,
+        g.number_of_src_nodes(),
+        g.num_edges(),
+        g.number_of_dst_nodes(),
+    )
+    rhs_shp = (len_rhs,) + shp[1]
+    feat_lhs = F.tensor(np.random.rand(*lhs_shp) + 1)
+    feat_rhs = F.tensor(np.random.rand(*rhs_shp) + 1)
+    print(
+        "lhs shape: {}, rhs shape: {}".format(
+            F.shape(feat_lhs), F.shape(feat_rhs)
+        )
+    )
+
+    lhs_frame = select(lhs_target, g.srcdata, g.edata, g.dstdata)
+    rhs_frame = select(rhs_target, g.srcdata, g.edata, g.dstdata)
+    lhs_frame["x"] = F.attach_grad(F.clone(feat_lhs))
+    rhs_frame["y"] = F.attach_grad(F.clone(feat_rhs))
+    msg_func = lhs_target + "_" + msg + "_" + rhs_target
+    print("SDDMM(message func: {})".format(msg_func))
+
+    lhs = F.attach_grad(F.clone(feat_lhs))
+    rhs = F.attach_grad(F.clone(feat_rhs))
+    with F.record_grad():
+        e = gsddmm(
+            g, msg, lhs, rhs, lhs_target=lhs_target, rhs_target=rhs_target
+        )
+        F.backward(F.reduce_sum(e))
+        grad_lhs = F.grad(lhs)
+        grad_rhs = F.grad(rhs)
+
+    with F.record_grad():
+        g.apply_edges(udf_apply_edges[msg_func])
+        if g.num_edges() > 0:
+            e1 = g.edata["m"]
+            assert F.allclose(e, e1)
+            print("forward passed")
+
+            F.backward(F.reduce_sum(e1))
+            if msg != "copy_rhs":
+                assert F.allclose(F.grad(lhs_frame["x"]), grad_lhs)
+            if msg != "copy_lhs":
+                assert F.allclose(F.grad(rhs_frame["y"]), grad_rhs)
+            print("backward passed")
+
+    lhs_frame.pop("x")
+    rhs_frame.pop("y")
+    if "m" in g.edata:
+        g.edata.pop("m")
+
+
+@pytest.mark.parametrize("reducer", ["sum", "max", "min", "mean"])
+def test_segment_reduce(reducer):
+    ctx = F.ctx()
+    value = F.tensor(np.random.rand(10, 5))
+    v1 = F.attach_grad(F.clone(value))
+    v2 = F.attach_grad(F.clone(value))
+    seglen = F.tensor([2, 3, 0, 4, 1, 0, 0])
+    u = F.copy_to(F.arange(0, F.shape(value)[0], F.int32), ctx)
+    v = F.repeat(
+        F.copy_to(F.arange(0, len(seglen), F.int32), ctx), seglen, dim=0
+    )
+
+    num_nodes = {"_U": len(u), "_V": len(seglen)}
+    g = dgl.convert.heterograph(
+        {("_U", "_E", "_V"): (u, v)}, num_nodes_dict=num_nodes
+    )
+    with F.record_grad():
+        rst1 = gspmm(g, "copy_lhs", reducer, v1, None)
+        if reducer in ["max", "min"]:
+            rst1 = F.replace_inf_with_zero(rst1)
+        F.backward(F.reduce_sum(rst1))
+        grad1 = F.grad(v1)
+
+    with F.record_grad():
+        rst2 = segment_reduce(seglen, v2, reducer=reducer)
+        F.backward(F.reduce_sum(rst2))
+        assert F.allclose(rst1, rst2)
+        print("forward passed")
+
+        grad2 = F.grad(v2)
+        assert F.allclose(grad1, grad2)
+        print("backward passed")
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [1, 8, 16, 64, 256])
+@pytest.mark.parametrize(
+    "dtype, tol",
+    [
+        (torch.float16, 1e-2),
+        (torch.bfloat16, 1e-2),
+        (torch.float32, 3e-3),
+        (torch.float64, 1e-4),
+    ],
+)
+def test_segment_mm(idtype, feat_size, dtype, tol):
+    if F._default_context_str == "cpu" and dtype == torch.float16:
+        pytest.skip("float16 is not supported on CPU.")
+    if (
+        F._default_context_str == "gpu"
+        and dtype == torch.bfloat16
+        and not torch.cuda.is_bf16_supported()
+    ):
+        pytest.skip("BF16 is not supported.")
+    dev = F.ctx()
+    # input
+    a = torch.tensor(np.random.rand(100, feat_size)).to(dev).to(dtype)
+    a.requires_grad_()
+    b = (
+        torch.tensor(np.random.rand(10, feat_size, feat_size + 1))
+        .to(dev)
+        .to(dtype)
+    )
+    b.requires_grad_()
+    seglen_a = torch.tensor([10, 15, 8, 0, 1, 9, 18, 24, 15, 0]).to(idtype)
+    dc = torch.tensor(np.random.rand(100, feat_size + 1)).to(dev).to(dtype)
+    # compute
+    c = dgl.ops.segment_mm(a, b, seglen_a)
+    c.backward(dc)
+    da = a.grad.clone()
+    db = b.grad.clone()
+    # ground truth
+    c_t = []
+    off = 0
+    for i, l in enumerate(seglen_a):
+        c_t.append(a[off : off + l] @ b[i])
+        off += l
+    c_t = torch.cat(c_t).to(dtype)
+    a.grad.zero_()
+    b.grad.zero_()
+    c_t.backward(dc)
+    da_t = a.grad
+    db_t = b.grad
+
+    assert torch.allclose(c, c_t, atol=tol, rtol=tol)
+    assert torch.allclose(da, da_t, atol=tol, rtol=tol)
+    assert torch.allclose(db, db_t, atol=tol, rtol=tol)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@pytest.mark.parametrize("feat_size", [1, 8, 16, 64, 256])
+@pytest.mark.parametrize(
+    "dtype, tol",
+    [
+        (torch.float16, 1e-2),
+        (torch.bfloat16, 2e-2),
+        (torch.float32, 3e-3),
+        (torch.float64, 1e-4),
+    ],
+)
+def test_gather_mm_idx_b(feat_size, dtype, tol):
+    if F._default_context_str == "cpu" and dtype == torch.float16:
+        pytest.skip("float16 is not supported on CPU.")
+
+    if F._default_context_str == "gpu":
+        if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+            pytest.skip("BF16 is not supported.")
+
+        if (
+            dtype == torch.float16
+            and torch.cuda.get_device_capability() < (7, 0)
+        ) or (
+            dtype == torch.bfloat16
+            and torch.cuda.get_device_capability() < (8, 0)
+        ):
+            pytest.skip(
+                f"{dtype} is not supported for atomic operations on GPU with "
+                f"cuda capability ({torch.cuda.get_device_capability()})."
+            )
+
+    dev = F.ctx()
+    # input
+    a = torch.tensor(np.random.rand(100, feat_size)).to(dev).to(dtype)
+    a.requires_grad_()
+    b = (
+        torch.tensor(np.random.rand(10, feat_size, feat_size + 1))
+        .to(dev)
+        .to(dtype)
+    )
+    b.requires_grad_()
+    idx = torch.tensor(np.random.randint(0, 10, 100)).to(dev).long()
+    dc = torch.tensor(np.random.rand(100, feat_size + 1)).to(dev).to(dtype)
+    # compute
+    c = gather_mm(a, b, idx_b=idx)
+    c.backward(dc)
+    da = a.grad.clone()
+    db = b.grad.clone()
+    # ground truth
+    c_t = torch.bmm(a.unsqueeze(1), b[idx]).squeeze(1)
+    a.grad.zero_()
+    b.grad.zero_()
+    c_t.backward(dc)
+    da_t = a.grad
+    db_t = b.grad
+
+    assert torch.allclose(c, c_t, atol=tol, rtol=tol)
+    assert torch.allclose(da, da_t, atol=tol, rtol=tol)
+    assert torch.allclose(db, db_t, atol=tol, rtol=tol)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [1, 8, 16, 64, 256])
+def _test_gather_mm_idx_a(idtype, feat_size):
+    # TODO(minjie): currently disabled due to bugs in the CUDA kernel. Need to fix it later.
+    import torch
+
+    dev = F.ctx()
+    # input
+    a = torch.tensor(np.random.rand(10, feat_size)).to(dev)
+    a.requires_grad_()
+    b = torch.tensor(np.random.rand(100, feat_size, feat_size + 1)).to(dev)
+    b.requires_grad_()
+    idx = torch.tensor(np.random.randint(0, 10, 100)).to(dev)
+    dc = torch.tensor(np.random.rand(100, feat_size + 1)).to(dev)
+    # compute
+    c = gather_mm(a, b, idx_a=idx)
+    c.backward(dc)
+    da = a.grad.clone()
+    db = b.grad.clone()
+    # ground truth
+    c_t = torch.bmm(a[idx].unsqueeze(1), b).squeeze(1)
+    a.grad.zero_()
+    b.grad.zero_()
+    c_t.backward(dc)
+    da_t = a.grad
+    db_t = b.grad
+
+    assert torch.allclose(c, c_t, atol=1e-4, rtol=1e-4)
+    assert torch.allclose(da, da_t, atol=1e-4, rtol=1e-4)
+    assert torch.allclose(db, db_t, atol=1e-4, rtol=1e-4)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@unittest.skipIf(
+    F._default_context_str == "gpu", reason="Libxsmm only fit in CPU."
+)
+def test_use_libxsmm_switch():
+    import torch
+
+    g = dgl.graph(([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]))
+    x = torch.ones(3, 2, requires_grad=True)
+    y = torch.arange(1, 13).float().view(6, 2).requires_grad_()
+
+    dgl.use_libxsmm(False)
+    assert ~dgl.is_libxsmm_enabled()
+    dgl.ops.u_mul_e_sum(g, x, y)
+    dgl.use_libxsmm(True)
+    assert dgl.is_libxsmm_enabled()
+    dgl.ops.u_mul_e_sum(g, x, y)

tests/python/common/test_batch-graph.py

@@ -0,0 +1,332 @@
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+from utils import parametrize_idtype
+
+
+def tree1(idtype):
+    """Generate a tree
+         0
+        / \
+       1   2
+      / \
+     3   4
+    Edges are from leaves to root.
+    """
+    g = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g.add_nodes(5)
+    g.add_edges(3, 1)
+    g.add_edges(4, 1)
+    g.add_edges(1, 0)
+    g.add_edges(2, 0)
+    g.ndata["h"] = F.tensor([0, 1, 2, 3, 4])
+    g.edata["h"] = F.randn((4, 10))
+    return g
+
+
+def tree2(idtype):
+    """Generate a tree
+         1
+        / \
+       4   3
+      / \
+     2   0
+    Edges are from leaves to root.
+    """
+    g = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g.add_nodes(5)
+    g.add_edges(2, 4)
+    g.add_edges(0, 4)
+    g.add_edges(4, 1)
+    g.add_edges(3, 1)
+    g.ndata["h"] = F.tensor([0, 1, 2, 3, 4])
+    g.edata["h"] = F.randn((4, 10))
+    return g
+
+
+@parametrize_idtype
+def test_batch_unbatch(idtype):
+    t1 = tree1(idtype)
+    t2 = tree2(idtype)
+
+    bg = dgl.batch([t1, t2])
+    assert bg.num_nodes() == 10
+    assert bg.num_edges() == 8
+    assert bg.batch_size == 2
+    assert F.allclose(bg.batch_num_nodes(), F.tensor([5, 5]))
+    assert F.allclose(bg.batch_num_edges(), F.tensor([4, 4]))
+
+    tt1, tt2 = dgl.unbatch(bg)
+    assert F.allclose(t1.ndata["h"], tt1.ndata["h"])
+    assert F.allclose(t1.edata["h"], tt1.edata["h"])
+    assert F.allclose(t2.ndata["h"], tt2.ndata["h"])
+    assert F.allclose(t2.edata["h"], tt2.edata["h"])
+
+
+@parametrize_idtype
+def test_batch_unbatch1(idtype):
+    t1 = tree1(idtype)
+    t2 = tree2(idtype)
+    b1 = dgl.batch([t1, t2])
+    b2 = dgl.batch([t2, b1])
+    assert b2.num_nodes() == 15
+    assert b2.num_edges() == 12
+    assert b2.batch_size == 3
+    assert F.allclose(b2.batch_num_nodes(), F.tensor([5, 5, 5]))
+    assert F.allclose(b2.batch_num_edges(), F.tensor([4, 4, 4]))
+
+    s1, s2, s3 = dgl.unbatch(b2)
+    assert F.allclose(t2.ndata["h"], s1.ndata["h"])
+    assert F.allclose(t2.edata["h"], s1.edata["h"])
+    assert F.allclose(t1.ndata["h"], s2.ndata["h"])
+    assert F.allclose(t1.edata["h"], s2.edata["h"])
+    assert F.allclose(t2.ndata["h"], s3.ndata["h"])
+    assert F.allclose(t2.edata["h"], s3.edata["h"])
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="TF doesn't support inplace update",
+)
+@parametrize_idtype
+def test_batch_unbatch_frame(idtype):
+    """Test module of node/edge frames of batched/unbatched DGLGraphs.
+    Also address the bug mentioned in https://github.com/dmlc/dgl/issues/1475.
+    """
+    t1 = tree1(idtype)
+    t2 = tree2(idtype)
+    N1 = t1.num_nodes()
+    E1 = t1.num_edges()
+    N2 = t2.num_nodes()
+    E2 = t2.num_edges()
+    D = 10
+    t1.ndata["h"] = F.randn((N1, D))
+    t1.edata["h"] = F.randn((E1, D))
+    t2.ndata["h"] = F.randn((N2, D))
+    t2.edata["h"] = F.randn((E2, D))
+
+    b1 = dgl.batch([t1, t2])
+    b2 = dgl.batch([t2])
+    b1.ndata["h"][:N1] = F.zeros((N1, D))
+    b1.edata["h"][:E1] = F.zeros((E1, D))
+    b2.ndata["h"][:N2] = F.zeros((N2, D))
+    b2.edata["h"][:E2] = F.zeros((E2, D))
+    assert not F.allclose(t1.ndata["h"], F.zeros((N1, D)))
+    assert not F.allclose(t1.edata["h"], F.zeros((E1, D)))
+    assert not F.allclose(t2.ndata["h"], F.zeros((N2, D)))
+    assert not F.allclose(t2.edata["h"], F.zeros((E2, D)))
+
+    g1, g2 = dgl.unbatch(b1)
+    (_g2,) = dgl.unbatch(b2)
+    assert F.allclose(g1.ndata["h"], F.zeros((N1, D)))
+    assert F.allclose(g1.edata["h"], F.zeros((E1, D)))
+    assert F.allclose(g2.ndata["h"], t2.ndata["h"])
+    assert F.allclose(g2.edata["h"], t2.edata["h"])
+    assert F.allclose(_g2.ndata["h"], F.zeros((N2, D)))
+    assert F.allclose(_g2.edata["h"], F.zeros((E2, D)))
+
+
+@parametrize_idtype
+def test_batch_unbatch2(idtype):
+    # test setting/getting features after batch
+    a = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    a.add_nodes(4)
+    a.add_edges(0, [1, 2, 3])
+    b = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    b.add_nodes(3)
+    b.add_edges(0, [1, 2])
+    c = dgl.batch([a, b])
+    c.ndata["h"] = F.ones((7, 1))
+    c.edata["w"] = F.ones((5, 1))
+    assert F.allclose(c.ndata["h"], F.ones((7, 1)))
+    assert F.allclose(c.edata["w"], F.ones((5, 1)))
+
+
+@parametrize_idtype
+def test_batch_send_and_recv(idtype):
+    t1 = tree1(idtype)
+    t2 = tree2(idtype)
+
+    bg = dgl.batch([t1, t2])
+    _mfunc = lambda edges: {"m": edges.src["h"]}
+    _rfunc = lambda nodes: {"h": F.sum(nodes.mailbox["m"], 1)}
+    u = [3, 4, 2 + 5, 0 + 5]
+    v = [1, 1, 4 + 5, 4 + 5]
+
+    bg.send_and_recv((u, v), _mfunc, _rfunc)
+
+    t1, t2 = dgl.unbatch(bg)
+    assert F.asnumpy(t1.ndata["h"][1]) == 7
+    assert F.asnumpy(t2.ndata["h"][4]) == 2
+
+
+@parametrize_idtype
+def test_batch_propagate(idtype):
+    t1 = tree1(idtype)
+    t2 = tree2(idtype)
+
+    bg = dgl.batch([t1, t2])
+    _mfunc = lambda edges: {"m": edges.src["h"]}
+    _rfunc = lambda nodes: {"h": F.sum(nodes.mailbox["m"], 1)}
+    # get leaves.
+
+    order = []
+
+    # step 1
+    u = [3, 4, 2 + 5, 0 + 5]
+    v = [1, 1, 4 + 5, 4 + 5]
+    order.append((u, v))
+
+    # step 2
+    u = [1, 2, 4 + 5, 3 + 5]
+    v = [0, 0, 1 + 5, 1 + 5]
+    order.append((u, v))
+
+    bg.prop_edges(order, _mfunc, _rfunc)
+    t1, t2 = dgl.unbatch(bg)
+
+    assert F.asnumpy(t1.ndata["h"][0]) == 9
+    assert F.asnumpy(t2.ndata["h"][1]) == 5
+
+
+@parametrize_idtype
+def test_batched_edge_ordering(idtype):
+    g1 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g1.add_nodes(6)
+    g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
+    e1 = F.randn((5, 10))
+    g1.edata["h"] = e1
+    g2 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g2.add_nodes(6)
+    g2.add_edges([0, 1, 2, 5, 4, 5], [1, 2, 3, 4, 3, 0])
+    e2 = F.randn((6, 10))
+    g2.edata["h"] = e2
+    g = dgl.batch([g1, g2])
+    r1 = g.edata["h"][g.edge_ids(4, 5)]
+    r2 = g1.edata["h"][g1.edge_ids(4, 5)]
+    assert F.array_equal(r1, r2)
+
+
+@parametrize_idtype
+def test_batch_no_edge(idtype):
+    g1 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g1.add_nodes(6)
+    g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
+    g2 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g2.add_nodes(6)
+    g2.add_edges([0, 1, 2, 5, 4, 5], [1, 2, 3, 4, 3, 0])
+    g3 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g3.add_nodes(1)  # no edges
+    g = dgl.batch([g1, g3, g2])  # should not throw an error
+
+
+@parametrize_idtype
+def test_batch_keeps_empty_data(idtype):
+    g1 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g1.ndata["nh"] = F.tensor([])
+    g1.edata["eh"] = F.tensor([])
+    g2 = dgl.graph(([], [])).astype(idtype).to(F.ctx())
+    g2.ndata["nh"] = F.tensor([])
+    g2.edata["eh"] = F.tensor([])
+    g = dgl.batch([g1, g2])
+    assert "nh" in g.ndata
+    assert "eh" in g.edata
+
+
+def _get_subgraph_batch_info(keys, induced_indices_arr, batch_num_objs):
+    """Internal function to compute batch information for subgraphs.
+    Parameters
+    ----------
+    keys : List[str]
+        The node/edge type keys.
+    induced_indices_arr : List[Tensor]
+        The induced node/edge index tensor for all node/edge types.
+    batch_num_objs : Tensor
+        Number of nodes/edges for each graph in the original batch.
+    Returns
+    -------
+    Mapping[str, Tensor]
+        A dictionary mapping all node/edge type keys to the ``batch_num_objs``
+        array of corresponding graph.
+    """
+    bucket_offset = np.expand_dims(
+        np.cumsum(F.asnumpy(batch_num_objs), 0), -1
+    )  # (num_bkts, 1)
+    ret = {}
+    for key, induced_indices in zip(keys, induced_indices_arr):
+        # NOTE(Zihao): this implementation is not efficient and we can replace it with
+        # binary search in the future.
+        induced_indices = np.expand_dims(
+            F.asnumpy(induced_indices), 0
+        )  # (1, num_nodes)
+        new_offset = np.sum((induced_indices < bucket_offset), 1)  # (num_bkts,)
+        # start_offset = [0] + [new_offset[i-1] for i in range(1, n_bkts)]
+        start_offset = np.concatenate([np.zeros((1,)), new_offset[:-1]], 0)
+        new_batch_num_objs = new_offset - start_offset
+        ret[key] = F.tensor(new_batch_num_objs, dtype=F.dtype(batch_num_objs))
+    return ret
+
+
+@parametrize_idtype
+def test_set_batch_info(idtype):
+    ctx = F.ctx()
+
+    g1 = dgl.rand_graph(30, 100).astype(idtype).to(F.ctx())
+    g2 = dgl.rand_graph(40, 200).astype(idtype).to(F.ctx())
+    bg = dgl.batch([g1, g2])
+    batch_num_nodes = F.astype(bg.batch_num_nodes(), idtype)
+    batch_num_edges = F.astype(bg.batch_num_edges(), idtype)
+
+    # test homogeneous node subgraph
+    sg_n = dgl.node_subgraph(bg, list(range(10, 20)) + list(range(50, 60)))
+    induced_nodes = sg_n.ndata["_ID"]
+    induced_edges = sg_n.edata["_ID"]
+    new_batch_num_nodes = _get_subgraph_batch_info(
+        bg.ntypes, [induced_nodes], batch_num_nodes
+    )
+    new_batch_num_edges = _get_subgraph_batch_info(
+        bg.canonical_etypes, [induced_edges], batch_num_edges
+    )
+    sg_n.set_batch_num_nodes(new_batch_num_nodes)
+    sg_n.set_batch_num_edges(new_batch_num_edges)
+    subg_n1, subg_n2 = dgl.unbatch(sg_n)
+    subg1 = dgl.node_subgraph(g1, list(range(10, 20)))
+    subg2 = dgl.node_subgraph(g2, list(range(20, 30)))
+    assert subg_n1.num_edges() == subg1.num_edges()
+    assert subg_n2.num_edges() == subg2.num_edges()
+
+    # test homogeneous edge subgraph
+    sg_e = dgl.edge_subgraph(
+        bg, list(range(40, 70)) + list(range(150, 200)), relabel_nodes=False
+    )
+    induced_nodes = F.arange(0, bg.num_nodes(), idtype)
+    induced_edges = sg_e.edata["_ID"]
+    new_batch_num_nodes = _get_subgraph_batch_info(
+        bg.ntypes, [induced_nodes], batch_num_nodes
+    )
+    new_batch_num_edges = _get_subgraph_batch_info(
+        bg.canonical_etypes, [induced_edges], batch_num_edges
+    )
+    sg_e.set_batch_num_nodes(new_batch_num_nodes)
+    sg_e.set_batch_num_edges(new_batch_num_edges)
+    subg_e1, subg_e2 = dgl.unbatch(sg_e)
+    subg1 = dgl.edge_subgraph(g1, list(range(40, 70)), relabel_nodes=False)
+    subg2 = dgl.edge_subgraph(g2, list(range(50, 100)), relabel_nodes=False)
+    assert subg_e1.num_nodes() == subg1.num_nodes()
+    assert subg_e2.num_nodes() == subg2.num_nodes()
+
+
+if __name__ == "__main__":
+    # test_batch_unbatch()
+    # test_batch_unbatch1()
+    # test_batch_unbatch_frame()
+    # test_batch_unbatch2()
+    # test_batched_edge_ordering()
+    # test_batch_send_then_recv()
+    # test_batch_send_and_recv()
+    # test_batch_propagate()
+    # test_batch_no_edge()
+    test_set_batch_info(F.int32)

tests/python/common/test_batch-heterograph.py

@@ -0,0 +1,565 @@
+import unittest
+
+import backend as F
+
+import dgl
+import pytest
+from dgl.base import ALL
+from utils import check_graph_equal, get_cases, parametrize_idtype
+
+
+def check_equivalence_between_heterographs(
+    g1, g2, node_attrs=None, edge_attrs=None
+):
+    assert g1.ntypes == g2.ntypes
+    assert g1.etypes == g2.etypes
+    assert g1.canonical_etypes == g2.canonical_etypes
+
+    for nty in g1.ntypes:
+        assert g1.num_nodes(nty) == g2.num_nodes(nty)
+
+    for ety in g1.etypes:
+        if len(g1._etype2canonical[ety]) > 0:
+            assert g1.num_edges(ety) == g2.num_edges(ety)
+
+    for ety in g1.canonical_etypes:
+        assert g1.num_edges(ety) == g2.num_edges(ety)
+        src1, dst1, eid1 = g1.edges(etype=ety, form="all")
+        src2, dst2, eid2 = g2.edges(etype=ety, form="all")
+        assert F.allclose(src1, src2)
+        assert F.allclose(dst1, dst2)
+        assert F.allclose(eid1, eid2)
+
+    if node_attrs is not None:
+        for nty in node_attrs.keys():
+            if g1.num_nodes(nty) == 0:
+                continue
+            for feat_name in node_attrs[nty]:
+                assert F.allclose(
+                    g1.nodes[nty].data[feat_name], g2.nodes[nty].data[feat_name]
+                )
+
+    if edge_attrs is not None:
+        for ety in edge_attrs.keys():
+            if g1.num_edges(ety) == 0:
+                continue
+            for feat_name in edge_attrs[ety]:
+                assert F.allclose(
+                    g1.edges[ety].data[feat_name], g2.edges[ety].data[feat_name]
+                )
+
+
+@pytest.mark.parametrize("gs", get_cases(["two_hetero_batch"]))
+@parametrize_idtype
+def test_topology(gs, idtype):
+    """Test batching two DGLGraphs where some nodes are isolated in some relations"""
+    g1, g2 = gs
+    g1 = g1.astype(idtype).to(F.ctx())
+    g2 = g2.astype(idtype).to(F.ctx())
+    bg = dgl.batch([g1, g2])
+
+    assert bg.idtype == idtype
+    assert bg.device == F.ctx()
+    assert bg.ntypes == g2.ntypes
+    assert bg.etypes == g2.etypes
+    assert bg.canonical_etypes == g2.canonical_etypes
+    assert bg.batch_size == 2
+
+    # Test number of nodes
+    for ntype in bg.ntypes:
+        print(ntype)
+        assert F.asnumpy(bg.batch_num_nodes(ntype)).tolist() == [
+            g1.num_nodes(ntype),
+            g2.num_nodes(ntype),
+        ]
+        assert bg.num_nodes(ntype) == (
+            g1.num_nodes(ntype) + g2.num_nodes(ntype)
+        )
+
+    # Test number of edges
+    for etype in bg.canonical_etypes:
+        assert F.asnumpy(bg.batch_num_edges(etype)).tolist() == [
+            g1.num_edges(etype),
+            g2.num_edges(etype),
+        ]
+        assert bg.num_edges(etype) == (
+            g1.num_edges(etype) + g2.num_edges(etype)
+        )
+
+    # Test relabeled nodes
+    for ntype in bg.ntypes:
+        assert list(F.asnumpy(bg.nodes(ntype))) == list(
+            range(bg.num_nodes(ntype))
+        )
+
+    # Test relabeled edges
+    src, dst = bg.edges(etype=("user", "follows", "user"))
+    assert list(F.asnumpy(src)) == [0, 1, 4, 5]
+    assert list(F.asnumpy(dst)) == [1, 2, 5, 6]
+    src, dst = bg.edges(etype=("user", "follows", "developer"))
+    assert list(F.asnumpy(src)) == [0, 1, 4, 5]
+    assert list(F.asnumpy(dst)) == [1, 2, 4, 5]
+    src, dst, eid = bg.edges(etype="plays", form="all")
+    assert list(F.asnumpy(src)) == [0, 1, 2, 3, 4, 5, 6]
+    assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2, 2, 3]
+    assert list(F.asnumpy(eid)) == [0, 1, 2, 3, 4, 5, 6]
+
+    # Test unbatching graphs
+    g3, g4 = dgl.unbatch(bg)
+    check_equivalence_between_heterographs(g1, g3)
+    check_equivalence_between_heterographs(g2, g4)
+
+    # Test dtype cast
+    if idtype == "int32":
+        bg_cast = bg.long()
+    else:
+        bg_cast = bg.int()
+    assert bg.batch_size == bg_cast.batch_size
+
+    # Test local var
+    bg_local = bg.local_var()
+    assert bg.batch_size == bg_local.batch_size
+
+
+@parametrize_idtype
+def test_batching_batched(idtype):
+    """Test batching a DGLGraph and a batched DGLGraph."""
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    bg1 = dgl.batch([g1, g2])
+    g3 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0], [1]),
+            ("user", "plays", "game"): ([1], [0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    bg2 = dgl.batch([bg1, g3])
+    assert bg2.idtype == idtype
+    assert bg2.device == F.ctx()
+    assert bg2.ntypes == g3.ntypes
+    assert bg2.etypes == g3.etypes
+    assert bg2.canonical_etypes == g3.canonical_etypes
+    assert bg2.batch_size == 3
+
+    # Test number of nodes
+    for ntype in bg2.ntypes:
+        assert F.asnumpy(bg2.batch_num_nodes(ntype)).tolist() == [
+            g1.num_nodes(ntype),
+            g2.num_nodes(ntype),
+            g3.num_nodes(ntype),
+        ]
+        assert bg2.num_nodes(ntype) == (
+            g1.num_nodes(ntype) + g2.num_nodes(ntype) + g3.num_nodes(ntype)
+        )
+
+    # Test number of edges
+    for etype in bg2.canonical_etypes:
+        assert F.asnumpy(bg2.batch_num_edges(etype)).tolist() == [
+            g1.num_edges(etype),
+            g2.num_edges(etype),
+            g3.num_edges(etype),
+        ]
+        assert bg2.num_edges(etype) == (
+            g1.num_edges(etype) + g2.num_edges(etype) + g3.num_edges(etype)
+        )
+
+    # Test relabeled nodes
+    for ntype in bg2.ntypes:
+        assert list(F.asnumpy(bg2.nodes(ntype))) == list(
+            range(bg2.num_nodes(ntype))
+        )
+
+    # Test relabeled edges
+    src, dst = bg2.edges(etype="follows")
+    assert list(F.asnumpy(src)) == [0, 1, 3, 4, 6]
+    assert list(F.asnumpy(dst)) == [1, 2, 4, 5, 7]
+    src, dst = bg2.edges(etype="plays")
+    assert list(F.asnumpy(src)) == [0, 1, 3, 4, 7]
+    assert list(F.asnumpy(dst)) == [0, 0, 1, 1, 2]
+
+    # Test unbatching graphs
+    g4, g5, g6 = dgl.unbatch(bg2)
+    check_equivalence_between_heterographs(g1, g4)
+    check_equivalence_between_heterographs(g2, g5)
+    check_equivalence_between_heterographs(g3, g6)
+
+
+@parametrize_idtype
+def test_features(idtype):
+    """Test the features of batched DGLGraphs"""
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g1.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
+    g1.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
+    g1.nodes["game"].data["h1"] = F.tensor([[0.0]])
+    g1.nodes["game"].data["h2"] = F.tensor([[1.0]])
+    g1.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
+    g1.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
+    g1.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])
+
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g2.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
+    g2.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
+    g2.nodes["game"].data["h1"] = F.tensor([[0.0]])
+    g2.nodes["game"].data["h2"] = F.tensor([[1.0]])
+    g2.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
+    g2.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
+    g2.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])
+
+    # test default setting
+    bg = dgl.batch([g1, g2])
+    assert F.allclose(
+        bg.nodes["user"].data["h1"],
+        F.cat(
+            [g1.nodes["user"].data["h1"], g2.nodes["user"].data["h1"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.nodes["user"].data["h2"],
+        F.cat(
+            [g1.nodes["user"].data["h2"], g2.nodes["user"].data["h2"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.nodes["game"].data["h1"],
+        F.cat(
+            [g1.nodes["game"].data["h1"], g2.nodes["game"].data["h1"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.nodes["game"].data["h2"],
+        F.cat(
+            [g1.nodes["game"].data["h2"], g2.nodes["game"].data["h2"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.edges["follows"].data["h1"],
+        F.cat(
+            [g1.edges["follows"].data["h1"], g2.edges["follows"].data["h1"]],
+            dim=0,
+        ),
+    )
+    assert F.allclose(
+        bg.edges["follows"].data["h2"],
+        F.cat(
+            [g1.edges["follows"].data["h2"], g2.edges["follows"].data["h2"]],
+            dim=0,
+        ),
+    )
+    assert F.allclose(
+        bg.edges["plays"].data["h1"],
+        F.cat(
+            [g1.edges["plays"].data["h1"], g2.edges["plays"].data["h1"]], dim=0
+        ),
+    )
+
+    # test specifying ndata/edata
+    bg = dgl.batch([g1, g2], ndata=["h2"], edata=["h1"])
+    assert F.allclose(
+        bg.nodes["user"].data["h2"],
+        F.cat(
+            [g1.nodes["user"].data["h2"], g2.nodes["user"].data["h2"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.nodes["game"].data["h2"],
+        F.cat(
+            [g1.nodes["game"].data["h2"], g2.nodes["game"].data["h2"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.edges["follows"].data["h1"],
+        F.cat(
+            [g1.edges["follows"].data["h1"], g2.edges["follows"].data["h1"]],
+            dim=0,
+        ),
+    )
+    assert F.allclose(
+        bg.edges["plays"].data["h1"],
+        F.cat(
+            [g1.edges["plays"].data["h1"], g2.edges["plays"].data["h1"]], dim=0
+        ),
+    )
+    assert "h1" not in bg.nodes["user"].data
+    assert "h1" not in bg.nodes["game"].data
+    assert "h2" not in bg.edges["follows"].data
+
+    # Test unbatching graphs
+    g3, g4 = dgl.unbatch(bg)
+    check_equivalence_between_heterographs(
+        g1,
+        g3,
+        node_attrs={"user": ["h2"], "game": ["h2"]},
+        edge_attrs={("user", "follows", "user"): ["h1"]},
+    )
+    check_equivalence_between_heterographs(
+        g2,
+        g4,
+        node_attrs={"user": ["h2"], "game": ["h2"]},
+        edge_attrs={("user", "follows", "user"): ["h1"]},
+    )
+
+
+@unittest.skipIf(
+    F.backend_name == "mxnet",
+    reason="MXNet does not support split array with zero-length segment.",
+)
+@parametrize_idtype
+def test_empty_relation(idtype):
+    """Test the features of batched DGLGraphs"""
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([], []),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g1.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
+    g1.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
+    g1.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
+    g1.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
+
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g2.nodes["user"].data["h1"] = F.tensor([[0.0], [1.0], [2.0]])
+    g2.nodes["user"].data["h2"] = F.tensor([[3.0], [4.0], [5.0]])
+    g2.nodes["game"].data["h1"] = F.tensor([[0.0]])
+    g2.nodes["game"].data["h2"] = F.tensor([[1.0]])
+    g2.edges["follows"].data["h1"] = F.tensor([[0.0], [1.0]])
+    g2.edges["follows"].data["h2"] = F.tensor([[2.0], [3.0]])
+    g2.edges["plays"].data["h1"] = F.tensor([[0.0], [1.0]])
+
+    bg = dgl.batch([g1, g2])
+
+    # Test number of nodes
+    for ntype in bg.ntypes:
+        assert F.asnumpy(bg.batch_num_nodes(ntype)).tolist() == [
+            g1.num_nodes(ntype),
+            g2.num_nodes(ntype),
+        ]
+
+    # Test number of edges
+    for etype in bg.canonical_etypes:
+        assert F.asnumpy(bg.batch_num_edges(etype)).tolist() == [
+            g1.num_edges(etype),
+            g2.num_edges(etype),
+        ]
+
+    # Test features
+    assert F.allclose(
+        bg.nodes["user"].data["h1"],
+        F.cat(
+            [g1.nodes["user"].data["h1"], g2.nodes["user"].data["h1"]], dim=0
+        ),
+    )
+    assert F.allclose(
+        bg.nodes["user"].data["h2"],
+        F.cat(
+            [g1.nodes["user"].data["h2"], g2.nodes["user"].data["h2"]], dim=0
+        ),
+    )
+    assert F.allclose(bg.nodes["game"].data["h1"], g2.nodes["game"].data["h1"])
+    assert F.allclose(bg.nodes["game"].data["h2"], g2.nodes["game"].data["h2"])
+    assert F.allclose(
+        bg.edges["follows"].data["h1"],
+        F.cat(
+            [g1.edges["follows"].data["h1"], g2.edges["follows"].data["h1"]],
+            dim=0,
+        ),
+    )
+    assert F.allclose(
+        bg.edges["plays"].data["h1"], g2.edges["plays"].data["h1"]
+    )
+
+    # Test unbatching graphs
+    g3, g4 = dgl.unbatch(bg)
+    check_equivalence_between_heterographs(
+        g1,
+        g3,
+        node_attrs={"user": ["h1", "h2"], "game": ["h1", "h2"]},
+        edge_attrs={("user", "follows", "user"): ["h1"]},
+    )
+    check_equivalence_between_heterographs(
+        g2,
+        g4,
+        node_attrs={"user": ["h1", "h2"], "game": ["h1", "h2"]},
+        edge_attrs={("user", "follows", "user"): ["h1"]},
+    )
+
+    # Test graphs without edges
+    g1 = dgl.heterograph({("u", "r", "v"): ([], [])}, {"u": 0, "v": 4})
+    g2 = dgl.heterograph({("u", "r", "v"): ([], [])}, {"u": 1, "v": 5})
+    dgl.batch([g1, g2])
+
+
+@parametrize_idtype
+def test_unbatch2(idtype):
+    # batch 3 graphs but unbatch to 2
+    g1 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
+    g2 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
+    g3 = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
+    bg = dgl.batch([g1, g2, g3])
+    bnn = F.tensor([8, 4])
+    bne = F.tensor([6, 3])
+    f1, f2 = dgl.unbatch(bg, node_split=bnn, edge_split=bne)
+    u, v = f1.edges(order="eid")
+    assert F.allclose(u, F.tensor([0, 1, 2, 4, 5, 6]))
+    assert F.allclose(v, F.tensor([1, 2, 3, 5, 6, 7]))
+    u, v = f2.edges(order="eid")
+    assert F.allclose(u, F.tensor([0, 1, 2]))
+    assert F.allclose(v, F.tensor([1, 2, 3]))
+
+    # batch 2 but unbatch to 3
+    bg = dgl.batch([f1, f2])
+    gg1, gg2, gg3 = dgl.unbatch(bg, F.tensor([4, 4, 4]), F.tensor([3, 3, 3]))
+    check_graph_equal(g1, gg1)
+    check_graph_equal(g2, gg2)
+    check_graph_equal(g3, gg3)
+
+
+@parametrize_idtype
+def test_slice_batch(idtype):
+    g1 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([], []),
+            ("user", "follows", "game"): ([0, 0], [1, 4]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1], [0, 0]),
+            ("user", "follows", "game"): ([0, 1], [1, 4]),
+        },
+        num_nodes_dict={"user": 4, "game": 6},
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g3 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0], [2]),
+            ("user", "plays", "game"): ([1, 2], [3, 4]),
+            ("user", "follows", "game"): ([], []),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g_list = [g1, g2, g3]
+    bg = dgl.batch(g_list)
+    bg.nodes["user"].data["h1"] = F.randn((bg.num_nodes("user"), 2))
+    bg.nodes["user"].data["h2"] = F.randn((bg.num_nodes("user"), 5))
+    bg.edges[("user", "follows", "user")].data["h1"] = F.randn(
+        (bg.num_edges(("user", "follows", "user")), 2)
+    )
+    for fmat in ["coo", "csr", "csc"]:
+        bg = bg.formats(fmat)
+        for i in range(len(g_list)):
+            g_i = g_list[i]
+            g_slice = dgl.slice_batch(bg, i)
+            assert g_i.ntypes == g_slice.ntypes
+            assert g_i.canonical_etypes == g_slice.canonical_etypes
+            assert g_i.idtype == g_slice.idtype
+            assert g_i.device == g_slice.device
+            for nty in g_i.ntypes:
+                assert g_i.num_nodes(nty) == g_slice.num_nodes(nty)
+                for feat in g_i.nodes[nty].data:
+                    assert F.allclose(
+                        g_i.nodes[nty].data[feat], g_slice.nodes[nty].data[feat]
+                    )
+
+            for ety in g_i.canonical_etypes:
+                assert g_i.num_edges(ety) == g_slice.num_edges(ety)
+                for feat in g_i.edges[ety].data:
+                    assert F.allclose(
+                        g_i.edges[ety].data[feat], g_slice.edges[ety].data[feat]
+                    )
+
+
+@parametrize_idtype
+def test_batch_keeps_empty_data(idtype):
+    g1 = (
+        dgl.heterograph({("a", "to", "a"): ([], [])}).astype(idtype).to(F.ctx())
+    )
+    g1.nodes["a"].data["nh"] = F.tensor([])
+    g1.edges[("a", "to", "a")].data["eh"] = F.tensor([])
+    g2 = (
+        dgl.heterograph({("a", "to", "a"): ([], [])}).astype(idtype).to(F.ctx())
+    )
+    g2.nodes["a"].data["nh"] = F.tensor([])
+    g2.edges[("a", "to", "a")].data["eh"] = F.tensor([])
+    g = dgl.batch([g1, g2])
+    assert "nh" in g.nodes["a"].data
+    assert "eh" in g.edges[("a", "to", "a")].data
+
+
+def test_batch_netypes():
+    # Test for https://github.com/dmlc/dgl/issues/2808
+    import networkx as nx
+
+    B = nx.DiGraph()
+    B.add_nodes_from(
+        [1, 2, 3, 4],
+        bipartite=0,
+        some_attr=F.tensor([1, 2, 3, 4], dtype=F.float32),
+    )
+    B.add_nodes_from(["a", "b", "c"], bipartite=1)
+    B.add_edges_from(
+        [(1, "a"), (1, "b"), (2, "b"), (2, "c"), (3, "c"), (4, "a")]
+    )
+
+    g_dict = {
+        0: dgl.bipartite_from_networkx(B, "A", "e", "B"),
+        1: dgl.bipartite_from_networkx(B, "B", "e", "A"),
+        2: dgl.bipartite_from_networkx(B, "A", "e", "B", u_attrs=["some_attr"]),
+        3: dgl.bipartite_from_networkx(B, "B", "e", "A", u_attrs=["some_attr"]),
+    }
+    for _, g in g_dict.items():
+        dgl.batch((g, g, g))
+
+
+if __name__ == "__main__":
+    # test_topology('int32')
+    # test_batching_batched('int32')
+    # test_batched_features('int32')
+    # test_empty_relation('int64')
+    # test_to_device('int32')
+    pass

tests/python/common/test_convert.py

@@ -0,0 +1,122 @@
+import unittest
+
+import backend as F
+import dgl
+
+from utils import parametrize_idtype
+
+
+def get_nodes_by_ntype(nodes, ntype):
+    return dict((k, v) for k, v in nodes.items() if v["ntype"] == ntype)
+
+
+def edge_attrs(edge):
+    # Edges in Networkx are in the format (src, dst, attrs)
+    return edge[2]
+
+
+def get_edges_by_etype(edges, etype):
+    return [e for e in edges if edge_attrs(e)["etype"] == etype]
+
+
+def check_attrs_for_nodes(nodes, attrs):
+    return all(v.keys() == attrs for v in nodes.values())
+
+
+def check_attr_values_for_nodes(nodes, attr_name, values):
+    return F.allclose(
+        F.stack([v[attr_name] for v in nodes.values()], 0), values
+    )
+
+
+def check_attrs_for_edges(edges, attrs):
+    return all(edge_attrs(e).keys() == attrs for e in edges)
+
+
+def check_attr_values_for_edges(edges, attr_name, values):
+    return F.allclose(
+        F.stack([edge_attrs(e)[attr_name] for e in edges], 0), values
+    )
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="`to_networkx` does not support graphs on GPU",
+)
+@parametrize_idtype
+def test_to_networkx(idtype):
+    # TODO: adapt and move code from the _test_nx_conversion function in
+    # tests/python/common/function/test_basics.py to here
+    # (pending resolution of https://github.com/dmlc/dgl/issues/5735).
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "follows", "topic"): ([1, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 3], [3, 4]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+
+    n1 = F.randn((5, 3))
+    n2 = F.randn((4, 2))
+    e1 = F.randn((2, 3))
+    e2 = F.randn((2, 2))
+
+    g.nodes["game"].data["n"] = F.copy_to(n1, ctx=F.ctx())
+    g.nodes["user"].data["n"] = F.copy_to(n2, ctx=F.ctx())
+    g.edges[("user", "follows", "user")].data["e"] = F.copy_to(e1, ctx=F.ctx())
+    g.edges["plays"].data["e"] = F.copy_to(e2, ctx=F.ctx())
+
+    nxg = dgl.to_networkx(
+        g,
+        node_attrs=["n"],
+        edge_attrs=["e"],
+    )
+
+    # Test nodes
+    nxg_nodes = dict(nxg.nodes(data=True))
+    assert len(nxg_nodes) == g.num_nodes()
+    assert {v["ntype"] for v in nxg_nodes.values()} == set(g.ntypes)
+
+    nxg_nodes_by_ntype = {}
+    for ntype in g.ntypes:
+        nxg_nodes_by_ntype[ntype] = get_nodes_by_ntype(nxg_nodes, ntype)
+        assert g.num_nodes(ntype) == len(nxg_nodes_by_ntype[ntype])
+
+    assert check_attrs_for_nodes(nxg_nodes_by_ntype["game"], {"ntype", "n"})
+    assert check_attr_values_for_nodes(nxg_nodes_by_ntype["game"], "n", n1)
+    assert check_attrs_for_nodes(nxg_nodes_by_ntype["user"], {"ntype", "n"})
+    assert check_attr_values_for_nodes(nxg_nodes_by_ntype["user"], "n", n2)
+    # Nodes without node attributes
+    assert check_attrs_for_nodes(nxg_nodes_by_ntype["topic"], {"ntype"})
+
+    # Test edges
+    nxg_edges = list(nxg.edges(data=True))
+    assert len(nxg_edges) == g.num_edges()
+    assert {edge_attrs(e)["etype"] for e in nxg_edges} == set(
+        g.canonical_etypes
+    )
+
+    nxg_edges_by_etype = {}
+    for etype in g.canonical_etypes:
+        nxg_edges_by_etype[etype] = get_edges_by_etype(nxg_edges, etype)
+        assert g.num_edges(etype) == len(nxg_edges_by_etype[etype])
+
+    assert check_attrs_for_edges(
+        nxg_edges_by_etype[("user", "follows", "user")],
+        {"id", "etype", "e"},
+    )
+    assert check_attr_values_for_edges(
+        nxg_edges_by_etype[("user", "follows", "user")], "e", e1
+    )
+    assert check_attrs_for_edges(
+        nxg_edges_by_etype[("user", "plays", "game")], {"id", "etype", "e"}
+    )
+    assert check_attr_values_for_edges(
+        nxg_edges_by_etype[("user", "plays", "game")], "e", e2
+    )
+    # Edges without edge attributes
+    assert check_attrs_for_edges(
+        nxg_edges_by_etype[("user", "follows", "topic")], {"id", "etype"}
+    )

tests/python/common/test_ffi.py

@@ -0,0 +1,55 @@
+import os
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+import pytest
+
+
+@unittest.skipIf(os.name == "nt", reason="Cython only works on linux")
+def test_cython():
+    import dgl._ffi._cy3.core
+
+
+@pytest.mark.parametrize("arg", [1, 2.3])
+def test_callback(arg):
+    def cb(x):
+        return x + 1
+
+    ret = dgl._api_internal._TestPythonCallback(cb, arg)
+    assert ret == arg + 1
+
+
+@pytest.mark.parametrize("dtype", [F.float32, F.float64, F.int32, F.int64])
+def _test_callback_array(dtype):
+    def cb(x):
+        return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
+
+    arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
+    ret = F.from_dgl_nd(
+        dgl._api_internal._TestPythonCallback(cb, F.to_dgl_nd(arg))
+    )
+    assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)
+
+
+@pytest.mark.parametrize("arg", [1, 2.3])
+def test_callback_thread(arg):
+    def cb(x):
+        return x + 1
+
+    ret = dgl._api_internal._TestPythonCallbackThread(cb, arg)
+    assert ret == arg + 1
+
+
+@pytest.mark.parametrize("dtype", [F.float32, F.float64, F.int32, F.int64])
+def _test_callback_array_thread(dtype):
+    def cb(x):
+        return F.to_dgl_nd(F.from_dgl_nd(x) + 1)
+
+    arg = F.copy_to(F.tensor([1, 2, 3], dtype=dtype), F.ctx())
+    ret = F.from_dgl_nd(
+        dgl._api_internal._TestPythonCallbackThread(cb, F.to_dgl_nd(arg))
+    )
+    assert np.allclose(F.asnumpy(ret), F.asnumpy(arg) + 1)

tests/python/common/test_frame.py

@@ -0,0 +1,120 @@
+import pickle
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.ndarray as nd
+import numpy as np
+from dgl.frame import Column
+from utils import parametrize_idtype
+
+
+def test_column_subcolumn():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    # subcolumn from cpu context
+    i1 = F.tensor([0, 2, 1, 3], dtype=F.int64)
+    l1 = original.subcolumn(i1)
+
+    assert len(l1) == i1.shape[0]
+    assert F.array_equal(l1.data, F.gather_row(data, i1))
+
+    # next subcolumn from target context
+    i2 = F.copy_to(F.tensor([0, 2], dtype=F.int64), F.ctx())
+    l2 = l1.subcolumn(i2)
+
+    assert len(l2) == i2.shape[0]
+    i1i2 = F.copy_to(F.gather_row(i1, F.copy_to(i2, F.context(i1))), F.ctx())
+    assert F.array_equal(l2.data, F.gather_row(data, i1i2))
+
+    # next subcolumn also from target context
+    i3 = F.copy_to(F.tensor([1], dtype=F.int64), F.ctx())
+    l3 = l2.subcolumn(i3)
+
+    assert len(l3) == i3.shape[0]
+    i1i2i3 = F.copy_to(
+        F.gather_row(i1i2, F.copy_to(i3, F.context(i1i2))), F.ctx()
+    )
+    assert F.array_equal(l3.data, F.gather_row(data, i1i2i3))
+
+
+def test_serialize_deserialize_plain():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    serial = pickle.dumps(original)
+    new = pickle.loads(serial)
+    print("new = {}".format(new))
+
+    assert F.array_equal(new.data, original.data)
+
+
+def test_serialize_deserialize_subcolumn():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    # subcolumn from cpu context
+    i1 = F.tensor([0, 2, 1, 3], dtype=F.int64)
+    l1 = original.subcolumn(i1)
+
+    serial = pickle.dumps(l1)
+    new = pickle.loads(serial)
+
+    assert F.array_equal(new.data, l1.data)
+
+
+def test_serialize_deserialize_dtype():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+    original = original.astype(F.int64)
+
+    serial = pickle.dumps(original)
+    new = pickle.loads(serial)
+
+    assert new.dtype == F.int64

tests/python/common/test_generators.py

@@ -0,0 +1,28 @@
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu", reason="GPU random choice not implemented"
+)
+def test_rand_graph():
+    g = dgl.rand_graph(10000, 100000)
+    assert g.num_nodes() == 10000
+    assert g.num_edges() == 100000
+    # test random seed
+    dgl.random.seed(42)
+    g1 = dgl.rand_graph(100, 30)
+    dgl.random.seed(42)
+    g2 = dgl.rand_graph(100, 30)
+    u1, v1 = g1.edges()
+    u2, v2 = g2.edges()
+    assert F.array_equal(u1, u2)
+    assert F.array_equal(v1, v2)
+
+
+if __name__ == "__main__":
+    test_rand_graph()

tests/python/common/test_heterograph-apply-edges.py

@@ -0,0 +1,301 @@
+import itertools
+import unittest
+from collections import Counter
+from itertools import product
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as spsp
+import torch
+
+from dgl import DGLError
+from scipy.sparse import rand
+from utils import get_cases, parametrize_idtype
+
+rfuncs = {"sum": fn.sum, "max": fn.max, "min": fn.min, "mean": fn.mean}
+fill_value = {"sum": 0, "max": float("-inf")}
+feat_size = 2
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+def create_test_heterograph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 1, 1], [0, 0, 1]),
+            ("developer", "develops", "game"): ([0, 1, 0], [0, 1, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+def create_random_hetero_with_single_source_node_type(idtype):
+    num_nodes = {"n1": 5, "n2": 10, "n3": 15}
+    etypes = [("n1", "r1", "n2"), ("n1", "r2", "n3"), ("n1", "r3", "n2")]
+    edges = {}
+    for etype in etypes:
+        src_ntype, _, dst_ntype = etype
+        arr = spsp.random(
+            num_nodes[src_ntype],
+            num_nodes[dst_ntype],
+            density=1,
+            format="coo",
+            random_state=100,
+        )
+        edges[etype] = (arr.row, arr.col)
+    return dgl.heterograph(edges, idtype=idtype, device=F.ctx())
+
+
+@parametrize_idtype
+def test_unary_copy_u(idtype):
+    def _test(mfunc):
+        g = create_test_heterograph(idtype)
+
+        x1 = F.randn((g.num_nodes("user"), feat_size))
+        x2 = F.randn((g.num_nodes("developer"), feat_size))
+
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        g.nodes["user"].data["h"] = x1
+        g.nodes["developer"].data["h"] = x2
+
+        #################################################################
+        #  apply_edges() is called on each relation type separately
+        #################################################################
+
+        with F.record_grad():
+            [
+                g.apply_edges(fn.copy_u("h", "m"), etype=rel)
+                for rel in g.canonical_etypes
+            ]
+            r1 = g["plays"].edata["m"]
+            F.backward(r1, F.ones(r1.shape))
+            n_grad1 = F.grad(g.ndata["h"]["user"])
+        # TODO (Israt): clear not working
+        g.edata["m"].clear()
+
+        #################################################################
+        #  apply_edges() is called on all relation types
+        #################################################################
+
+        g.apply_edges(fn.copy_u("h", "m"))
+        r2 = g["plays"].edata["m"]
+        F.backward(r2, F.ones(r2.shape))
+        n_grad2 = F.grad(g.nodes["user"].data["h"])
+
+        # correctness check
+        def _print_error(a, b):
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        if not F.allclose(n_grad1, n_grad2):
+            print("node grad")
+            _print_error(n_grad1, n_grad2)
+        assert F.allclose(n_grad1, n_grad2)
+
+    _test(fn.copy_u)
+
+
+@parametrize_idtype
+def test_unary_copy_e(idtype):
+    def _test(mfunc):
+        g = create_test_heterograph(idtype)
+        feat_size = 2
+
+        x1 = F.randn((4, feat_size))
+        x2 = F.randn((4, feat_size))
+        x3 = F.randn((3, feat_size))
+        x4 = F.randn((3, feat_size))
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+        g["plays"].edata["eid"] = x1
+        g["follows"].edata["eid"] = x2
+        g["develops"].edata["eid"] = x3
+        g["wishes"].edata["eid"] = x4
+
+        #################################################################
+        #  apply_edges() is called on each relation type separately
+        #################################################################
+        with F.record_grad():
+            [
+                g.apply_edges(fn.copy_e("eid", "m"), etype=rel)
+                for rel in g.canonical_etypes
+            ]
+            r1 = g["develops"].edata["m"]
+            F.backward(r1, F.ones(r1.shape))
+            e_grad1 = F.grad(g["develops"].edata["eid"])
+
+        #################################################################
+        #  apply_edges() is called on all relation types
+        #################################################################
+
+        g.apply_edges(fn.copy_e("eid", "m"))
+        r2 = g["develops"].edata["m"]
+        F.backward(r2, F.ones(r2.shape))
+        e_grad2 = F.grad(g["develops"].edata["eid"])
+
+        # # correctness check
+        def _print_error(a, b):
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        if not F.allclose(e_grad1, e_grad2):
+            print("edge grad")
+            _print_error(e_grad1, e_grad2)
+        assert F.allclose(e_grad1, e_grad2)
+
+    _test(fn.copy_e)
+
+
+@parametrize_idtype
+def test_binary_op(idtype):
+    def _test(lhs, rhs, binary_op):
+        g = create_test_heterograph(idtype)
+
+        n1 = F.randn((g.num_nodes("user"), feat_size))
+        n2 = F.randn((g.num_nodes("developer"), feat_size))
+        n3 = F.randn((g.num_nodes("game"), feat_size))
+
+        x1 = F.randn((g.num_edges("plays"), feat_size))
+        x2 = F.randn((g.num_edges("follows"), feat_size))
+        x3 = F.randn((g.num_edges("develops"), feat_size))
+        x4 = F.randn((g.num_edges("wishes"), feat_size))
+
+        builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
+        builtin_msg = getattr(fn, builtin_msg_name)
+
+        #################################################################
+        #  apply_edges() is called on each relation type separately
+        #################################################################
+
+        F.attach_grad(n1)
+        F.attach_grad(n2)
+        F.attach_grad(n3)
+        g.nodes["user"].data["h"] = n1
+        g.nodes["developer"].data["h"] = n2
+        g.nodes["game"].data["h"] = n3
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+        g["plays"].edata["h"] = x1
+        g["follows"].edata["h"] = x2
+        g["develops"].edata["h"] = x3
+        g["wishes"].edata["h"] = x4
+
+        with F.record_grad():
+            [
+                g.apply_edges(builtin_msg("h", "h", "m"), etype=rel)
+                for rel in g.canonical_etypes
+            ]
+            r1 = g["plays"].edata["m"]
+            loss = F.sum(r1.view(-1), 0)
+            F.backward(loss)
+            n_grad1 = F.grad(g.nodes["game"].data["h"])
+
+        #################################################################
+        #  apply_edges() is called on all relation types
+        #################################################################
+
+        F.attach_grad(n1)
+        F.attach_grad(n2)
+        F.attach_grad(n3)
+        g.nodes["user"].data["h"] = n1
+        g.nodes["developer"].data["h"] = n2
+        g.nodes["game"].data["h"] = n3
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+        g["plays"].edata["h"] = x1
+        g["follows"].edata["h"] = x2
+        g["develops"].edata["h"] = x3
+        g["wishes"].edata["h"] = x4
+
+        with F.record_grad():
+            g.apply_edges(builtin_msg("h", "h", "m"))
+            r2 = g["plays"].edata["m"]
+            loss = F.sum(r2.view(-1), 0)
+            F.backward(loss)
+            n_grad2 = F.grad(g.nodes["game"].data["h"])
+
+        # correctness check
+        def _print_error(a, b):
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        if n_grad1 is not None or n_grad2 is not None:
+            if not F.allclose(n_grad1, n_grad2):
+                print("node grad")
+                _print_error(n_grad1, n_grad2)
+            assert F.allclose(n_grad1, n_grad2)
+
+    target = ["u", "v", "e"]
+    for lhs, rhs in product(target, target):
+        if lhs == rhs:
+            continue
+        for binary_op in ["add", "sub", "mul", "div", "dot"]:
+            print(lhs, rhs, binary_op)
+            _test(lhs, rhs, binary_op)
+
+
+# Here we test heterograph with only single source node type because the format
+# of node feature is a tensor.
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_heterograph_with_single_source_node_type_apply_edges(idtype):
+    hg = create_random_hetero_with_single_source_node_type(idtype)
+
+    hg.nodes["n1"].data["h"] = F.randn((hg.num_nodes("n1"), 1))
+    hg.nodes["n2"].data["h"] = F.randn((hg.num_nodes("n2"), 1))
+    hg.nodes["n3"].data["h"] = F.randn((hg.num_nodes("n3"), 1))
+
+    assert type(hg.srcdata["h"]) == torch.Tensor
+    hg.apply_edges(fn.u_add_v("h", "h", "x"))
+
+
+if __name__ == "__main__":
+    test_unary_copy_u()
+    test_unary_copy_e()

tests/python/common/test_heterograph-index.py

@@ -0,0 +1,84 @@
+import unittest
+
+import backend as F
+
+import dgl
+import pytest
+from dgl import DGLError
+from utils import parametrize_idtype
+
+
+def create_test_heterograph(idtype):
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="Need gpu for this test"
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Pinning graph outplace only supported for PyTorch",
+)
+@parametrize_idtype
+def test_pin_memory(idtype):
+    g = create_test_heterograph(idtype)
+    g.nodes["user"].data["h"] = F.ones((3, 5))
+    g.nodes["game"].data["i"] = F.ones((2, 5))
+    g.edges["plays"].data["e"] = F.ones((4, 4))
+    g = g.to(F.cpu())
+    assert not g.is_pinned()
+
+    # Test pinning a CPU graph.
+    g._graph.pin_memory()
+    assert not g.is_pinned()
+    g._graph = g._graph.pin_memory()
+    assert g.is_pinned()
+    assert g.device == F.cpu()
+
+    # when clone with a new (different) formats, e.g., g.formats("csc")
+    # ensure the new graphs are not pinned
+    assert not g.formats("csc").is_pinned()
+    assert not g.formats("csr").is_pinned()
+    # 'coo' formats is the default and thus not cloned
+    assert g.formats("coo").is_pinned()
+
+    # Test pinning a GPU graph will cause error raised.
+    g1 = g.to(F.cuda())
+    with pytest.raises(DGLError):
+        g1._graph.pin_memory()
+
+    # Test pinning an empty homograph
+    g2 = dgl.graph(([], []))
+    assert not g2.is_pinned()
+    g2._graph = g2._graph.pin_memory()
+    assert g2.is_pinned()
+
+    # Test pinning heterograph with 0 edge of one relation type
+    g3 = dgl.heterograph(
+        {("a", "b", "c"): ([0, 1], [1, 2]), ("c", "d", "c"): ([], [])}
+    ).astype(idtype)
+    g3._graph = g3._graph.pin_memory()
+    assert g3.is_pinned()
+
+
+if __name__ == "__main__":
+    pass

tests/python/common/test_heterograph-kernel.py

@@ -0,0 +1,414 @@
+from itertools import product
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import numpy as np
+import pytest
+from utils import get_cases, parametrize_idtype
+
+
+def udf_copy_src(edges):
+    return {"m": edges.src["u"]}
+
+
+def udf_copy_edge(edges):
+    return {"m": edges.data["e"]}
+
+
+def udf_mean(nodes):
+    return {"r2": F.mean(nodes.mailbox["m"], 1)}
+
+
+def udf_sum(nodes):
+    return {"r2": F.sum(nodes.mailbox["m"], 1)}
+
+
+def udf_max(nodes):
+    return {"r2": F.max(nodes.mailbox["m"], 1)}
+
+
+D1 = 5
+D2 = 3
+D3 = 4
+D4 = 10  # NOTE(xiang): used to dot feature vector
+builtin = {"sum": fn.sum, "max": fn.max, "mean": fn.mean}
+udf_reduce = {"sum": udf_sum, "max": udf_max, "mean": udf_mean}
+fill_value = {"sum": 0, "max": float("-inf")}
+
+
+def generate_feature(g, broadcast="none", binary_op="none"):
+    """Create graph with src, edge, dst feature. broadcast can be 'u',
+    'e', 'v', 'none'
+    """
+    np.random.seed(31)
+    nv = g.num_nodes()
+    ne = g.num_edges()
+    if binary_op == "dot":
+        if broadcast == "e":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D2, 1, D4)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+        elif broadcast == "u":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1, D4)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+        elif broadcast == "v":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1, D4)))
+        else:
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
+    else:
+        if broadcast == "e":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D2, 1)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+        elif broadcast == "u":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+        elif broadcast == "v":
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
+        else:
+            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
+            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
+    return (
+        F.astype(u, F.float32),
+        F.astype(v, F.float32),
+        F.astype(e, F.float32),
+    )
+
+
+def test_copy_src_reduce():
+    def _test(red, partial):
+        g = dgl.from_networkx(nx.erdos_renyi_graph(100, 0.1))
+        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+        # https://github.com/dmlc/dgl/issues/761
+        g.add_edges(g.nodes(), g.nodes())
+        g = g.to(F.ctx())
+        hu, hv, he = generate_feature(g, "none", "none")
+        if partial:
+            nid = F.tensor(list(range(0, 100, 2)), g.idtype)
+
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        with F.record_grad():
+            if partial:
+                g.pull(
+                    nid,
+                    fn.copy_u(u="u", out="m"),
+                    builtin[red](msg="m", out="r1"),
+                )
+            else:
+                g.update_all(
+                    fn.copy_u(u="u", out="m"), builtin[red](msg="m", out="r1")
+                )
+            r1 = g.ndata["r1"]
+            F.backward(F.reduce_sum(r1))
+            n_grad1 = F.grad(g.ndata["u"])
+
+        # reset grad
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        with F.record_grad():
+            if partial:
+                g.pull(nid, udf_copy_src, udf_reduce[red])
+            else:
+                g.update_all(udf_copy_src, udf_reduce[red])
+            r2 = g.ndata["r2"]
+            F.backward(F.reduce_sum(r2))
+            n_grad2 = F.grad(g.ndata["u"])
+
+        def _print_error(a, b):
+            print("ERROR: Test copy_src_{} partial: {}".format(red, partial))
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        if not F.allclose(n_grad1, n_grad2):
+            print("node grad")
+            _print_error(n_grad1, n_grad2)
+        assert F.allclose(n_grad1, n_grad2)
+
+    _test("sum", False)
+    _test("max", False)
+    _test("mean", False)
+    _test("sum", True)
+    _test("max", True)
+    _test("mean", True)
+
+
+def test_copy_edge_reduce():
+    def _test(red, partial):
+        g = dgl.from_networkx(nx.erdos_renyi_graph(100, 0.1))
+        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+        g.add_edges(g.nodes(), g.nodes())
+        g = g.to(F.ctx())
+        hu, hv, he = generate_feature(g, "none", "none")
+        if partial:
+            nid = F.tensor(list(range(0, 100, 2)), g.idtype)
+
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        with F.record_grad():
+            if partial:
+                g.pull(
+                    nid,
+                    fn.copy_e(e="e", out="m"),
+                    builtin[red](msg="m", out="r1"),
+                )
+            else:
+                g.update_all(
+                    fn.copy_e(e="e", out="m"), builtin[red](msg="m", out="r1")
+                )
+            r1 = g.ndata["r1"]
+            F.backward(F.reduce_sum(r1))
+            e_grad1 = F.grad(g.edata["e"])
+
+        # reset grad
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        with F.record_grad():
+            if partial:
+                g.pull(nid, udf_copy_edge, udf_reduce[red])
+            else:
+                g.update_all(udf_copy_edge, udf_reduce[red])
+            r2 = g.ndata["r2"]
+            F.backward(F.reduce_sum(r2))
+            e_grad2 = F.grad(g.edata["e"])
+
+        def _print_error(a, b):
+            print("ERROR: Test copy_edge_{} partial: {}".format(red, partial))
+            return
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        if not F.allclose(e_grad1, e_grad2):
+            print("edge gradient")
+            _print_error(e_grad1, e_grad2)
+        assert F.allclose(e_grad1, e_grad2)
+
+    _test("sum", False)
+    _test("max", False)
+    _test("mean", False)
+    _test("sum", True)
+    _test("max", True)
+    _test("mean", True)
+
+
+def test_all_binary_builtins():
+    def _test(g, lhs, rhs, binary_op, reducer, partial, nid, broadcast="none"):
+        # initialize node/edge features with uniform(-1, 1)
+        hu, hv, he = generate_feature(g, broadcast, binary_op)
+        if binary_op == "div":
+            # op = div
+            # lhs range: [-1, 1]
+            # rhs range: [1, 2]
+            # result range: [-1, 1]
+            if rhs == "u":
+                hu = (hu + 3) / 2
+            elif rhs == "v":
+                hv = (hv + 3) / 2
+            elif rhs == "e":
+                he = (he + 3) / 2
+
+        if binary_op == "add" or binary_op == "sub":
+            # op = add, sub
+            # lhs range: [-1/2, 1/2]
+            # rhs range: [-1/2, 1/2]
+            # result range: [-1, 1]
+            hu = hu / 2
+            hv = hv / 2
+            he = he / 2
+
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
+        builtin_msg = getattr(fn, builtin_msg_name)
+        builtin_red = getattr(fn, reducer)
+
+        def target_feature_switch(g, target):
+            if target == "u":
+                return g.ndata["u"]
+            elif target == "v":
+                return g.ndata["v"]
+            else:
+                return g.edata["e"]
+
+        with F.record_grad():
+            if partial:
+                g.pull(nid, builtin_msg(lhs, rhs, "m"), builtin_red("m", "r1"))
+            else:
+                g.update_all(builtin_msg(lhs, rhs, "m"), builtin_red("m", "r1"))
+            r1 = g.ndata.pop("r1")
+            F.backward(F.reduce_sum(r1))
+            lhs_grad_1 = F.grad(target_feature_switch(g, lhs))
+            rhs_grad_1 = F.grad(target_feature_switch(g, rhs))
+
+        # reset grad
+        g.ndata["u"] = F.attach_grad(F.clone(hu))
+        g.ndata["v"] = F.attach_grad(F.clone(hv))
+        g.edata["e"] = F.attach_grad(F.clone(he))
+
+        def target_switch(edges, target):
+            if target == "u":
+                return edges.src
+            elif target == "v":
+                return edges.dst
+            elif target == "e":
+                return edges.data
+            else:
+                assert 0, "Unknown target {}".format(target)
+
+        def mfunc(edges):
+            op = getattr(F, binary_op)
+            lhs_data = target_switch(edges, lhs)[lhs]
+            rhs_data = target_switch(edges, rhs)[rhs]
+            # NOTE(zihao): we need to do batched broadcast
+            # e.g. (68, 3, 1) op (68, 5, 3, 4)
+            while F.ndim(lhs_data) < F.ndim(rhs_data):
+                lhs_data = F.unsqueeze(lhs_data, 1)
+            while F.ndim(rhs_data) < F.ndim(lhs_data):
+                rhs_data = F.unsqueeze(rhs_data, 1)
+            return {"m": op(lhs_data, rhs_data)}
+
+        def rfunc(nodes):
+            op = getattr(F, reducer)
+            return {"r2": op(nodes.mailbox["m"], 1)}
+
+        with F.record_grad():
+            if partial:
+                g.pull(nid, mfunc, rfunc)
+            else:
+                g.update_all(mfunc, rfunc)
+            r2 = g.ndata.pop("r2")
+            F.backward(F.reduce_sum(r2), F.tensor([1.0]))
+            lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
+            rhs_grad_2 = F.grad(target_feature_switch(g, rhs))
+
+        rtol = 1e-4
+        atol = 1e-4
+
+        def _print_error(a, b):
+            print(
+                "ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}".format(
+                    lhs, binary_op, rhs, reducer, broadcast, partial
+                )
+            )
+            return
+            if lhs == "u":
+                lhs_data = hu
+            elif lhs == "v":
+                lhs_data = hv
+            elif lhs == "e":
+                lhs_data = he
+
+            if rhs == "u":
+                rhs_data = hu
+            elif rhs == "v":
+                rhs_data = hv
+            elif rhs == "e":
+                rhs_data = he
+            print("lhs", F.asnumpy(lhs_data).tolist())
+            print("rhs", F.asnumpy(rhs_data).tolist())
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y, rtol, atol):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2, rtol, atol):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2, rtol, atol)
+
+        if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol):
+            print("left grad")
+            _print_error(lhs_grad_1, lhs_grad_2)
+        assert F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol)
+
+        if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol):
+            print("right grad")
+            _print_error(rhs_grad_1, rhs_grad_2)
+        assert F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol)
+
+    g = dgl.graph([])
+    g.add_nodes(20)
+    # NOTE(zihao): add self-loop to avoid zero-degree nodes.
+    g.add_edges(g.nodes(), g.nodes())
+    for i in range(2, 18):
+        g.add_edges(0, i)
+        g.add_edges(1, i)
+        g.add_edges(i, 18)
+        g.add_edges(i, 19)
+    g.add_edges(18, 0)
+    g.add_edges(18, 1)
+    g.add_edges(19, 0)
+    g.add_edges(19, 1)
+    g = g.to(F.ctx())
+    nid = F.tensor([0, 1, 4, 5, 7, 12, 14, 15, 18, 19], g.idtype)
+    target = ["u", "v", "e"]
+
+    for lhs, rhs in product(target, target):
+        if lhs == rhs:
+            continue
+        for binary_op in ["add", "sub", "mul", "div"]:
+            for reducer in ["sum", "max", "min", "mean"]:
+                for broadcast in ["none", lhs, rhs]:
+                    for partial in [False, True]:
+                        print(lhs, rhs, binary_op, reducer, broadcast, partial)
+                        _test(
+                            g,
+                            lhs,
+                            rhs,
+                            binary_op,
+                            reducer,
+                            partial,
+                            nid,
+                            broadcast=broadcast,
+                        )
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo-zero-degree"]))
+def test_mean_zero_degree(g, idtype):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata["h"] = F.ones((g.num_nodes(), 3))
+    g.update_all(fn.copy_u("h", "m"), fn.mean("m", "x"))
+    deg = F.asnumpy(g.in_degrees())
+    v = F.tensor(np.where(deg == 0)[0])
+    assert F.allclose(F.gather_row(g.ndata["x"], v), F.zeros((len(v), 3)))
+
+
+if __name__ == "__main__":
+    test_copy_src_reduce()
+    test_copy_edge_reduce()
+    test_all_binary_builtins()

tests/python/common/test_heterograph-misc.py

@@ -0,0 +1,538 @@
+import math
+import numbers
+
+import backend as F
+
+import dgl
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as sp
+from dgl import DGLError
+
+
+# graph generation: a random graph with 10 nodes
+#  and 20 edges.
+#  - has self loop
+#  - no multi edge
+def edge_pair_input(sort=False):
+    if sort:
+        src = [0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 5, 5, 6, 7, 7, 7, 9]
+        dst = [4, 6, 9, 3, 5, 3, 7, 5, 8, 1, 3, 4, 9, 1, 9, 6, 2, 8, 9, 2]
+        return src, dst
+    else:
+        src = [0, 0, 4, 5, 0, 4, 7, 4, 4, 3, 2, 7, 7, 5, 3, 2, 1, 9, 6, 1]
+        dst = [9, 6, 3, 9, 4, 4, 9, 9, 1, 8, 3, 2, 8, 1, 5, 7, 3, 2, 6, 5]
+        return src, dst
+
+
+def nx_input():
+    g = nx.DiGraph()
+    src, dst = edge_pair_input()
+    for i, e in enumerate(zip(src, dst)):
+        g.add_edge(*e, id=i)
+    return g
+
+
+def elist_input():
+    src, dst = edge_pair_input()
+    return list(zip(src, dst))
+
+
+def scipy_coo_input():
+    src, dst = edge_pair_input()
+    return sp.coo_matrix((np.ones((20,)), (src, dst)), shape=(10, 10))
+
+
+def scipy_csr_input():
+    src, dst = edge_pair_input()
+    csr = sp.coo_matrix((np.ones((20,)), (src, dst)), shape=(10, 10)).tocsr()
+    csr.sort_indices()
+    # src = [0 0 0 1 1 2 2 3 3 4 4 4 4 5 5 6 7 7 7 9]
+    # dst = [4 6 9 3 5 3 7 5 8 1 3 4 9 1 9 6 2 8 9 2]
+    return csr
+
+
+def gen_by_mutation():
+    g = dgl.graph([])
+    src, dst = edge_pair_input()
+    g.add_nodes(10)
+    g.add_edges(src, dst)
+    return g
+
+
+def test_query():
+    def _test_one(g):
+        assert g.num_nodes() == 10
+        assert g.num_edges() == 20
+
+        for i in range(10):
+            assert g.has_nodes(i)
+        assert not g.has_nodes(11)
+        assert F.allclose(g.has_nodes([0, 2, 10, 11]), F.tensor([1, 1, 0, 0]))
+
+        src, dst = edge_pair_input()
+        for u, v in zip(src, dst):
+            assert g.has_edges_between(u, v)
+        assert not g.has_edges_between(0, 0)
+        assert F.allclose(
+            g.has_edges_between([0, 0, 3], [0, 9, 8]), F.tensor([0, 1, 1])
+        )
+        assert set(F.asnumpy(g.predecessors(9))) == set([0, 5, 7, 4])
+        assert set(F.asnumpy(g.successors(2))) == set([7, 3])
+
+        assert g.edge_ids(4, 4) == 5
+        assert F.allclose(g.edge_ids([4, 0], [4, 9]), F.tensor([5, 0]))
+
+        src, dst = g.find_edges([3, 6, 5])
+        assert F.allclose(src, F.tensor([5, 7, 4]))
+        assert F.allclose(dst, F.tensor([9, 9, 4]))
+
+        src, dst, eid = g.in_edges(9, form="all")
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set([(0, 9, 0), (5, 9, 3), (7, 9, 6), (4, 9, 7)])
+        src, dst, eid = g.in_edges(
+            [9, 0, 8], form="all"
+        )  # test node#0 has no in edges
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(
+            [(0, 9, 0), (5, 9, 3), (7, 9, 6), (4, 9, 7), (3, 8, 9), (7, 8, 12)]
+        )
+
+        src, dst, eid = g.out_edges(0, form="all")
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set([(0, 9, 0), (0, 6, 1), (0, 4, 4)])
+        src, dst, eid = g.out_edges(
+            [0, 4, 8], form="all"
+        )  # test node#8 has no out edges
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(
+            [
+                (0, 9, 0),
+                (0, 6, 1),
+                (0, 4, 4),
+                (4, 3, 2),
+                (4, 4, 5),
+                (4, 9, 7),
+                (4, 1, 8),
+            ]
+        )
+
+        src, dst, eid = g.edges("all", "eid")
+        t_src, t_dst = edge_pair_input()
+        t_tup = list(zip(t_src, t_dst, list(range(20))))
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(t_tup)
+        assert list(F.asnumpy(eid)) == list(range(20))
+
+        src, dst, eid = g.edges("all", "srcdst")
+        t_src, t_dst = edge_pair_input()
+        t_tup = list(zip(t_src, t_dst, list(range(20))))
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(t_tup)
+        assert list(F.asnumpy(src)) == sorted(list(F.asnumpy(src)))
+
+        assert g.in_degrees(0) == 0
+        assert g.in_degrees(9) == 4
+        assert F.allclose(g.in_degrees([0, 9]), F.tensor([0, 4]))
+        assert g.out_degrees(8) == 0
+        assert g.out_degrees(9) == 1
+        assert F.allclose(g.out_degrees([8, 9]), F.tensor([0, 1]))
+
+        assert np.array_equal(
+            F.sparse_to_numpy(g.adj_external(transpose=True)),
+            scipy_coo_input().toarray().T,
+        )
+        assert np.array_equal(
+            F.sparse_to_numpy(g.adj_external(transpose=False)),
+            scipy_coo_input().toarray(),
+        )
+
+    def _test(g):
+        # test twice to see whether the cached format works or not
+        _test_one(g)
+        _test_one(g)
+
+    def _test_csr_one(g):
+        assert g.num_nodes() == 10
+        assert g.num_edges() == 20
+
+        for i in range(10):
+            assert g.has_nodes(i)
+        assert not g.has_nodes(11)
+        assert F.allclose(g.has_nodes([0, 2, 10, 11]), F.tensor([1, 1, 0, 0]))
+
+        src, dst = edge_pair_input(sort=True)
+        for u, v in zip(src, dst):
+            assert g.has_edges_between(u, v)
+        assert not g.has_edges_between(0, 0)
+        assert F.allclose(
+            g.has_edges_between([0, 0, 3], [0, 9, 8]), F.tensor([0, 1, 1])
+        )
+        assert set(F.asnumpy(g.predecessors(9))) == set([0, 5, 7, 4])
+        assert set(F.asnumpy(g.successors(2))) == set([7, 3])
+
+        # src = [0 0 0 1 1 2 2 3 3 4 4 4 4 5 5 6 7 7 7 9]
+        # dst = [4 6 9 3 5 3 7 5 8 1 3 4 9 1 9 6 2 8 9 2]
+        # eid = [0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9]
+        assert g.edge_ids(4, 4) == 11
+        assert F.allclose(g.edge_ids([4, 0], [4, 9]), F.tensor([11, 2]))
+
+        src, dst = g.find_edges([3, 6, 5])
+        assert F.allclose(src, F.tensor([1, 2, 2]))
+        assert F.allclose(dst, F.tensor([3, 7, 3]))
+
+        src, dst, eid = g.in_edges(9, form="all")
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set([(0, 9, 2), (5, 9, 14), (7, 9, 18), (4, 9, 12)])
+        src, dst, eid = g.in_edges(
+            [9, 0, 8], form="all"
+        )  # test node#0 has no in edges
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(
+            [
+                (0, 9, 2),
+                (5, 9, 14),
+                (7, 9, 18),
+                (4, 9, 12),
+                (3, 8, 8),
+                (7, 8, 17),
+            ]
+        )
+
+        src, dst, eid = g.out_edges(0, form="all")
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set([(0, 9, 2), (0, 6, 1), (0, 4, 0)])
+        src, dst, eid = g.out_edges(
+            [0, 4, 8], form="all"
+        )  # test node#8 has no out edges
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(
+            [
+                (0, 9, 2),
+                (0, 6, 1),
+                (0, 4, 0),
+                (4, 3, 10),
+                (4, 4, 11),
+                (4, 9, 12),
+                (4, 1, 9),
+            ]
+        )
+
+        src, dst, eid = g.edges("all", "eid")
+        t_src, t_dst = edge_pair_input(sort=True)
+        t_tup = list(zip(t_src, t_dst, list(range(20))))
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(t_tup)
+        assert list(F.asnumpy(eid)) == list(range(20))
+
+        src, dst, eid = g.edges("all", "srcdst")
+        t_src, t_dst = edge_pair_input(sort=True)
+        t_tup = list(zip(t_src, t_dst, list(range(20))))
+        tup = list(zip(F.asnumpy(src), F.asnumpy(dst), F.asnumpy(eid)))
+        assert set(tup) == set(t_tup)
+        assert list(F.asnumpy(src)) == sorted(list(F.asnumpy(src)))
+
+        assert g.in_degrees(0) == 0
+        assert g.in_degrees(9) == 4
+        assert F.allclose(g.in_degrees([0, 9]), F.tensor([0, 4]))
+        assert g.out_degrees(8) == 0
+        assert g.out_degrees(9) == 1
+        assert F.allclose(g.out_degrees([8, 9]), F.tensor([0, 1]))
+
+        assert np.array_equal(
+            F.sparse_to_numpy(g.adj_external(transpose=True)),
+            scipy_coo_input().toarray().T,
+        )
+        assert np.array_equal(
+            F.sparse_to_numpy(g.adj_external(transpose=False)),
+            scipy_coo_input().toarray(),
+        )
+
+    def _test_csr(g):
+        # test twice to see whether the cached format works or not
+        _test_csr_one(g)
+        _test_csr_one(g)
+
+    def _test_edge_ids():
+        g = gen_by_mutation()
+        eids = g.edge_ids([4, 0], [4, 9])
+        assert eids.shape[0] == 2
+        eid = g.edge_ids(4, 4)
+        assert isinstance(eid, numbers.Number)
+        with pytest.raises(DGLError):
+            eids = g.edge_ids([9, 0], [4, 9])
+
+        with pytest.raises(DGLError):
+            eid = g.edge_ids(4, 5)
+
+        g.add_edges(0, 4)
+        eids = g.edge_ids([0, 0], [4, 9])
+        eid = g.edge_ids(0, 4)
+
+    _test(gen_by_mutation())
+    _test(dgl.graph(elist_input()))
+    _test(dgl.from_scipy(scipy_coo_input()))
+    _test_csr(dgl.from_scipy(scipy_csr_input()))
+    _test_edge_ids()
+
+
+def test_mutation():
+    g = dgl.graph([])
+    g = g.to(F.ctx())
+    # test add nodes with data
+    g.add_nodes(5)
+    g.add_nodes(5, {"h": F.ones((5, 2))})
+    ans = F.cat([F.zeros((5, 2)), F.ones((5, 2))], 0)
+    assert F.allclose(ans, g.ndata["h"])
+    g.ndata["w"] = 2 * F.ones((10, 2))
+    assert F.allclose(2 * F.ones((10, 2)), g.ndata["w"])
+    # test add edges with data
+    g.add_edges([2, 3], [3, 4])
+    g.add_edges([0, 1], [1, 2], {"m": F.ones((2, 2))})
+    ans = F.cat([F.zeros((2, 2)), F.ones((2, 2))], 0)
+    assert F.allclose(ans, g.edata["m"])
+
+
+def test_scipy_adjmat():
+    g = dgl.graph([])
+    g.add_nodes(10)
+    g.add_edges(range(9), range(1, 10))
+
+    adj_0 = g.adj_external(scipy_fmt="csr")
+    adj_1 = g.adj_external(scipy_fmt="coo")
+    assert np.array_equal(adj_0.toarray(), adj_1.toarray())
+
+    adj_t0 = g.adj_external(transpose=False, scipy_fmt="csr")
+    adj_t_1 = g.adj_external(transpose=False, scipy_fmt="coo")
+    assert np.array_equal(adj_0.toarray(), adj_1.toarray())
+
+
+def test_incmat():
+    g = dgl.graph([])
+    g.add_nodes(4)
+    g.add_edges(0, 1)  # 0
+    g.add_edges(0, 2)  # 1
+    g.add_edges(0, 3)  # 2
+    g.add_edges(2, 3)  # 3
+    g.add_edges(1, 1)  # 4
+    inc_in = F.sparse_to_numpy(g.incidence_matrix("in"))
+    inc_out = F.sparse_to_numpy(g.incidence_matrix("out"))
+    inc_both = F.sparse_to_numpy(g.incidence_matrix("both"))
+    print(inc_in)
+    print(inc_out)
+    print(inc_both)
+    assert np.allclose(
+        inc_in,
+        np.array(
+            [
+                [0.0, 0.0, 0.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0, 0.0, 1.0],
+                [0.0, 1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 1.0, 0.0],
+            ]
+        ),
+    )
+    assert np.allclose(
+        inc_out,
+        np.array(
+            [
+                [1.0, 1.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 1.0],
+                [0.0, 0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0],
+            ]
+        ),
+    )
+    assert np.allclose(
+        inc_both,
+        np.array(
+            [
+                [-1.0, -1.0, -1.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, -1.0, 0.0],
+                [0.0, 0.0, 1.0, 1.0, 0.0],
+            ]
+        ),
+    )
+
+
+def test_find_edges():
+    g = dgl.graph([])
+    g.add_nodes(10)
+    g.add_edges(range(9), range(1, 10))
+    e = g.find_edges([1, 3, 2, 4])
+    assert (
+        F.asnumpy(e[0][0]) == 1
+        and F.asnumpy(e[0][1]) == 3
+        and F.asnumpy(e[0][2]) == 2
+        and F.asnumpy(e[0][3]) == 4
+    )
+    assert (
+        F.asnumpy(e[1][0]) == 2
+        and F.asnumpy(e[1][1]) == 4
+        and F.asnumpy(e[1][2]) == 3
+        and F.asnumpy(e[1][3]) == 5
+    )
+
+    try:
+        g.find_edges([10])
+        fail = False
+    except DGLError:
+        fail = True
+    finally:
+        assert fail
+
+
+def test_ismultigraph():
+    g = dgl.graph([])
+    g.add_nodes(10)
+    assert g.is_multigraph == False
+    g.add_edges([0], [0])
+    assert g.is_multigraph == False
+    g.add_edges([1], [2])
+    assert g.is_multigraph == False
+    g.add_edges([0, 2], [0, 3])
+    assert g.is_multigraph == True
+
+
+def test_hypersparse_query():
+    g = dgl.graph([])
+    g = g.to(F.ctx())
+    g.add_nodes(1000001)
+    g.add_edges([0], [1])
+    for i in range(10):
+        assert g.has_nodes(i)
+    assert not g.has_nodes(1000002)
+    assert g.edge_ids(0, 1) == 0
+    src, dst = g.find_edges([0])
+    src, dst, eid = g.in_edges(1, form="all")
+    src, dst, eid = g.out_edges(0, form="all")
+    src, dst = g.edges()
+    assert g.in_degrees(0) == 0
+    assert g.in_degrees(1) == 1
+    assert g.out_degrees(0) == 1
+    assert g.out_degrees(1) == 0
+
+
+def test_empty_data_initialized():
+    g = dgl.graph([])
+    g = g.to(F.ctx())
+    g.ndata["ha"] = F.tensor([])
+    g.add_nodes(1, {"hb": F.tensor([1])})
+    assert "ha" in g.ndata
+    assert len(g.ndata["ha"]) == 1
+
+
+def test_is_sorted():
+    u_src, u_dst = edge_pair_input(False)
+    s_src, s_dst = edge_pair_input(True)
+
+    u_src = F.tensor(u_src, dtype=F.int32)
+    u_dst = F.tensor(u_dst, dtype=F.int32)
+    s_src = F.tensor(s_src, dtype=F.int32)
+    s_dst = F.tensor(s_dst, dtype=F.int32)
+
+    src_sorted, dst_sorted = dgl.utils.is_sorted_srcdst(u_src, u_dst)
+    assert src_sorted == False
+    assert dst_sorted == False
+
+    src_sorted, dst_sorted = dgl.utils.is_sorted_srcdst(s_src, s_dst)
+    assert src_sorted == True
+    assert dst_sorted == True
+
+    src_sorted, dst_sorted = dgl.utils.is_sorted_srcdst(u_src, u_dst)
+    assert src_sorted == False
+    assert dst_sorted == False
+
+    src_sorted, dst_sorted = dgl.utils.is_sorted_srcdst(s_src, u_dst)
+    assert src_sorted == True
+    assert dst_sorted == False
+
+
+def test_default_types():
+    dg = dgl.graph([])
+    g = dgl.graph(([], []))
+    assert dg.ntypes == g.ntypes
+    assert dg.etypes == g.etypes
+
+
+def test_formats():
+    g = dgl.rand_graph(10, 20)
+    # in_degrees works if coo or csc available
+    # out_degrees works if coo or csr available
+    try:
+        g.in_degrees()
+        g.out_degrees()
+        g.formats("coo").in_degrees()
+        g.formats("coo").out_degrees()
+        g.formats("csc").in_degrees()
+        g.formats("csr").out_degrees()
+        fail = False
+    except DGLError:
+        fail = True
+    finally:
+        assert not fail
+    # in_degrees NOT works if csc available only
+    try:
+        g.formats("csc").out_degrees()
+        fail = True
+    except DGLError:
+        fail = False
+    finally:
+        assert not fail
+    # out_degrees NOT works if csr available only
+    try:
+        g.formats("csr").in_degrees()
+        fail = True
+    except DGLError:
+        fail = False
+    finally:
+        assert not fail
+
+    # If the intersection of created formats and allowed formats is
+    # not empty, then retain the intersection.
+    # Case1: intersection is not empty and intersected is equal to
+    # created formats.
+    g = g.formats(["coo", "csr"])
+    g.create_formats_()
+    g = g.formats(["coo", "csr", "csc"])
+    assert sorted(g.formats()["created"]) == sorted(["coo", "csr"])
+    assert sorted(g.formats()["not created"]) == sorted(["csc"])
+
+    # Case2: intersection is not empty and intersected is not equal
+    # to created formats.
+    g = g.formats(["coo", "csr"])
+    g.create_formats_()
+    g = g.formats(["coo", "csc"])
+    assert sorted(g.formats()["created"]) == sorted(["coo"])
+    assert sorted(g.formats()["not created"]) == sorted(["csc"])
+
+    # If the intersection of created formats and allowed formats is
+    # empty, then create a format in the order of `coo` -> `csr` ->
+    # `csc`.
+    # Case1: intersection is empty and just one format is allowed.
+    g = g.formats(["coo", "csr"])
+    g.create_formats_()
+    g = g.formats(["csc"])
+    assert sorted(g.formats()["created"]) == sorted(["csc"])
+    assert sorted(g.formats()["not created"]) == sorted([])
+
+    # Case2: intersection is empty and more than one format is allowed.
+    g = g.formats("csc")
+    g.create_formats_()
+    g = g.formats(["csr", "coo"])
+    assert sorted(g.formats()["created"]) == sorted(["coo"])
+    assert sorted(g.formats()["not created"]) == sorted(["csr"])
+
+
+if __name__ == "__main__":
+    test_query()
+    test_mutation()
+    test_scipy_adjmat()
+    test_incmat()
+    test_find_edges()
+    test_hypersparse_query()
+    test_is_sorted()
+    test_default_types()
+    test_formats()

tests/python/common/test_heterograph-pickle.py

@@ -0,0 +1,235 @@
+import io
+import pickle
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import pytest
+import scipy.sparse as ssp
+from dgl.graph_index import create_graph_index
+from dgl.utils import toindex
+from utils import (
+    assert_is_identical,
+    assert_is_identical_hetero,
+    check_graph_equal,
+    get_cases,
+    parametrize_idtype,
+)
+
+
+def _assert_is_identical_nodeflow(nf1, nf2):
+    assert nf1.num_nodes() == nf2.num_nodes()
+    src, dst = nf1.all_edges()
+    src2, dst2 = nf2.all_edges()
+    assert F.array_equal(src, src2)
+    assert F.array_equal(dst, dst2)
+
+    assert nf1.num_layers == nf2.num_layers
+    for i in range(nf1.num_layers):
+        assert nf1.layer_size(i) == nf2.layer_size(i)
+        assert nf1.layers[i].data.keys() == nf2.layers[i].data.keys()
+        for k in nf1.layers[i].data:
+            assert F.allclose(nf1.layers[i].data[k], nf2.layers[i].data[k])
+    assert nf1.num_blocks == nf2.num_blocks
+    for i in range(nf1.num_blocks):
+        assert nf1.block_size(i) == nf2.block_size(i)
+        assert nf1.blocks[i].data.keys() == nf2.blocks[i].data.keys()
+        for k in nf1.blocks[i].data:
+            assert F.allclose(nf1.blocks[i].data[k], nf2.blocks[i].data[k])
+
+
+def _assert_is_identical_batchedgraph(bg1, bg2):
+    assert_is_identical(bg1, bg2)
+    assert bg1.batch_size == bg2.batch_size
+    assert bg1.batch_num_nodes == bg2.batch_num_nodes
+    assert bg1.batch_num_edges == bg2.batch_num_edges
+
+
+def _assert_is_identical_batchedhetero(bg1, bg2):
+    assert_is_identical_hetero(bg1, bg2)
+    for ntype in bg1.ntypes:
+        assert bg1.batch_num_nodes(ntype) == bg2.batch_num_nodes(ntype)
+    for canonical_etype in bg1.canonical_etypes:
+        assert bg1.batch_num_edges(canonical_etype) == bg2.batch_num_edges(
+            canonical_etype
+        )
+
+
+def _assert_is_identical_index(i1, i2):
+    assert i1.slice_data() == i2.slice_data()
+    assert F.array_equal(i1.tousertensor(), i2.tousertensor())
+
+
+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_index():
+    # normal index
+    i = toindex([1, 2, 3])
+    i.tousertensor()
+    i.todgltensor()  # construct a dgl tensor which is unpicklable
+    i2 = _reconstruct_pickle(i)
+    _assert_is_identical_index(i, i2)
+
+    # slice index
+    i = toindex(slice(5, 10))
+    i2 = _reconstruct_pickle(i)
+    _assert_is_identical_index(i, i2)
+
+
+def test_pickling_graph_index():
+    gi = create_graph_index(None, False)
+    gi.add_nodes(3)
+    src_idx = toindex([0, 0])
+    dst_idx = toindex([1, 2])
+    gi.add_edges(src_idx, dst_idx)
+
+    gi2 = _reconstruct_pickle(gi)
+
+    assert gi2.num_nodes() == gi.num_nodes()
+    src_idx2, dst_idx2, _ = gi2.edges()
+    assert F.array_equal(src_idx.tousertensor(), src_idx2.tousertensor())
+    assert F.array_equal(dst_idx.tousertensor(), dst_idx2.tousertensor())
+
+
+def _global_message_func(nodes):
+    return {"x": nodes.data["x"]}
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@parametrize_idtype
+@pytest.mark.parametrize(
+    "g", get_cases(exclude=["dglgraph", "two_hetero_batch"])
+)
+def test_pickling_graph(g, idtype):
+    g = g.astype(idtype)
+    new_g = _reconstruct_pickle(g)
+    check_graph_equal(g, new_g, check_feature=True)
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+def test_pickling_batched_heterograph():
+    # copied from test_heterograph.create_test_heterograph()
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        }
+    )
+    g2 = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        }
+    )
+
+    g.nodes["user"].data["u_h"] = F.randn((3, 4))
+    g.nodes["game"].data["g_h"] = F.randn((2, 5))
+    g.edges["plays"].data["p_h"] = F.randn((4, 6))
+    g2.nodes["user"].data["u_h"] = F.randn((3, 4))
+    g2.nodes["game"].data["g_h"] = F.randn((2, 5))
+    g2.edges["plays"].data["p_h"] = F.randn((4, 6))
+
+    bg = dgl.batch([g, g2])
+    new_bg = _reconstruct_pickle(bg)
+    check_graph_equal(bg, new_bg)
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="GPU edge_subgraph w/ relabeling not implemented",
+)
+def test_pickling_subgraph():
+    f1 = io.BytesIO()
+    f2 = io.BytesIO()
+    g = dgl.rand_graph(10000, 100000)
+    g.ndata["x"] = F.randn((10000, 4))
+    g.edata["x"] = F.randn((100000, 5))
+    pickle.dump(g, f1)
+    sg = g.subgraph([0, 1])
+    sgx = sg.ndata["x"]  # materialize
+    pickle.dump(sg, f2)
+    # TODO(BarclayII): How should I test that the size of the subgraph pickle file should not
+    # be as large as the size of the original pickle file?
+    assert f1.tell() > f2.tell() * 50
+
+    f2.seek(0)
+    f2.truncate()
+    sgx = sg.edata["x"]  # materialize
+    pickle.dump(sg, f2)
+    assert f1.tell() > f2.tell() * 50
+
+    f2.seek(0)
+    f2.truncate()
+    sg = g.edge_subgraph([0])
+    sgx = sg.edata["x"]  # materialize
+    pickle.dump(sg, f2)
+    assert f1.tell() > f2.tell() * 50
+
+    f2.seek(0)
+    f2.truncate()
+    sgx = sg.ndata["x"]  # materialize
+    pickle.dump(sg, f2)
+    assert f1.tell() > f2.tell() * 50
+
+    f1.close()
+    f2.close()
+
+
+@unittest.skipIf(F._default_context_str != "gpu", reason="Need GPU for pin")
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="TensorFlow create graph on gpu when unpickle",
+)
+@parametrize_idtype
+def test_pickling_is_pinned(idtype):
+    from copy import deepcopy
+
+    g = dgl.rand_graph(10, 20, idtype=idtype, device=F.cpu())
+    hg = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        },
+        idtype=idtype,
+        device=F.cpu(),
+    )
+    for graph in [g, hg]:
+        assert not graph.is_pinned()
+        graph.pin_memory_()
+        assert graph.is_pinned()
+        pg = _reconstruct_pickle(graph)
+        assert pg.is_pinned()
+        pg.unpin_memory_()
+        dg = deepcopy(graph)
+        assert dg.is_pinned()
+        dg.unpin_memory_()
+        graph.unpin_memory_()
+
+
+if __name__ == "__main__":
+    test_pickling_index()
+    test_pickling_graph_index()
+    test_pickling_frame()
+    test_pickling_graph()
+    test_pickling_nodeflow()
+    test_pickling_batched_graph()
+    test_pickling_heterograph()
+    test_pickling_batched_heterograph()
+    test_pickling_is_pinned()

tests/python/common/test_heterograph-remove.py

@@ -0,0 +1,223 @@
+import backend as F
+
+import dgl
+import numpy as np
+from utils import parametrize_idtype
+
+
+def create_graph(idtype, num_node):
+    g = dgl.graph([])
+    g = g.astype(idtype).to(F.ctx())
+    g.add_nodes(num_node)
+    return g
+
+
+@parametrize_idtype
+def test_node_removal(idtype):
+    g = create_graph(idtype, 10)
+    g.add_edges(0, 0)
+    assert g.num_nodes() == 10
+    g.ndata["id"] = F.arange(0, 10)
+
+    # remove nodes
+    g.remove_nodes(range(4, 7))
+    assert g.num_nodes() == 7
+    assert F.array_equal(g.ndata["id"], F.tensor([0, 1, 2, 3, 7, 8, 9]))
+    assert dgl.NID not in g.ndata
+    assert dgl.EID not in g.edata
+
+    # add nodes
+    g.add_nodes(3)
+    assert g.num_nodes() == 10
+    assert F.array_equal(
+        g.ndata["id"], F.tensor([0, 1, 2, 3, 7, 8, 9, 0, 0, 0])
+    )
+
+    # remove nodes
+    g.remove_nodes(range(1, 4), store_ids=True)
+    assert g.num_nodes() == 7
+    assert F.array_equal(g.ndata["id"], F.tensor([0, 7, 8, 9, 0, 0, 0]))
+    assert dgl.NID in g.ndata
+    assert dgl.EID in g.edata
+
+
+@parametrize_idtype
+def test_multigraph_node_removal(idtype):
+    g = create_graph(idtype, 5)
+    for i in range(5):
+        g.add_edges(i, i)
+        g.add_edges(i, i)
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 10
+
+    # remove nodes
+    g.remove_nodes([2, 3])
+    assert g.num_nodes() == 3
+    assert g.num_edges() == 6
+
+    # add nodes
+    g.add_nodes(1)
+    g.add_edges(1, 1)
+    g.add_edges(1, 1)
+    assert g.num_nodes() == 4
+    assert g.num_edges() == 8
+
+    # remove nodes
+    g.remove_nodes([0])
+    assert g.num_nodes() == 3
+    assert g.num_edges() == 6
+
+
+@parametrize_idtype
+def test_multigraph_edge_removal(idtype):
+    g = create_graph(idtype, 5)
+    for i in range(5):
+        g.add_edges(i, i)
+        g.add_edges(i, i)
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 10
+
+    # remove edges
+    g.remove_edges([2, 3])
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 8
+
+    # add edges
+    g.add_edges(1, 1)
+    g.add_edges(1, 1)
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 10
+
+    # remove edges
+    g.remove_edges([0, 1])
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 8
+
+
+@parametrize_idtype
+def test_edge_removal(idtype):
+    g = create_graph(idtype, 5)
+    for i in range(5):
+        for j in range(5):
+            g.add_edges(i, j)
+    g.edata["id"] = F.arange(0, 25)
+
+    # remove edges
+    g.remove_edges(range(13, 20))
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 18
+    assert F.array_equal(
+        g.edata["id"], F.tensor(list(range(13)) + list(range(20, 25)))
+    )
+    assert dgl.NID not in g.ndata
+    assert dgl.EID not in g.edata
+
+    # add edges
+    g.add_edges(3, 3)
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 19
+    assert F.array_equal(
+        g.edata["id"], F.tensor(list(range(13)) + list(range(20, 25)) + [0])
+    )
+
+    # remove edges
+    g.remove_edges(range(2, 10), store_ids=True)
+    assert g.num_nodes() == 5
+    assert g.num_edges() == 11
+    assert F.array_equal(
+        g.edata["id"], F.tensor([0, 1, 10, 11, 12, 20, 21, 22, 23, 24, 0])
+    )
+    assert dgl.EID in g.edata
+
+
+@parametrize_idtype
+def test_node_and_edge_removal(idtype):
+    g = create_graph(idtype, 10)
+    for i in range(10):
+        for j in range(10):
+            g.add_edges(i, j)
+    g.edata["id"] = F.arange(0, 100)
+    assert g.num_nodes() == 10
+    assert g.num_edges() == 100
+
+    # remove nodes
+    g.remove_nodes([2, 4])
+    assert g.num_nodes() == 8
+    assert g.num_edges() == 64
+
+    # remove edges
+    g.remove_edges(range(10, 20))
+    assert g.num_nodes() == 8
+    assert g.num_edges() == 54
+
+    # add nodes
+    g.add_nodes(2)
+    assert g.num_nodes() == 10
+    assert g.num_edges() == 54
+
+    # add edges
+    for i in range(8, 10):
+        for j in range(8, 10):
+            g.add_edges(i, j)
+    assert g.num_nodes() == 10
+    assert g.num_edges() == 58
+
+    # remove edges
+    g.remove_edges(range(10, 20))
+    assert g.num_nodes() == 10
+    assert g.num_edges() == 48
+
+
+@parametrize_idtype
+def test_node_frame(idtype):
+    g = create_graph(idtype, 10)
+    data = np.random.rand(10, 3)
+    new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
+    g.ndata["h"] = F.tensor(data)
+
+    # remove nodes
+    g.remove_nodes(range(3, 7))
+    assert F.allclose(g.ndata["h"], F.tensor(new_data))
+
+
+@parametrize_idtype
+def test_edge_frame(idtype):
+    g = create_graph(idtype, 10)
+    g.add_edges(list(range(10)), list(range(1, 10)) + [0])
+    data = np.random.rand(10, 3)
+    new_data = data.take([0, 1, 2, 7, 8, 9], axis=0)
+    g.edata["h"] = F.tensor(data)
+
+    # remove edges
+    g.remove_edges(range(3, 7))
+    assert F.allclose(g.edata["h"], F.tensor(new_data))
+
+
+@parametrize_idtype
+def test_issue1287(idtype):
+    # reproduce https://github.com/dmlc/dgl/issues/1287.
+    # setting features after remove nodes
+    g = create_graph(idtype, 5)
+    g.add_edges([0, 2, 3, 1, 1], [1, 0, 3, 1, 0])
+    g.remove_nodes([0, 1])
+    g.ndata["h"] = F.randn((g.num_nodes(), 3))
+    g.edata["h"] = F.randn((g.num_edges(), 2))
+
+    # remove edges
+    g = create_graph(idtype, 5)
+    g.add_edges([0, 2, 3, 1, 1], [1, 0, 3, 1, 0])
+    g.remove_edges([0, 1])
+    g = g.to(F.ctx())
+    g.ndata["h"] = F.randn((g.num_nodes(), 3))
+    g.edata["h"] = F.randn((g.num_edges(), 2))
+
+
+if __name__ == "__main__":
+    test_node_removal()
+    test_edge_removal()
+    test_multigraph_node_removal()
+    test_multigraph_edge_removal()
+    test_node_and_edge_removal()
+    test_node_frame()
+    test_edge_frame()
+    test_frame_size()

tests/python/common/test_heterograph-shared-memory.py

@@ -0,0 +1,108 @@
+import io
+import multiprocessing as mp
+import os
+import pickle
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import scipy.sparse as ssp
+from dgl.graph_index import create_graph_index
+from dgl.utils import toindex
+from utils import parametrize_idtype
+
+
+def create_test_graph(idtype):
+    g = dgl.heterograph(
+        (
+            {
+                ("user", "follows", "user"): ([0, 1], [1, 2]),
+                ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+                ("user", "wishes", "game"): ([0, 2], [1, 0]),
+                ("developer", "develops", "game"): ([0, 1], [0, 1]),
+            }
+        ),
+        idtype=idtype,
+    )
+    return g
+
+
+def _assert_is_identical_hetero(g, g2):
+    assert g.ntypes == g2.ntypes
+    assert g.canonical_etypes == g2.canonical_etypes
+
+    # check if two metagraphs are identical
+    for edges, features in g.metagraph().edges(keys=True).items():
+        assert g2.metagraph().edges(keys=True)[edges] == features
+
+    # check if node ID spaces and feature spaces are equal
+    for ntype in g.ntypes:
+        assert g.num_nodes(ntype) == g2.num_nodes(ntype)
+
+    # check if edge ID spaces and feature spaces are equal
+    for etype in g.canonical_etypes:
+        src, dst = g.all_edges(etype=etype, order="eid")
+        src2, dst2 = g2.all_edges(etype=etype, order="eid")
+        assert F.array_equal(src, src2)
+        assert F.array_equal(dst, dst2)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="Not support tensorflow for now",
+)
+@parametrize_idtype
+def test_single_process(idtype):
+    hg = create_test_graph(idtype=idtype)
+    hg_share = hg.shared_memory("hg")
+    hg_rebuild = dgl.hetero_from_shared_memory("hg")
+    hg_save_again = hg_rebuild.shared_memory("hg")
+    _assert_is_identical_hetero(hg, hg_share)
+    _assert_is_identical_hetero(hg, hg_rebuild)
+    _assert_is_identical_hetero(hg, hg_save_again)
+
+
+def sub_proc(hg_origin, name):
+    hg_rebuild = dgl.hetero_from_shared_memory(name)
+    hg_save_again = hg_rebuild.shared_memory(name)
+    _assert_is_identical_hetero(hg_origin, hg_rebuild)
+    _assert_is_identical_hetero(hg_origin, hg_save_again)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="Not support tensorflow for now",
+)
+@parametrize_idtype
+def test_multi_process(idtype):
+    hg = create_test_graph(idtype=idtype)
+    hg_share = hg.shared_memory("hg1")
+    p = mp.Process(target=sub_proc, args=(hg, "hg1"))
+    p.start()
+    p.join()
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="Need gpu for this test"
+)
+@unittest.skipIf(
+    dgl.backend.backend_name == "tensorflow",
+    reason="Not support tensorflow for now",
+)
+def test_copy_from_gpu():
+    hg = create_test_graph(idtype=F.int32)
+    hg_gpu = hg.to(F.cuda())
+    hg_share = hg_gpu.shared_memory("hg_gpu")
+    p = mp.Process(target=sub_proc, args=(hg, "hg_gpu"))
+    p.start()
+    p.join()
+
+
+# TODO: Test calling shared_memory with Blocks (a subclass of HeteroGraph)
+if __name__ == "__main__":
+    test_single_process(F.int64)
+    test_multi_process(F.int32)
+    test_copy_from_gpu()

tests/python/common/test_heterograph-specialization.py

@@ -0,0 +1,385 @@
+import backend as F
+
+import dgl
+import dgl.function as fn
+import numpy as np
+import scipy.sparse as sp
+from utils import parametrize_idtype
+
+D = 5
+
+
+def generate_graph(idtype):
+    g = dgl.graph([])
+    g = g.astype(idtype).to(F.ctx())
+    g.add_nodes(10)
+    # create a graph where 0 is the source and 9 is the sink
+    for i in range(1, 9):
+        g.add_edges(0, i)
+        g.add_edges(i, 9)
+    # add a back flow from 9 to 0
+    g.add_edges(9, 0)
+    g.ndata.update({"f1": F.randn((10,)), "f2": F.randn((10, D))})
+    weights = F.randn((17,))
+    g.edata.update({"e1": weights, "e2": F.unsqueeze(weights, 1)})
+    return g
+
+
+@parametrize_idtype
+def test_v2v_update_all(idtype):
+    def _test(fld):
+        def message_func(edges):
+            return {"m": edges.src[fld]}
+
+        def message_func_edge(edges):
+            if len(edges.src[fld].shape) == 1:
+                return {"m": edges.src[fld] * edges.data["e1"]}
+            else:
+                return {"m": edges.src[fld] * edges.data["e2"]}
+
+        def reduce_func(nodes):
+            return {fld: F.sum(nodes.mailbox["m"], 1)}
+
+        def apply_func(nodes):
+            return {fld: 2 * nodes.data[fld]}
+
+        g = generate_graph(idtype)
+        # update all
+        v1 = g.ndata[fld]
+        g.update_all(
+            fn.copy_u(u=fld, out="m"), fn.sum(msg="m", out=fld), apply_func
+        )
+        v2 = g.ndata[fld]
+        g.ndata.update({fld: v1})
+        g.update_all(message_func, reduce_func, apply_func)
+        v3 = g.ndata[fld]
+        assert F.allclose(v2, v3)
+        # update all with edge weights
+        v1 = g.ndata[fld]
+        g.update_all(
+            fn.u_mul_e(fld, "e1", "m"), fn.sum(msg="m", out=fld), apply_func
+        )
+        v2 = g.ndata[fld]
+        g.ndata.update({fld: v1})
+        g.update_all(message_func_edge, reduce_func, apply_func)
+        v4 = g.ndata[fld]
+        assert F.allclose(v2, v4)
+
+    # test 1d node features
+    _test("f1")
+    # test 2d node features
+    _test("f2")
+
+
+@parametrize_idtype
+def test_v2v_snr(idtype):
+    u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
+    v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
+
+    def _test(fld):
+        def message_func(edges):
+            return {"m": edges.src[fld]}
+
+        def message_func_edge(edges):
+            if len(edges.src[fld].shape) == 1:
+                return {"m": edges.src[fld] * edges.data["e1"]}
+            else:
+                return {"m": edges.src[fld] * edges.data["e2"]}
+
+        def reduce_func(nodes):
+            return {fld: F.sum(nodes.mailbox["m"], 1)}
+
+        def apply_func(nodes):
+            return {fld: 2 * nodes.data[fld]}
+
+        g = generate_graph(idtype)
+        # send and recv
+        v1 = g.ndata[fld]
+        g.send_and_recv(
+            (u, v),
+            fn.copy_u(u=fld, out="m"),
+            fn.sum(msg="m", out=fld),
+            apply_func,
+        )
+        v2 = g.ndata[fld]
+        g.ndata.update({fld: v1})
+        g.send_and_recv((u, v), message_func, reduce_func, apply_func)
+        v3 = g.ndata[fld]
+        assert F.allclose(v2, v3)
+        # send and recv with edge weights
+        v1 = g.ndata[fld]
+        g.send_and_recv(
+            (u, v),
+            fn.u_mul_e(fld, "e1", "m"),
+            fn.sum(msg="m", out=fld),
+            apply_func,
+        )
+        v2 = g.ndata[fld]
+        g.ndata.update({fld: v1})
+        g.send_and_recv((u, v), message_func_edge, reduce_func, apply_func)
+        v4 = g.ndata[fld]
+        assert F.allclose(v2, v4)
+
+    # test 1d node features
+    _test("f1")
+    # test 2d node features
+    _test("f2")
+
+
+@parametrize_idtype
+def test_v2v_pull(idtype):
+    nodes = F.tensor([1, 2, 3, 9], idtype)
+
+    def _test(fld):
+        def message_func(edges):
+            return {"m": edges.src[fld]}
+
+        def message_func_edge(edges):
+            if len(edges.src[fld].shape) == 1:
+                return {"m": edges.src[fld] * edges.data["e1"]}
+            else:
+                return {"m": edges.src[fld] * edges.data["e2"]}
+
+        def reduce_func(nodes):
+            return {fld: F.sum(nodes.mailbox["m"], 1)}
+
+        def apply_func(nodes):
+            return {fld: 2 * nodes.data[fld]}
+
+        g = generate_graph(idtype)
+        # send and recv
+        v1 = g.ndata[fld]
+        g.pull(
+            nodes,
+            fn.copy_u(u=fld, out="m"),
+            fn.sum(msg="m", out=fld),
+            apply_func,
+        )
+        v2 = g.ndata[fld]
+        g.ndata[fld] = v1
+        g.pull(nodes, message_func, reduce_func, apply_func)
+        v3 = g.ndata[fld]
+        assert F.allclose(v2, v3)
+        # send and recv with edge weights
+        v1 = g.ndata[fld]
+        g.pull(
+            nodes,
+            fn.u_mul_e(fld, "e1", "m"),
+            fn.sum(msg="m", out=fld),
+            apply_func,
+        )
+        v2 = g.ndata[fld]
+        g.ndata[fld] = v1
+        g.pull(nodes, message_func_edge, reduce_func, apply_func)
+        v4 = g.ndata[fld]
+        assert F.allclose(v2, v4)
+
+    # test 1d node features
+    _test("f1")
+    # test 2d node features
+    _test("f2")
+
+
+@parametrize_idtype
+def test_update_all_multi_fallback(idtype):
+    # create a graph with zero in degree nodes
+    g = dgl.graph([])
+    g = g.astype(idtype).to(F.ctx())
+    g.add_nodes(10)
+    for i in range(1, 9):
+        g.add_edges(0, i)
+        g.add_edges(i, 9)
+    g.ndata["h"] = F.randn((10, D))
+    g.edata["w1"] = F.randn((16,))
+    g.edata["w2"] = F.randn((16, D))
+
+    def _mfunc_hxw1(edges):
+        return {"m1": edges.src["h"] * F.unsqueeze(edges.data["w1"], 1)}
+
+    def _mfunc_hxw2(edges):
+        return {"m2": edges.src["h"] * edges.data["w2"]}
+
+    def _rfunc_m1(nodes):
+        return {"o1": F.sum(nodes.mailbox["m1"], 1)}
+
+    def _rfunc_m2(nodes):
+        return {"o2": F.sum(nodes.mailbox["m2"], 1)}
+
+    def _rfunc_m1max(nodes):
+        return {"o3": F.max(nodes.mailbox["m1"], 1)}
+
+    def _afunc(nodes):
+        ret = {}
+        for k, v in nodes.data.items():
+            if k.startswith("o"):
+                ret[k] = 2 * v
+        return ret
+
+    # compute ground truth
+    g.update_all(_mfunc_hxw1, _rfunc_m1, _afunc)
+    o1 = g.ndata.pop("o1")
+    g.update_all(_mfunc_hxw2, _rfunc_m2, _afunc)
+    o2 = g.ndata.pop("o2")
+    g.update_all(_mfunc_hxw1, _rfunc_m1max, _afunc)
+    o3 = g.ndata.pop("o3")
+    # v2v spmv
+    g.update_all(
+        fn.u_mul_e("h", "w1", "m1"), fn.sum(msg="m1", out="o1"), _afunc
+    )
+    assert F.allclose(o1, g.ndata.pop("o1"))
+    # v2v fallback to e2v
+    g.update_all(
+        fn.u_mul_e("h", "w2", "m2"), fn.sum(msg="m2", out="o2"), _afunc
+    )
+    assert F.allclose(o2, g.ndata.pop("o2"))
+
+
+@parametrize_idtype
+def test_pull_multi_fallback(idtype):
+    # create a graph with zero in degree nodes
+    g = dgl.graph([])
+    g = g.astype(idtype).to(F.ctx())
+    g.add_nodes(10)
+    for i in range(1, 9):
+        g.add_edges(0, i)
+        g.add_edges(i, 9)
+    g.ndata["h"] = F.randn((10, D))
+    g.edata["w1"] = F.randn((16,))
+    g.edata["w2"] = F.randn((16, D))
+
+    def _mfunc_hxw1(edges):
+        return {"m1": edges.src["h"] * F.unsqueeze(edges.data["w1"], 1)}
+
+    def _mfunc_hxw2(edges):
+        return {"m2": edges.src["h"] * edges.data["w2"]}
+
+    def _rfunc_m1(nodes):
+        return {"o1": F.sum(nodes.mailbox["m1"], 1)}
+
+    def _rfunc_m2(nodes):
+        return {"o2": F.sum(nodes.mailbox["m2"], 1)}
+
+    def _rfunc_m1max(nodes):
+        return {"o3": F.max(nodes.mailbox["m1"], 1)}
+
+    def _afunc(nodes):
+        ret = {}
+        for k, v in nodes.data.items():
+            if k.startswith("o"):
+                ret[k] = 2 * v
+        return ret
+
+    # nodes to pull
+    def _pull_nodes(nodes):
+        # compute ground truth
+        g.pull(nodes, _mfunc_hxw1, _rfunc_m1, _afunc)
+        o1 = g.ndata.pop("o1")
+        g.pull(nodes, _mfunc_hxw2, _rfunc_m2, _afunc)
+        o2 = g.ndata.pop("o2")
+        g.pull(nodes, _mfunc_hxw1, _rfunc_m1max, _afunc)
+        o3 = g.ndata.pop("o3")
+        # v2v spmv
+        g.pull(
+            nodes,
+            fn.u_mul_e("h", "w1", "m1"),
+            fn.sum(msg="m1", out="o1"),
+            _afunc,
+        )
+        assert F.allclose(o1, g.ndata.pop("o1"))
+        # v2v fallback to e2v
+        g.pull(
+            nodes,
+            fn.u_mul_e("h", "w2", "m2"),
+            fn.sum(msg="m2", out="o2"),
+            _afunc,
+        )
+        assert F.allclose(o2, g.ndata.pop("o2"))
+
+    # test#1: non-0deg nodes
+    nodes = [1, 2, 9]
+    _pull_nodes(nodes)
+    # test#2: 0deg nodes + non-0deg nodes
+    nodes = [0, 1, 2, 9]
+    _pull_nodes(nodes)
+
+
+@parametrize_idtype
+def test_spmv_3d_feat(idtype):
+    def src_mul_edge_udf(edges):
+        return {
+            "sum": edges.src["h"]
+            * F.unsqueeze(F.unsqueeze(edges.data["h"], 1), 1)
+        }
+
+    def sum_udf(nodes):
+        return {"h": F.sum(nodes.mailbox["sum"], 1)}
+
+    n = 100
+    p = 0.1
+    a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
+    g = dgl.from_scipy(a)
+    g = g.astype(idtype).to(F.ctx())
+    m = g.num_edges()
+
+    # test#1: v2v with adj data
+    h = F.randn((n, 5, 5))
+    e = F.randn((m,))
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(
+        message_func=fn.u_mul_e("h", "h", "sum"), reduce_func=fn.sum("sum", "h")
+    )  # 1
+    ans = g.ndata["h"]
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(
+        message_func=src_mul_edge_udf, reduce_func=fn.sum("sum", "h")
+    )  # 2
+    assert F.allclose(g.ndata["h"], ans)
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf)  # 3
+    assert F.allclose(g.ndata["h"], ans)
+
+    # test#2: e2v
+    def src_mul_edge_udf(edges):
+        return {"sum": edges.src["h"] * edges.data["h"]}
+
+    h = F.randn((n, 5, 5))
+    e = F.randn((m, 5, 5))
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(
+        message_func=fn.u_mul_e("h", "h", "sum"), reduce_func=fn.sum("sum", "h")
+    )  # 1
+    ans = g.ndata["h"]
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(
+        message_func=src_mul_edge_udf, reduce_func=fn.sum("sum", "h")
+    )  # 2
+    assert F.allclose(g.ndata["h"], ans)
+
+    g.ndata["h"] = h
+    g.edata["h"] = e
+    g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf)  # 3
+    assert F.allclose(g.ndata["h"], ans)
+
+
+if __name__ == "__main__":
+    test_v2v_update_all()
+    test_v2v_snr()
+    test_v2v_pull()
+    test_v2v_update_all_multi_fn()
+    test_v2v_snr_multi_fn()
+    test_e2v_update_all_multi_fn()
+    test_e2v_snr_multi_fn()
+    test_e2v_recv_multi_fn()
+    test_update_all_multi_fallback()
+    test_pull_multi_fallback()
+    test_spmv_3d_feat()

tests/python/common/test_heterograph-update-all.py

@@ -0,0 +1,372 @@
+import itertools
+import unittest
+from collections import Counter
+from itertools import product
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as ssp
+from dgl import DGLError
+from scipy.sparse import rand
+from utils import get_cases, parametrize_idtype
+
+rfuncs = {"sum": fn.sum, "max": fn.max, "min": fn.min, "mean": fn.mean}
+feat_size = 2
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+def create_test_heterograph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 1, 1], [0, 0, 1]),
+            ("developer", "develops", "game"): ([0, 1, 0], [0, 1, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+def create_test_heterograph_2(idtype):
+    src = np.random.randint(0, 50, 25)
+    dst = np.random.randint(0, 50, 25)
+    src1 = np.random.randint(0, 25, 10)
+    dst1 = np.random.randint(0, 25, 10)
+    src2 = np.random.randint(0, 100, 1000)
+    dst2 = np.random.randint(0, 100, 1000)
+    g = dgl.heterograph(
+        {
+            ("user", "becomes", "player"): (src, dst),
+            ("user", "follows", "user"): (src, dst),
+            ("user", "plays", "game"): (src, dst),
+            ("user", "wishes", "game"): (src1, dst1),
+            ("developer", "develops", "game"): (src2, dst2),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+def create_test_heterograph_large(idtype):
+    src = np.random.randint(0, 50, 2500)
+    dst = np.random.randint(0, 50, 2500)
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): (src, dst),
+            ("user", "plays", "game"): (src, dst),
+            ("user", "wishes", "game"): (src, dst),
+            ("developer", "develops", "game"): (src, dst),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+@parametrize_idtype
+def test_unary_copy_u(idtype):
+    def _test(mfunc, rfunc):
+        g = create_test_heterograph_2(idtype)
+        g0 = create_test_heterograph(idtype)
+        g1 = create_test_heterograph_large(idtype)
+        cross_reducer = rfunc.__name__
+        x1 = F.randn((g.num_nodes("user"), feat_size))
+        x2 = F.randn((g.num_nodes("developer"), feat_size))
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        g.nodes["user"].data["h"] = x1
+        g.nodes["developer"].data["h"] = x2
+
+        #################################################################
+        #  multi_update_all(): call msg_passing separately for each etype
+        #################################################################
+
+        with F.record_grad():
+            g.multi_update_all(
+                {
+                    etype: (mfunc("h", "m"), rfunc("m", "y"))
+                    for etype in g.canonical_etypes
+                },
+                cross_reducer,
+            )
+            r1 = g.nodes["game"].data["y"].clone()
+            r2 = g.nodes["user"].data["y"].clone()
+            r3 = g.nodes["player"].data["y"].clone()
+            loss = r1.sum() + r2.sum() + r3.sum()
+            F.backward(loss)
+            n_grad1 = F.grad(g.nodes["user"].data["h"]).clone()
+            n_grad2 = F.grad(g.nodes["developer"].data["h"]).clone()
+
+        g.nodes["user"].data.clear()
+        g.nodes["developer"].data.clear()
+        g.nodes["game"].data.clear()
+        g.nodes["player"].data.clear()
+
+        #################################################################
+        #  update_all(): call msg_passing for all etypes
+        #################################################################
+
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        g.nodes["user"].data["h"] = x1
+        g.nodes["developer"].data["h"] = x2
+
+        with F.record_grad():
+            g.update_all(mfunc("h", "m"), rfunc("m", "y"))
+            r4 = g.nodes["game"].data["y"]
+            r5 = g.nodes["user"].data["y"]
+            r6 = g.nodes["player"].data["y"]
+            loss = r4.sum() + r5.sum() + r6.sum()
+            F.backward(loss)
+            n_grad3 = F.grad(g.nodes["user"].data["h"])
+            n_grad4 = F.grad(g.nodes["developer"].data["h"])
+
+        assert F.allclose(r1, r4)
+        assert F.allclose(r2, r5)
+        assert F.allclose(r3, r6)
+        assert F.allclose(n_grad1, n_grad3)
+        assert F.allclose(n_grad2, n_grad4)
+
+    _test(fn.copy_u, fn.sum)
+    _test(fn.copy_u, fn.max)
+    _test(fn.copy_u, fn.min)
+    # _test('copy_u', 'mean')
+
+
+@parametrize_idtype
+def test_unary_copy_e(idtype):
+    def _test(mfunc, rfunc):
+        g = create_test_heterograph_large(idtype)
+        g0 = create_test_heterograph_2(idtype)
+        g1 = create_test_heterograph(idtype)
+        cross_reducer = rfunc.__name__
+        x1 = F.randn((g.num_edges("plays"), feat_size))
+        x2 = F.randn((g.num_edges("follows"), feat_size))
+        x3 = F.randn((g.num_edges("develops"), feat_size))
+        x4 = F.randn((g.num_edges("wishes"), feat_size))
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+        g["plays"].edata["eid"] = x1
+        g["follows"].edata["eid"] = x2
+        g["develops"].edata["eid"] = x3
+        g["wishes"].edata["eid"] = x4
+
+        #################################################################
+        #  multi_update_all(): call msg_passing separately for each etype
+        #################################################################
+
+        with F.record_grad():
+            g.multi_update_all(
+                {
+                    "plays": (mfunc("eid", "m"), rfunc("m", "y")),
+                    "follows": (mfunc("eid", "m"), rfunc("m", "y")),
+                    "develops": (mfunc("eid", "m"), rfunc("m", "y")),
+                    "wishes": (mfunc("eid", "m"), rfunc("m", "y")),
+                },
+                cross_reducer,
+            )
+            r1 = g.nodes["game"].data["y"].clone()
+            r2 = g.nodes["user"].data["y"].clone()
+            loss = r1.sum() + r2.sum()
+            F.backward(loss)
+            e_grad1 = F.grad(g["develops"].edata["eid"]).clone()
+            e_grad2 = F.grad(g["plays"].edata["eid"]).clone()
+            e_grad3 = F.grad(g["wishes"].edata["eid"]).clone()
+            e_grad4 = F.grad(g["follows"].edata["eid"]).clone()
+        {etype: (g[etype].edata.clear()) for _, etype, _ in g.canonical_etypes},
+
+        #################################################################
+        #  update_all(): call msg_passing for all etypes
+        #################################################################
+
+        # TODO(Israt): output type can be None in multi_update and empty
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+
+        g["plays"].edata["eid"] = x1
+        g["follows"].edata["eid"] = x2
+        g["develops"].edata["eid"] = x3
+        g["wishes"].edata["eid"] = x4
+
+        with F.record_grad():
+            g.update_all(mfunc("eid", "m"), rfunc("m", "y"))
+            r3 = g.nodes["game"].data["y"]
+            r4 = g.nodes["user"].data["y"]
+            loss = r3.sum() + r4.sum()
+            F.backward(loss)
+            e_grad5 = F.grad(g["develops"].edata["eid"])
+            e_grad6 = F.grad(g["plays"].edata["eid"])
+            e_grad7 = F.grad(g["wishes"].edata["eid"])
+            e_grad8 = F.grad(g["follows"].edata["eid"])
+
+        # # correctness check
+        def _print_error(a, b):
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        assert F.allclose(r1, r3)
+        assert F.allclose(r2, r4)
+        assert F.allclose(e_grad1, e_grad5)
+        assert F.allclose(e_grad2, e_grad6)
+        assert F.allclose(e_grad3, e_grad7)
+        assert F.allclose(e_grad4, e_grad8)
+
+    _test(fn.copy_e, fn.sum)
+    _test(fn.copy_e, fn.max)
+    _test(fn.copy_e, fn.min)
+    # _test('copy_e', 'mean')
+
+
+@parametrize_idtype
+def test_binary_op(idtype):
+    def _test(lhs, rhs, binary_op, reducer):
+        g = create_test_heterograph(idtype)
+
+        x1 = F.randn((g.num_nodes("user"), feat_size))
+        x2 = F.randn((g.num_nodes("developer"), feat_size))
+        x3 = F.randn((g.num_nodes("game"), feat_size))
+
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        g.nodes["user"].data["h"] = x1
+        g.nodes["developer"].data["h"] = x2
+        g.nodes["game"].data["h"] = x3
+
+        x1 = F.randn((4, feat_size))
+        x2 = F.randn((4, feat_size))
+        x3 = F.randn((3, feat_size))
+        x4 = F.randn((3, feat_size))
+        F.attach_grad(x1)
+        F.attach_grad(x2)
+        F.attach_grad(x3)
+        F.attach_grad(x4)
+        g["plays"].edata["h"] = x1
+        g["follows"].edata["h"] = x2
+        g["develops"].edata["h"] = x3
+        g["wishes"].edata["h"] = x4
+
+        builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
+        builtin_msg = getattr(fn, builtin_msg_name)
+        builtin_red = getattr(fn, reducer)
+
+        #################################################################
+        #  multi_update_all(): call msg_passing separately for each etype
+        #################################################################
+
+        with F.record_grad():
+            g.multi_update_all(
+                {
+                    etype: (builtin_msg("h", "h", "m"), builtin_red("m", "y"))
+                    for etype in g.canonical_etypes
+                },
+                "sum",
+            )
+            r1 = g.nodes["game"].data["y"]
+            F.backward(r1, F.ones(r1.shape))
+            n_grad1 = F.grad(r1)
+
+        #################################################################
+        #  update_all(): call msg_passing for all etypes
+        #################################################################
+
+        g.update_all(builtin_msg("h", "h", "m"), builtin_red("m", "y"))
+        r2 = g.nodes["game"].data["y"]
+        F.backward(r2, F.ones(r2.shape))
+        n_grad2 = F.grad(r2)
+
+        # correctness check
+        def _print_error(a, b):
+            for i, (x, y) in enumerate(
+                zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())
+            ):
+                if not np.allclose(x, y):
+                    print("@{} {} v.s. {}".format(i, x, y))
+
+        if not F.allclose(r1, r2):
+            _print_error(r1, r2)
+        assert F.allclose(r1, r2)
+        # TODO (Israt): r1 and r2 have different frad func associated with
+        # if not F.allclose(n_grad1, n_grad2):
+        #     print('node grad')
+        #     _print_error(n_grad1, n_grad2)
+        # assert(F.allclose(n_grad1, n_grad2))
+
+    target = ["u", "v", "e"]
+    for lhs, rhs in product(target, target):
+        if lhs == rhs:
+            continue
+        for binary_op in ["add", "sub", "mul", "div"]:
+            # TODO(Israt) :Add support for reduce func "max", "min", "mean"
+            for reducer in ["sum"]:
+                print(lhs, rhs, binary_op, reducer)
+                _test(lhs, rhs, binary_op, reducer)
+
+
+# Issue #5873
+def test_multi_update_all_minmax_reduce_with_isolated_nodes():
+    g = dgl.heterograph(
+        {
+            ("A", "AB", "B"): ([0, 1, 2, 3], [0, 0, 1, 1]),
+            ("C", "CB", "B"): ([0, 1, 2, 3], [2, 2, 3, 3]),
+        },
+        device=F.ctx(),
+    )
+    g.nodes["A"].data["x"] = F.randn((4, 16))
+    g.nodes["C"].data["x"] = F.randn((4, 16))
+    g.multi_update_all(
+        {
+            "AB": (dgl.function.copy_u("x", "m"), dgl.function.min("m", "a1")),
+            "CB": (dgl.function.copy_u("x", "m"), dgl.function.min("m", "a2")),
+        },
+        cross_reducer="min",
+    )
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a1"])).any()
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a2"])).any()
+
+    g.multi_update_all(
+        {
+            "AB": (dgl.function.copy_u("x", "m"), dgl.function.max("m", "a1")),
+            "CB": (dgl.function.copy_u("x", "m"), dgl.function.max("m", "a2")),
+        },
+        cross_reducer="max",
+    )
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a1"])).any()
+    assert not np.isinf(F.asnumpy(g.nodes["B"].data["a2"])).any()
+
+
+if __name__ == "__main__":
+    test_unary_copy_u()
+    test_unary_copy_e()
+    test_binary_op()

tests/python/common/test_homophily.py

@@ -0,0 +1,68 @@
+import math
+import unittest
+
+import backend as F
+
+import dgl
+from utils import parametrize_idtype
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_node_homophily(idtype):
+    # IfChangeThenChange: python/dgl/homophily.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(
+        ([1, 2, 0, 4], [0, 1, 2, 3]), idtype=idtype, device=device
+    )
+    y = F.tensor([0, 0, 0, 0, 1])
+    assert math.isclose(dgl.node_homophily(graph, y), 0.6000000238418579)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_edge_homophily(idtype):
+    # IfChangeThenChange: python/dgl/homophily.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(
+        ([1, 2, 0, 4], [0, 1, 2, 3]), idtype=idtype, device=device
+    )
+    y = F.tensor([0, 0, 0, 0, 1])
+    assert math.isclose(dgl.edge_homophily(graph, y), 0.75)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_linkx_homophily(idtype):
+    # IfChangeThenChange: python/dgl/homophily.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(([0, 1, 2, 3], [1, 2, 0, 4]), device=device)
+    y = F.tensor([0, 0, 0, 0, 1])
+    assert math.isclose(dgl.linkx_homophily(graph, y), 0.19999998807907104)
+
+    y = F.tensor([0, 1, 2, 3, 4])
+    assert math.isclose(dgl.linkx_homophily(graph, y), 0.0000000000000000)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_adjusted_homophily(idtype):
+    # IfChangeThenChange: python/dgl/homophily.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(
+        ([1, 2, 0, 4], [0, 1, 2, 3]), idtype=idtype, device=device
+    )
+    y = F.tensor([0, 0, 0, 0, 1])
+    assert math.isclose(dgl.adjusted_homophily(graph, y), -0.1428571492433548)

tests/python/common/test_label_informativeness.py

@@ -0,0 +1,45 @@
+import math
+import unittest
+
+import backend as F
+
+import dgl
+from utils import parametrize_idtype
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_edge_label_informativeness(idtype):
+    # IfChangeThenChange: python/dgl/label_informativeness.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(
+        ([0, 1, 2, 2, 3, 4], [1, 2, 0, 3, 4, 5]), idtype=idtype, device=device
+    )
+    y = F.tensor([0, 0, 0, 0, 1, 1])
+    assert math.isclose(
+        dgl.edge_label_informativeness(graph, y),
+        0.25177597999572754,
+        abs_tol=1e-6,
+    )
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="Only support PyTorch for now"
+)
+@parametrize_idtype
+def test_node_label_informativeness(idtype):
+    # IfChangeThenChange: python/dgl/label_informativeness.py
+    # Update the docstring example.
+    device = F.ctx()
+    graph = dgl.graph(
+        ([0, 1, 2, 2, 3, 4], [1, 2, 0, 3, 4, 5]), idtype=idtype, device=device
+    )
+    y = F.tensor([0, 0, 0, 0, 1, 1])
+    assert math.isclose(
+        dgl.node_label_informativeness(graph, y),
+        0.3381872773170471,
+        abs_tol=1e-6,
+    )

tests/python/common/test_merge.py

@@ -0,0 +1,55 @@
+import backend as F
+
+import dgl
+from utils import parametrize_idtype
+
+
+@parametrize_idtype
+def test_heterograph_merge(idtype):
+    g1 = (
+        dgl.heterograph({("a", "to", "b"): ([0, 1], [1, 0])})
+        .astype(idtype)
+        .to(F.ctx())
+    )
+    g1_n_edges = g1.num_edges(etype="to")
+    g1.nodes["a"].data["nh"] = F.randn((2, 3))
+    g1.nodes["b"].data["nh"] = F.randn((2, 3))
+    g1.edges["to"].data["eh"] = F.randn((2, 3))
+
+    g2 = (
+        dgl.heterograph({("a", "to", "b"): ([1, 2, 3], [2, 3, 5])})
+        .astype(idtype)
+        .to(F.ctx())
+    )
+    g2.nodes["a"].data["nh"] = F.randn((4, 3))
+    g2.nodes["b"].data["nh"] = F.randn((6, 3))
+    g2.edges["to"].data["eh"] = F.randn((3, 3))
+    g2.add_nodes(3, ntype="a")
+    g2.add_nodes(3, ntype="b")
+
+    m = dgl.merge([g1, g2])
+
+    # Check g2's edges and nodes were added to g1's in m.
+    m_us = F.asnumpy(m.edges()[0][g1_n_edges:])
+    g2_us = F.asnumpy(g2.edges()[0])
+    assert all(m_us == g2_us)
+    m_vs = F.asnumpy(m.edges()[1][g1_n_edges:])
+    g2_vs = F.asnumpy(g2.edges()[1])
+    assert all(m_vs == g2_vs)
+    for ntype in m.ntypes:
+        assert m.num_nodes(ntype=ntype) == max(
+            g1.num_nodes(ntype=ntype), g2.num_nodes(ntype=ntype)
+        )
+
+        # Check g1's node data was updated with g2's in m.
+        for key in m.nodes[ntype].data:
+            g2_n_nodes = g2.num_nodes(ntype=ntype)
+            updated_g1_ndata = F.asnumpy(m.nodes[ntype].data[key][:g2_n_nodes])
+            g2_ndata = F.asnumpy(g2.nodes[ntype].data[key])
+            assert all((updated_g1_ndata == g2_ndata).flatten())
+
+    # Check g1's edge data was updated with g2's in m.
+    for key in m.edges["to"].data:
+        updated_g1_edata = F.asnumpy(m.edges["to"].data[key][g1_n_edges:])
+        g2_edata = F.asnumpy(g2.edges["to"].data[key])
+        assert all((updated_g1_edata == g2_edata).flatten())

tests/python/common/test_partition.py

@@ -0,0 +1,55 @@
+import unittest
+
+import backend as F
+
+from dgl.distributed import graph_partition_book as gpb
+from dgl.partition import NDArrayPartition
+from utils import parametrize_idtype
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="NDArrayPartition only works on GPU.",
+)
+@parametrize_idtype
+def test_get_node_partition_from_book(idtype):
+    node_map = {"_N": F.tensor([[0, 3], [4, 5], [6, 10]], dtype=idtype)}
+    edge_map = {
+        ("_N", "_E", "_N"): F.tensor([[0, 9], [10, 15], [16, 25]], dtype=idtype)
+    }
+    ntypes = {ntype: i for i, ntype in enumerate(node_map)}
+    etypes = {etype: i for i, etype in enumerate(edge_map)}
+    book = gpb.RangePartitionBook(0, 3, node_map, edge_map, ntypes, etypes)
+    partition = gpb.get_node_partition_from_book(book, F.ctx())
+    assert partition.num_parts() == 3
+    assert partition.array_size() == 11
+
+    # Test map_to_local
+    test_ids = F.copy_to(F.tensor([0, 2, 6, 7, 10], dtype=idtype), F.ctx())
+    act_ids = partition.map_to_local(test_ids)
+    exp_ids = F.copy_to(F.tensor([0, 2, 0, 1, 4], dtype=idtype), F.ctx())
+    assert F.array_equal(act_ids, exp_ids)
+
+    # Test map_to_global
+    test_ids = F.copy_to(F.tensor([0, 2], dtype=idtype), F.ctx())
+    act_ids = partition.map_to_global(test_ids, 0)
+    exp_ids = F.copy_to(F.tensor([0, 2], dtype=idtype), F.ctx())
+    assert F.array_equal(act_ids, exp_ids)
+
+    test_ids = F.copy_to(F.tensor([0, 1], dtype=idtype), F.ctx())
+    act_ids = partition.map_to_global(test_ids, 1)
+    exp_ids = F.copy_to(F.tensor([4, 5], dtype=idtype), F.ctx())
+    assert F.array_equal(act_ids, exp_ids)
+
+    test_ids = F.copy_to(F.tensor([0, 1, 4], dtype=idtype), F.ctx())
+    act_ids = partition.map_to_global(test_ids, 2)
+    exp_ids = F.copy_to(F.tensor([6, 7, 10], dtype=idtype), F.ctx())
+    assert F.array_equal(act_ids, exp_ids)
+
+    # Test generate_permutation
+    test_ids = F.copy_to(F.tensor([6, 0, 7, 2, 10], dtype=idtype), F.ctx())
+    perm, split_sum = partition.generate_permutation(test_ids)
+    exp_perm = F.copy_to(F.tensor([1, 3, 0, 2, 4], dtype=idtype), F.ctx())
+    exp_sum = F.copy_to(F.tensor([2, 0, 3]), F.ctx())
+    assert F.array_equal(perm, exp_perm)
+    assert F.array_equal(split_sum, exp_sum)

tests/python/common/test_propagate.py

@@ -0,0 +1,121 @@
+import unittest
+
+import backend as F
+
+import dgl
+import networkx as nx
+from utils import check_fail, parametrize_idtype
+
+
+def create_graph(idtype):
+    g = dgl.from_networkx(nx.path_graph(5), idtype=idtype, device=F.ctx())
+    return g
+
+
+def mfunc(edges):
+    return {"m": edges.src["x"]}
+
+
+def rfunc(nodes):
+    msg = F.sum(nodes.mailbox["m"], 1)
+    return {"x": nodes.data["x"] + msg}
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@parametrize_idtype
+def test_prop_nodes_bfs(idtype):
+    g = create_graph(idtype)
+    g.ndata["x"] = F.ones((5, 2))
+    dgl.prop_nodes_bfs(
+        g, 0, message_func=mfunc, reduce_func=rfunc, apply_node_func=None
+    )
+    # pull nodes using bfs order will result in a cumsum[i] + data[i] + data[i+1]
+    assert F.allclose(
+        g.ndata["x"],
+        F.tensor([[2.0, 2.0], [4.0, 4.0], [6.0, 6.0], [8.0, 8.0], [9.0, 9.0]]),
+    )
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@parametrize_idtype
+def test_prop_edges_dfs(idtype):
+    g = create_graph(idtype)
+    g.ndata["x"] = F.ones((5, 2))
+    dgl.prop_edges_dfs(
+        g, 0, message_func=mfunc, reduce_func=rfunc, apply_node_func=None
+    )
+    # snr using dfs results in a cumsum
+    assert F.allclose(
+        g.ndata["x"],
+        F.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0], [5.0, 5.0]]),
+    )
+
+    g.ndata["x"] = F.ones((5, 2))
+    dgl.prop_edges_dfs(
+        g,
+        0,
+        has_reverse_edge=True,
+        message_func=mfunc,
+        reduce_func=rfunc,
+        apply_node_func=None,
+    )
+    # result is cumsum[i] + cumsum[i-1]
+    assert F.allclose(
+        g.ndata["x"],
+        F.tensor([[1.0, 1.0], [3.0, 3.0], [5.0, 5.0], [7.0, 7.0], [9.0, 9.0]]),
+    )
+
+    g.ndata["x"] = F.ones((5, 2))
+    dgl.prop_edges_dfs(
+        g,
+        0,
+        has_nontree_edge=True,
+        message_func=mfunc,
+        reduce_func=rfunc,
+        apply_node_func=None,
+    )
+    # result is cumsum[i] + cumsum[i+1]
+    assert F.allclose(
+        g.ndata["x"],
+        F.tensor([[3.0, 3.0], [5.0, 5.0], [7.0, 7.0], [9.0, 9.0], [5.0, 5.0]]),
+    )
+
+
+@unittest.skipIf(F._default_context_str == "gpu", reason="GPU not implemented")
+@parametrize_idtype
+def test_prop_nodes_topo(idtype):
+    # bi-directional chain
+    g = create_graph(idtype)
+    assert check_fail(dgl.prop_nodes_topo, g)  # has loop
+
+    # tree
+    tree = dgl.graph([])
+    tree.add_nodes(5)
+    tree.add_edges(1, 0)
+    tree.add_edges(2, 0)
+    tree.add_edges(3, 2)
+    tree.add_edges(4, 2)
+    tree = dgl.graph(tree.edges())
+    # init node feature data
+    tree.ndata["x"] = F.zeros((5, 2))
+    # set all leaf nodes to be ones
+    tree.nodes[[1, 3, 4]].data["x"] = F.ones((3, 2))
+
+    # Filtering DGLWarning:
+    #    The input graph for the user-defined edge
+    #    function does not contain valid edges
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        dgl.prop_nodes_topo(
+            tree, message_func=mfunc, reduce_func=rfunc, apply_node_func=None
+        )
+    # root node get the sum
+    assert F.allclose(tree.nodes[0].data["x"], F.tensor([[3.0, 3.0]]))
+
+
+if __name__ == "__main__":
+    test_prop_nodes_bfs()
+    test_prop_edges_dfs()
+    test_prop_nodes_topo()

tests/python/common/test_random.py

@@ -0,0 +1,47 @@
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu", reason="GPU random choice not implemented"
+)
+def test_random_choice():
+    # test 1
+    a = F.arange(0, 100)
+    x = dgl.random.choice(a, 10, replace=True, prob=None)
+    assert len(x) == 10
+    for i in range(len(x)):
+        assert F.asnumpy(x[i]) >= 0 and F.asnumpy(x[i]) < 100
+    # test 2, replace=False, small num
+    a = F.arange(0, 100)
+    x = dgl.random.choice(a, 10, replace=False, prob=None)
+    assert len(x) == 10
+    for i in range(len(x)):
+        assert F.asnumpy(x[i]) >= 0 and F.asnumpy(x[i]) < 100
+    # test 3, replace=False, large num
+    a = F.arange(0, 100)
+    x = dgl.random.choice(a, 100, replace=False, prob=None)
+    assert len(x) == 100
+    assert np.array_equal(np.sort(F.asnumpy(x)), F.asnumpy(a))
+    # test 4, first arg is integer
+    x = dgl.random.choice(100, 100, replace=False, prob=None)
+    assert len(x) == 100
+    assert np.array_equal(np.sort(F.asnumpy(x)), F.asnumpy(a))
+    # test 5, with prob
+    prob = np.ones((100,))
+    prob[37:40] = 0.0
+    prob -= prob.min()
+    prob /= prob.sum()
+    prob = F.tensor(prob)
+    x = dgl.random.choice(100, 97, replace=False, prob=prob)
+    assert len(x) == 97
+    for i in range(len(x)):
+        assert F.asnumpy(x[i]) < 37 or F.asnumpy(x[i]) >= 40
+
+
+if __name__ == "__main__":
+    test_random_choice()

tests/python/common/test_readout.py

@@ -0,0 +1,237 @@
+import unittest
+
+import backend as F
+
+import dgl
+import networkx as nx
+import numpy as np
+import pytest
+from utils import parametrize_idtype
+from utils.graph_cases import get_cases
+
+
+@parametrize_idtype
+def test_sum_case1(idtype):
+    # NOTE: If you want to update this test case, remember to update the docstring
+    #  example too!!!
+    g1 = dgl.graph(([0, 1], [1, 0]), idtype=idtype, device=F.ctx())
+    g1.ndata["h"] = F.tensor([1.0, 2.0])
+    g2 = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
+    g2.ndata["h"] = F.tensor([1.0, 2.0, 3.0])
+    bg = dgl.batch([g1, g2])
+    bg.ndata["w"] = F.tensor([0.1, 0.2, 0.1, 0.5, 0.2])
+    assert F.allclose(F.tensor([3.0]), dgl.sum_nodes(g1, "h"))
+    assert F.allclose(F.tensor([3.0, 6.0]), dgl.sum_nodes(bg, "h"))
+    assert F.allclose(F.tensor([0.5, 1.7]), dgl.sum_nodes(bg, "h", "w"))
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+@pytest.mark.parametrize("reducer", ["sum", "max", "mean"])
+def test_reduce_readout(g, idtype, reducer):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata["h"] = F.randn((g.num_nodes(), 3))
+    g.edata["h"] = F.randn((g.num_edges(), 2))
+
+    # Test.1: node readout
+    x = dgl.readout_nodes(g, "h", op=reducer)
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = dgl.readout_nodes(sg, "h", op=reducer)
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    x = getattr(dgl, "{}_nodes".format(reducer))(g, "h")
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = getattr(dgl, "{}_nodes".format(reducer))(sg, "h")
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    # Test.2: edge readout
+    x = dgl.readout_edges(g, "h", op=reducer)
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = dgl.readout_edges(sg, "h", op=reducer)
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    x = getattr(dgl, "{}_edges".format(reducer))(g, "h")
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = getattr(dgl, "{}_edges".format(reducer))(sg, "h")
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+@pytest.mark.parametrize("reducer", ["sum", "max", "mean"])
+def test_weighted_reduce_readout(g, idtype, reducer):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata["h"] = F.randn((g.num_nodes(), 3))
+    g.ndata["w"] = F.randn((g.num_nodes(), 1))
+    g.edata["h"] = F.randn((g.num_edges(), 2))
+    g.edata["w"] = F.randn((g.num_edges(), 1))
+
+    # Test.1: node readout
+    x = dgl.readout_nodes(g, "h", "w", op=reducer)
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = dgl.readout_nodes(sg, "h", "w", op=reducer)
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    x = getattr(dgl, "{}_nodes".format(reducer))(g, "h", "w")
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = getattr(dgl, "{}_nodes".format(reducer))(sg, "h", "w")
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    # Test.2: edge readout
+    x = dgl.readout_edges(g, "h", "w", op=reducer)
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = dgl.readout_edges(sg, "h", "w", op=reducer)
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    x = getattr(dgl, "{}_edges".format(reducer))(g, "h", "w")
+    # check correctness
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        sx = getattr(dgl, "{}_edges".format(reducer))(sg, "h", "w")
+        subx.append(sx)
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+@pytest.mark.parametrize("descending", [True, False])
+def test_topk(g, idtype, descending):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata["x"] = F.randn((g.num_nodes(), 3))
+
+    # Test.1: to test the case where k > number of nodes.
+    dgl.topk_nodes(g, "x", 100, sortby=-1)
+
+    # Test.2: test correctness
+    min_nnodes = F.asnumpy(g.batch_num_nodes()).min()
+    if min_nnodes <= 1:
+        return
+    k = min_nnodes - 1
+    val, indices = dgl.topk_nodes(g, "x", k, descending=descending, sortby=-1)
+    print(k)
+    print(g.ndata["x"])
+    print("val", val)
+    print("indices", indices)
+    subg = dgl.unbatch(g)
+    subval, subidx = [], []
+    for sg in subg:
+        subx = F.asnumpy(sg.ndata["x"])
+        ai = np.argsort(subx[:, -1:].flatten())
+        if descending:
+            ai = np.ascontiguousarray(ai[::-1])
+        subx = np.expand_dims(subx[ai[:k]], 0)
+        subval.append(F.tensor(subx))
+        subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
+    print(F.cat(subval, dim=0))
+    assert F.allclose(val, F.cat(subval, dim=0))
+    assert F.allclose(indices, F.cat(subidx, dim=0))
+
+    # Test.3: sorby=None
+    dgl.topk_nodes(g, "x", k, sortby=None)
+
+    g.edata["x"] = F.randn((g.num_edges(), 3))
+
+    # Test.4: topk edges where k > number of edges.
+    dgl.topk_edges(g, "x", 100, sortby=-1)
+
+    # Test.5: topk edges test correctness
+    min_nedges = F.asnumpy(g.batch_num_edges()).min()
+    if min_nedges <= 1:
+        return
+    k = min_nedges - 1
+    val, indices = dgl.topk_edges(g, "x", k, descending=descending, sortby=-1)
+    print(k)
+    print(g.edata["x"])
+    print("val", val)
+    print("indices", indices)
+    subg = dgl.unbatch(g)
+    subval, subidx = [], []
+    for sg in subg:
+        subx = F.asnumpy(sg.edata["x"])
+        ai = np.argsort(subx[:, -1:].flatten())
+        if descending:
+            ai = np.ascontiguousarray(ai[::-1])
+        subx = np.expand_dims(subx[ai[:k]], 0)
+        subval.append(F.tensor(subx))
+        subidx.append(F.tensor(np.expand_dims(ai[:k], 0)))
+    print(F.cat(subval, dim=0))
+    assert F.allclose(val, F.cat(subval, dim=0))
+    assert F.allclose(indices, F.cat(subidx, dim=0))
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+def test_softmax(g, idtype):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata["h"] = F.randn((g.num_nodes(), 3))
+    g.edata["h"] = F.randn((g.num_edges(), 2))
+
+    # Test.1: node readout
+    x = dgl.softmax_nodes(g, "h")
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        subx.append(F.softmax(sg.ndata["h"], dim=0))
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+    # Test.2: edge readout
+    x = dgl.softmax_edges(g, "h")
+    subg = dgl.unbatch(g)
+    subx = []
+    for sg in subg:
+        subx.append(F.softmax(sg.edata["h"], dim=0))
+    assert F.allclose(x, F.cat(subx, dim=0))
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
+def test_broadcast(idtype, g):
+    g = g.astype(idtype).to(F.ctx())
+    gfeat = F.randn((g.batch_size, 3))
+
+    # Test.0: broadcast_nodes
+    g.ndata["h"] = dgl.broadcast_nodes(g, gfeat)
+    subg = dgl.unbatch(g)
+    for i, sg in enumerate(subg):
+        assert F.allclose(
+            sg.ndata["h"],
+            F.repeat(F.reshape(gfeat[i], (1, 3)), sg.num_nodes(), dim=0),
+        )
+
+    # Test.1: broadcast_edges
+    g.edata["h"] = dgl.broadcast_edges(g, gfeat)
+    subg = dgl.unbatch(g)
+    for i, sg in enumerate(subg):
+        assert F.allclose(
+            sg.edata["h"],
+            F.repeat(F.reshape(gfeat[i], (1, 3)), sg.num_edges(), dim=0),
+        )

tests/python/common/test_sparse_ops-csr.py

@@ -0,0 +1,268 @@
+import backend as F
+
+import dgl
+import numpy as np
+import pytest
+import scipy.sparse as ssp
+from utils import parametrize_idtype
+
+if F.backend_name == "pytorch":
+    import torch
+
+    torch.backends.cuda.matmul.allow_tf32 = False
+
+
+def _random_simple_graph(
+    idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype
+):
+    src = np.random.randint(0, M, (max_nnz,))
+    dst = np.random.randint(0, N, (max_nnz,))
+    val = np.random.randn(max_nnz)
+    a = ssp.csr_matrix((val, (src, dst)), shape=(M, N))
+    a.sum_duplicates()
+    a = a.tocoo()
+    # shuffle edges
+    perm = np.random.permutation(a.nnz)
+    row = a.row[perm]
+    col = a.col[perm]
+    val = a.data[perm]
+    a = ssp.csr_matrix((val, (row, col)), shape=(M, N))
+
+    A = dgl.heterograph(
+        {
+            (srctype, etype, dsttype): (
+                F.copy_to(F.tensor(row, dtype=idtype), ctx),
+                F.copy_to(F.tensor(col, dtype=idtype), ctx),
+            )
+        },
+        num_nodes_dict={srctype: a.shape[0], dsttype: a.shape[1]},
+    )
+    A.edata["w"] = F.copy_to(F.tensor(val, dtype=dtype), ctx)
+    return a, A
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+@pytest.mark.parametrize("return_edge_ids", [True, False])
+def test_csrmm(idtype, dtype, return_edge_ids):
+    a, A = _random_simple_graph(
+        idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
+    )
+    b, B = _random_simple_graph(
+        idtype, dtype, F.ctx(), 600, 700, 9000, "B", "C", "BC"
+    )
+    C, C_weights = dgl._sparse_ops._csrmm(
+        A._graph, A.edata["w"], B._graph, B.edata["w"], 2
+    )
+    C_adj = C.adjacency_matrix_scipy(0, False, "csr", return_edge_ids)
+    C_adj.data = F.asnumpy(C_weights)
+    C_adj = F.tensor(C_adj.todense(), dtype=dtype)
+    c = F.tensor((a * b).todense(), dtype=dtype)
+    assert F.allclose(C_adj, c)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+@pytest.mark.parametrize("num_vtypes", [1, 2])
+def test_csrmm_backward(idtype, dtype, num_vtypes):
+    a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
+    b, B = _random_simple_graph(
+        idtype,
+        dtype,
+        F.ctx(),
+        4,
+        3,
+        6,
+        "B",
+        "A" if num_vtypes == 1 else "C",
+        "BA",
+    )
+    A_row, A_col = A.edges(order="eid")
+    B_row, B_col = B.edges(order="eid")
+    A_row = F.asnumpy(A_row)
+    A_col = F.asnumpy(A_col)
+    B_row = F.asnumpy(B_row)
+    B_col = F.asnumpy(B_col)
+    a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
+    b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
+
+    A.edata["w"] = F.attach_grad(A.edata["w"])
+    B.edata["w"] = F.attach_grad(B.edata["w"])
+
+    with F.record_grad():
+        C = dgl.adj_product_graph(A, B, "w")
+        assert len(C.ntypes) == num_vtypes
+        assert len(C.etypes) == 1
+        C_dense = np.zeros((3, 3))
+        C_row, C_col = C.edges(order="eid")
+        C_row = F.asnumpy(C_row)
+        C_col = F.asnumpy(C_col)
+        C_dense[C_row, C_col] = F.asnumpy(C.edata["w"])
+        c_dense = F.matmul(a_dense, b_dense)
+        assert np.allclose(C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4)
+
+        F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense))
+        a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
+        b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
+        A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
+        B_spspmm_grad = F.asnumpy(F.grad(B.edata["w"]))
+        assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
+        assert np.allclose(b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+@pytest.mark.parametrize("return_edge_ids", [True, False])
+def test_csrsum(idtype, dtype, return_edge_ids):
+    a, A = _random_simple_graph(
+        idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
+    )
+    b, B = _random_simple_graph(
+        idtype, dtype, F.ctx(), 500, 600, 9000, "A", "B", "AB"
+    )
+    C, C_weights = dgl._sparse_ops._csrsum(
+        [A._graph, B._graph], [A.edata["w"], B.edata["w"]]
+    )
+    C_adj = C.adjacency_matrix_scipy(0, False, "csr", return_edge_ids)
+    C_adj.data = F.asnumpy(C_weights)
+    C_adj = F.tensor(C_adj.todense(), dtype=dtype)
+    c = F.tensor((a + b).todense(), dtype=dtype)
+    assert F.allclose(C_adj, c)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+@pytest.mark.parametrize("nelems", [1, 2])
+def test_csrsum_backward(idtype, dtype, nelems):
+    a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
+    b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
+    A_row, A_col = A.edges(order="eid")
+    B_row, B_col = B.edges(order="eid")
+    A_row = F.asnumpy(A_row)
+    A_col = F.asnumpy(A_col)
+    B_row = F.asnumpy(B_row)
+    B_col = F.asnumpy(B_col)
+    a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
+    b_dense = F.attach_grad(F.tensor(b.todense(), dtype=dtype))
+
+    A.edata["w"] = F.attach_grad(A.edata["w"])
+    B.edata["w"] = F.attach_grad(B.edata["w"])
+
+    with F.record_grad():
+        if nelems == 2:
+            # Test for two element case
+            C = dgl.adj_sum_graph([A, B], "w")
+            assert C.canonical_etypes == A.canonical_etypes
+            C_dense = np.zeros((3, 4))
+            C_row, C_col = C.edges(order="eid")
+            C_row = F.asnumpy(C_row)
+            C_col = F.asnumpy(C_col)
+            C_dense[C_row, C_col] = F.asnumpy(C.edata["w"])
+            c_dense = a_dense + b_dense
+            assert np.allclose(
+                C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4
+            )
+
+            F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense))
+            a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
+            b_dense_grad = F.asnumpy(F.grad(b_dense))[B_row, B_col]
+            A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
+            B_spspmm_grad = F.asnumpy(F.grad(B.edata["w"]))
+            assert np.allclose(
+                a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4
+            )
+            assert np.allclose(
+                b_dense_grad, B_spspmm_grad, rtol=1e-4, atol=1e-4
+            )
+        elif nelems == 1:
+            # Test for single element case
+            C = dgl.adj_sum_graph([A], "w")
+            assert C.canonical_etypes == A.canonical_etypes
+            C_dense = np.zeros((3, 4))
+            C_row, C_col = C.edges(order="eid")
+            C_row = F.asnumpy(C_row)
+            C_col = F.asnumpy(C_col)
+            C_dense[C_row, C_col] = F.asnumpy(C.edata["w"])
+            c_dense = a_dense
+            assert np.allclose(
+                C_dense, F.asnumpy(c_dense), rtol=1e-4, atol=1e-4
+            )
+
+            F.backward(F.reduce_sum(C.edata["w"]) + F.reduce_sum(c_dense))
+            a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
+            A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
+            assert np.allclose(
+                a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4
+            )
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+@pytest.mark.parametrize("A_nnz", [9000, 0])
+@pytest.mark.parametrize("B_nnz", [9000, 0])
+def test_csrmask(idtype, dtype, A_nnz, B_nnz):
+    a, A = _random_simple_graph(
+        idtype, dtype, F.ctx(), 500, 600, A_nnz, "A", "B", "AB"
+    )
+    b, B = _random_simple_graph(
+        idtype, dtype, F.ctx(), 500, 600, B_nnz, "A", "B", "AB"
+    )
+    C = dgl._sparse_ops._csrmask(A._graph, A.edata["w"], B._graph)
+    B_row, B_col = B.edges(order="eid")
+    B_row = F.asnumpy(B_row)
+    B_col = F.asnumpy(B_col)
+    c = F.tensor(a.todense()[B_row, B_col], dtype)
+    assert F.allclose(C, c)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("dtype", [F.float32, F.float64])
+def test_csrmask_backward(idtype, dtype):
+    a, A = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
+    b, B = _random_simple_graph(idtype, dtype, F.ctx(), 3, 4, 6, "A", "B", "AB")
+    A_row, A_col = A.edges(order="eid")
+    B_row, B_col = B.edges(order="eid")
+    A_row = F.asnumpy(A_row)
+    A_col = F.asnumpy(A_col)
+    B_row = F.asnumpy(B_row)
+    B_col = F.asnumpy(B_col)
+    a_dense = F.attach_grad(F.tensor(a.todense(), dtype=dtype))
+
+    A.edata["w"] = F.attach_grad(A.edata["w"])
+
+    with F.record_grad():
+        # Test for two element case
+        C1 = F.csrmask(A._graph, A.edata["w"], B._graph)
+        if dgl.backend.backend_name == "tensorflow":
+            import tensorflow as tf
+
+            C2 = tf.gather_nd(a_dense, tf.stack([B_row, B_col], 1))
+        else:
+            C2 = a_dense[B_row, B_col]
+        assert F.allclose(C1, C2, rtol=1e-4, atol=1e-4)
+
+        F.backward(F.reduce_sum(C1) + F.reduce_sum(C2))
+        a_dense_grad = F.asnumpy(F.grad(a_dense))[A_row, A_col]
+        A_spspmm_grad = F.asnumpy(F.grad(A.edata["w"]))
+        assert np.allclose(a_dense_grad, A_spspmm_grad, rtol=1e-4, atol=1e-4)
+
+
+if __name__ == "__main__":
+    test_csrmm(F.int32, F.float32)
+    test_csrmm(F.int64, F.float32)
+    test_csrsum(F.int32, F.float32)
+    test_csrsum(F.int64, F.float32)
+    test_csrmask(F.int32, F.float32, 9000, 9000)
+    test_csrmask(F.int64, F.float32, 9000, 0)
+    test_csrmask(F.int32, F.float32, 0, 9000)
+    test_csrmask(F.int64, F.float32, 0, 0)
+    test_csrmm_backward(F.int32, F.float32, 1)
+    test_csrmm_backward(F.int64, F.float32, 1)
+    test_csrmm_backward(F.int32, F.float32, 2)
+    test_csrmm_backward(F.int64, F.float32, 2)
+    test_csrsum_backward(F.int32, F.float32, 1)
+    test_csrsum_backward(F.int64, F.float32, 1)
+    test_csrsum_backward(F.int32, F.float32, 2)
+    test_csrsum_backward(F.int64, F.float32, 2)
+    test_csrmask_backward(F.int32, F.float32)
+    test_csrmask_backward(F.int64, F.float32)

tests/python/common/test_subgraph.py

@@ -0,0 +1,858 @@
+import unittest
+
+import backend as F
+
+import dgl
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as ssp
+from utils import parametrize_idtype
+
+D = 5
+
+
+def generate_graph(grad=False, add_data=True):
+    g = dgl.graph([]).to(F.ctx())
+    g.add_nodes(10)
+    # create a graph where 0 is the source and 9 is the sink
+    for i in range(1, 9):
+        g.add_edges(0, i)
+        g.add_edges(i, 9)
+    # add a back flow from 9 to 0
+    g.add_edges(9, 0)
+    if add_data:
+        ncol = F.randn((10, D))
+        ecol = F.randn((17, D))
+        if grad:
+            ncol = F.attach_grad(ncol)
+            ecol = F.attach_grad(ecol)
+        g.ndata["h"] = ncol
+        g.edata["l"] = ecol
+    return g
+
+
+def test_edge_subgraph():
+    # Test when the graph has no node data and edge data.
+    g = generate_graph(add_data=False)
+    eid = [0, 2, 3, 6, 7, 9]
+
+    # relabel=True
+    sg = g.edge_subgraph(eid)
+    assert F.array_equal(
+        sg.ndata[dgl.NID], F.tensor([0, 2, 4, 5, 1, 9], g.idtype)
+    )
+    assert F.array_equal(sg.edata[dgl.EID], F.tensor(eid, g.idtype))
+    sg.ndata["h"] = F.arange(0, sg.num_nodes())
+    sg.edata["h"] = F.arange(0, sg.num_edges())
+
+    # relabel=False
+    sg = g.edge_subgraph(eid, relabel_nodes=False)
+    assert g.num_nodes() == sg.num_nodes()
+    assert F.array_equal(sg.edata[dgl.EID], F.tensor(eid, g.idtype))
+    sg.ndata["h"] = F.arange(0, sg.num_nodes())
+    sg.edata["h"] = F.arange(0, sg.num_edges())
+
+
+@pytest.mark.parametrize("relabel_nodes", [True, False])
+def test_subgraph_relabel_nodes(relabel_nodes):
+    g = generate_graph()
+    h = g.ndata["h"]
+    l = g.edata["l"]
+    nid = [0, 2, 3, 6, 7, 9]
+    sg = g.subgraph(nid, relabel_nodes=relabel_nodes)
+    eid = {2, 3, 4, 5, 10, 11, 12, 13, 16}
+    assert set(F.asnumpy(sg.edata[dgl.EID])) == eid
+    eid = sg.edata[dgl.EID]
+    # the subgraph is empty initially except for EID field
+    # the subgraph is empty initially except for NID field if relabel_nodes
+    if relabel_nodes:
+        assert len(sg.ndata) == 2
+    assert len(sg.edata) == 2
+    sh = sg.ndata["h"]
+    # The node number is not reduced if relabel_node=False.
+    # The subgraph keeps the same node information as the original graph.
+    if relabel_nodes:
+        assert F.allclose(F.gather_row(h, F.tensor(nid)), sh)
+    else:
+        assert F.allclose(
+            F.gather_row(h, F.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])), sh
+        )
+    # The s,d,eid means the source node, destination node and edge id of the subgraph.
+    # The edges labeled 1 are those selected by the subgraph.
+    """
+    s, d, eid
+    0, 1, 0
+    1, 9, 1
+    0, 2, 2    1
+    2, 9, 3    1
+    0, 3, 4    1
+    3, 9, 5    1
+    0, 4, 6
+    4, 9, 7
+    0, 5, 8
+    5, 9, 9       3
+    0, 6, 10   1
+    6, 9, 11   1  3
+    0, 7, 12   1
+    7, 9, 13   1  3
+    0, 8, 14
+    8, 9, 15      3
+    9, 0, 16   1
+    """
+    assert F.allclose(F.gather_row(l, eid), sg.edata["l"])
+    # update the node/edge features on the subgraph should NOT
+    # reflect to the parent graph.
+    if relabel_nodes:
+        sg.ndata["h"] = F.zeros((6, D))
+    else:
+        sg.ndata["h"] = F.zeros((10, D))
+    assert F.allclose(h, g.ndata["h"])
+
+
+def _test_map_to_subgraph():
+    g = dgl.graph([])
+    g.add_nodes(10)
+    g.add_edges(F.arange(0, 9), F.arange(1, 10))
+    h = g.subgraph([0, 1, 2, 5, 8])
+    v = h.map_to_subgraph_nid([0, 8, 2])
+    assert np.array_equal(F.asnumpy(v), np.array([0, 4, 2]))
+
+
+def create_test_heterograph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([0, 1], [0, 1]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    for etype in g.etypes:
+        g.edges[etype].data["weight"] = F.randn((g.num_edges(etype),))
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+def create_test_heterograph2(idtype):
+    """test heterograph from the docstring, with an empty relation"""
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')
+
+    g = dgl.heterograph(
+        {
+            ("user", "follows", "user"): ([0, 1], [1, 2]),
+            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
+            ("user", "wishes", "game"): ([0, 2], [1, 0]),
+            ("developer", "develops", "game"): ([], []),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    for etype in g.etypes:
+        g.edges[etype].data["weight"] = F.randn((g.num_edges(etype),))
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name == "mxnet",
+    reason="MXNet doesn't support bool tensor",
+)
+@parametrize_idtype
+def test_subgraph_mask(idtype):
+    g = create_test_heterograph(idtype)
+    g_graph = g["follows"]
+    g_bipartite = g["plays"]
+
+    x = F.randn((3, 5))
+    y = F.randn((2, 4))
+    g.nodes["user"].data["h"] = x
+    g.edges["follows"].data["h"] = y
+
+    def _check_subgraph(g, sg):
+        assert sg.idtype == g.idtype
+        assert sg.device == g.device
+        assert sg.ntypes == g.ntypes
+        assert sg.etypes == g.etypes
+        assert sg.canonical_etypes == g.canonical_etypes
+        assert F.array_equal(
+            F.tensor(sg.nodes["user"].data[dgl.NID]), F.tensor([1, 2], idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.nodes["game"].data[dgl.NID]), F.tensor([0], idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["plays"].data[dgl.EID]), F.tensor([1], idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["wishes"].data[dgl.EID]), F.tensor([1], idtype)
+        )
+        assert sg.num_nodes("developer") == 0
+        assert sg.num_edges("develops") == 0
+        assert F.array_equal(
+            sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
+        )
+        assert F.array_equal(
+            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
+        )
+
+    sg1 = g.subgraph(
+        {
+            "user": F.tensor([False, True, True], dtype=F.bool),
+            "game": F.tensor([True, False, False, False], dtype=F.bool),
+        }
+    )
+    _check_subgraph(g, sg1)
+    sg2 = g.edge_subgraph(
+        {
+            "follows": F.tensor([False, True], dtype=F.bool),
+            "plays": F.tensor([False, True, False, False], dtype=F.bool),
+            "wishes": F.tensor([False, True], dtype=F.bool),
+        }
+    )
+    _check_subgraph(g, sg2)
+
+
+@parametrize_idtype
+def test_subgraph1(idtype):
+    g = create_test_heterograph(idtype)
+    g_graph = g["follows"]
+    g_bipartite = g["plays"]
+
+    x = F.randn((3, 5))
+    y = F.randn((2, 4))
+    g.nodes["user"].data["h"] = x
+    g.edges["follows"].data["h"] = y
+
+    def _check_subgraph(g, sg):
+        assert sg.idtype == g.idtype
+        assert sg.device == g.device
+        assert sg.ntypes == g.ntypes
+        assert sg.etypes == g.etypes
+        assert sg.canonical_etypes == g.canonical_etypes
+        assert F.array_equal(
+            F.tensor(sg.nodes["user"].data[dgl.NID]), F.tensor([1, 2], g.idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.nodes["game"].data[dgl.NID]), F.tensor([0], g.idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], g.idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["plays"].data[dgl.EID]), F.tensor([1], g.idtype)
+        )
+        assert F.array_equal(
+            F.tensor(sg.edges["wishes"].data[dgl.EID]), F.tensor([1], g.idtype)
+        )
+        assert sg.num_nodes("developer") == 0
+        assert sg.num_edges("develops") == 0
+        assert F.array_equal(
+            sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
+        )
+        assert F.array_equal(
+            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
+        )
+
+    sg1 = g.subgraph({"user": [1, 2], "game": [0]})
+    _check_subgraph(g, sg1)
+    sg2 = g.edge_subgraph({"follows": [1], "plays": [1], "wishes": [1]})
+    _check_subgraph(g, sg2)
+
+    # backend tensor input
+    sg1 = g.subgraph(
+        {
+            "user": F.tensor([1, 2], dtype=idtype),
+            "game": F.tensor([0], dtype=idtype),
+        }
+    )
+    _check_subgraph(g, sg1)
+    sg2 = g.edge_subgraph(
+        {
+            "follows": F.tensor([1], dtype=idtype),
+            "plays": F.tensor([1], dtype=idtype),
+            "wishes": F.tensor([1], dtype=idtype),
+        }
+    )
+    _check_subgraph(g, sg2)
+
+    # numpy input
+    sg1 = g.subgraph({"user": np.array([1, 2]), "game": np.array([0])})
+    _check_subgraph(g, sg1)
+    sg2 = g.edge_subgraph(
+        {
+            "follows": np.array([1]),
+            "plays": np.array([1]),
+            "wishes": np.array([1]),
+        }
+    )
+    _check_subgraph(g, sg2)
+
+    def _check_subgraph_single_ntype(g, sg, preserve_nodes=False):
+        assert sg.idtype == g.idtype
+        assert sg.device == g.device
+        assert sg.ntypes == g.ntypes
+        assert sg.etypes == g.etypes
+        assert sg.canonical_etypes == g.canonical_etypes
+
+        if not preserve_nodes:
+            assert F.array_equal(
+                F.tensor(sg.nodes["user"].data[dgl.NID]),
+                F.tensor([1, 2], g.idtype),
+            )
+        else:
+            for ntype in sg.ntypes:
+                assert g.num_nodes(ntype) == sg.num_nodes(ntype)
+
+        assert F.array_equal(
+            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], g.idtype)
+        )
+
+        if not preserve_nodes:
+            assert F.array_equal(
+                sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
+            )
+        assert F.array_equal(
+            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
+        )
+
+    def _check_subgraph_single_etype(g, sg, preserve_nodes=False):
+        assert sg.ntypes == g.ntypes
+        assert sg.etypes == g.etypes
+        assert sg.canonical_etypes == g.canonical_etypes
+
+        if not preserve_nodes:
+            assert F.array_equal(
+                F.tensor(sg.nodes["user"].data[dgl.NID]),
+                F.tensor([0, 1], g.idtype),
+            )
+            assert F.array_equal(
+                F.tensor(sg.nodes["game"].data[dgl.NID]),
+                F.tensor([0], g.idtype),
+            )
+        else:
+            for ntype in sg.ntypes:
+                assert g.num_nodes(ntype) == sg.num_nodes(ntype)
+
+        assert F.array_equal(
+            F.tensor(sg.edges["plays"].data[dgl.EID]),
+            F.tensor([0, 1], g.idtype),
+        )
+
+    sg1_graph = g_graph.subgraph([1, 2])
+    _check_subgraph_single_ntype(g_graph, sg1_graph)
+    sg1_graph = g_graph.edge_subgraph([1])
+    _check_subgraph_single_ntype(g_graph, sg1_graph)
+    sg1_graph = g_graph.edge_subgraph([1], relabel_nodes=False)
+    _check_subgraph_single_ntype(g_graph, sg1_graph, True)
+    sg2_bipartite = g_bipartite.edge_subgraph([0, 1])
+    _check_subgraph_single_etype(g_bipartite, sg2_bipartite)
+    sg2_bipartite = g_bipartite.edge_subgraph([0, 1], relabel_nodes=False)
+    _check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)
+
+    def _check_typed_subgraph1(g, sg):
+        assert g.idtype == sg.idtype
+        assert g.device == sg.device
+        assert set(sg.ntypes) == {"user", "game"}
+        assert set(sg.etypes) == {"follows", "plays", "wishes"}
+        for ntype in sg.ntypes:
+            assert sg.num_nodes(ntype) == g.num_nodes(ntype)
+        for etype in sg.etypes:
+            src_sg, dst_sg = sg.all_edges(etype=etype, order="eid")
+            src_g, dst_g = g.all_edges(etype=etype, order="eid")
+            assert F.array_equal(src_sg, src_g)
+            assert F.array_equal(dst_sg, dst_g)
+        assert F.array_equal(
+            sg.nodes["user"].data["h"], g.nodes["user"].data["h"]
+        )
+        assert F.array_equal(
+            sg.edges["follows"].data["h"], g.edges["follows"].data["h"]
+        )
+        g.nodes["user"].data["h"] = F.scatter_row(
+            g.nodes["user"].data["h"], F.tensor([2]), F.randn((1, 5))
+        )
+        g.edges["follows"].data["h"] = F.scatter_row(
+            g.edges["follows"].data["h"], F.tensor([1]), F.randn((1, 4))
+        )
+        assert F.array_equal(
+            sg.nodes["user"].data["h"], g.nodes["user"].data["h"]
+        )
+        assert F.array_equal(
+            sg.edges["follows"].data["h"], g.edges["follows"].data["h"]
+        )
+
+    def _check_typed_subgraph2(g, sg):
+        assert set(sg.ntypes) == {"developer", "game"}
+        assert set(sg.etypes) == {"develops"}
+        for ntype in sg.ntypes:
+            assert sg.num_nodes(ntype) == g.num_nodes(ntype)
+        for etype in sg.etypes:
+            src_sg, dst_sg = sg.all_edges(etype=etype, order="eid")
+            src_g, dst_g = g.all_edges(etype=etype, order="eid")
+            assert F.array_equal(src_sg, src_g)
+            assert F.array_equal(dst_sg, dst_g)
+
+    sg3 = g.node_type_subgraph(["user", "game"])
+    _check_typed_subgraph1(g, sg3)
+    sg4 = g.edge_type_subgraph(["develops"])
+    _check_typed_subgraph2(g, sg4)
+    sg5 = g.edge_type_subgraph(["follows", "plays", "wishes"])
+    _check_typed_subgraph1(g, sg5)
+
+    # Test for restricted format
+    for fmt in ["csr", "csc", "coo"]:
+        g = dgl.graph(([0, 1], [1, 2])).formats(fmt)
+        sg = g.subgraph({g.ntypes[0]: [1, 0]})
+        nids = F.asnumpy(sg.ndata[dgl.NID])
+        assert np.array_equal(nids, np.array([1, 0]))
+        src, dst = sg.edges(order="eid")
+        src = F.asnumpy(src)
+        dst = F.asnumpy(dst)
+        assert np.array_equal(src, np.array([1]))
+
+
+@parametrize_idtype
+def test_in_subgraph(idtype):
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): (
+                [1, 2, 3, 0, 2, 3, 0],
+                [0, 0, 0, 1, 1, 1, 2],
+            ),
+            ("user", "play", "game"): ([0, 0, 1, 3], [0, 1, 2, 2]),
+            ("game", "liked-by", "user"): (
+                [2, 2, 2, 1, 1, 0],
+                [0, 1, 2, 0, 3, 0],
+            ),
+            ("user", "flips", "coin"): ([0, 1, 2, 3], [0, 0, 0, 0]),
+        },
+        idtype=idtype,
+        num_nodes_dict={"user": 5, "game": 10, "coin": 8},
+    ).to(F.ctx())
+    subg = dgl.in_subgraph(hg, {"user": [0, 1], "game": 0})
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+    u, v = subg["follow"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(
+        hg["follow"].edge_ids(u, v), subg["follow"].edata[dgl.EID]
+    )
+    assert edge_set == {(1, 0), (2, 0), (3, 0), (0, 1), (2, 1), (3, 1)}
+    u, v = subg["play"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(hg["play"].edge_ids(u, v), subg["play"].edata[dgl.EID])
+    assert edge_set == {(0, 0)}
+    u, v = subg["liked-by"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(
+        hg["liked-by"].edge_ids(u, v), subg["liked-by"].edata[dgl.EID]
+    )
+    assert edge_set == {(2, 0), (2, 1), (1, 0), (0, 0)}
+    assert subg["flips"].num_edges() == 0
+    for ntype in subg.ntypes:
+        assert dgl.NID not in subg.nodes[ntype].data
+
+    # Test store_ids
+    subg = dgl.in_subgraph(hg, {"user": [0, 1], "game": 0}, store_ids=False)
+    for etype in ["follow", "play", "liked-by"]:
+        assert dgl.EID not in subg.edges[etype].data
+    for ntype in subg.ntypes:
+        assert dgl.NID not in subg.nodes[ntype].data
+
+    # Test relabel nodes
+    subg = dgl.in_subgraph(hg, {"user": [0, 1], "game": 0}, relabel_nodes=True)
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+
+    u, v = subg["follow"].edges()
+    old_u = F.gather_row(subg.nodes["user"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["user"].data[dgl.NID], v)
+    assert F.array_equal(
+        hg["follow"].edge_ids(old_u, old_v), subg["follow"].edata[dgl.EID]
+    )
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(1, 0), (2, 0), (3, 0), (0, 1), (2, 1), (3, 1)}
+
+    u, v = subg["play"].edges()
+    old_u = F.gather_row(subg.nodes["user"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["game"].data[dgl.NID], v)
+    assert F.array_equal(
+        hg["play"].edge_ids(old_u, old_v), subg["play"].edata[dgl.EID]
+    )
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(0, 0)}
+
+    u, v = subg["liked-by"].edges()
+    old_u = F.gather_row(subg.nodes["game"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["user"].data[dgl.NID], v)
+    assert F.array_equal(
+        hg["liked-by"].edge_ids(old_u, old_v), subg["liked-by"].edata[dgl.EID]
+    )
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(2, 0), (2, 1), (1, 0), (0, 0)}
+
+    assert subg.num_nodes("user") == 4
+    assert subg.num_nodes("game") == 3
+    assert subg.num_nodes("coin") == 0
+    assert subg.num_edges("flips") == 0
+
+
+@parametrize_idtype
+def test_out_subgraph(idtype):
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): (
+                [1, 2, 3, 0, 2, 3, 0],
+                [0, 0, 0, 1, 1, 1, 2],
+            ),
+            ("user", "play", "game"): ([0, 0, 1, 3], [0, 1, 2, 2]),
+            ("game", "liked-by", "user"): (
+                [2, 2, 2, 1, 1, 0],
+                [0, 1, 2, 0, 3, 0],
+            ),
+            ("user", "flips", "coin"): ([0, 1, 2, 3], [0, 0, 0, 0]),
+        },
+        idtype=idtype,
+    ).to(F.ctx())
+    subg = dgl.out_subgraph(hg, {"user": [0, 1], "game": 0})
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+    u, v = subg["follow"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(1, 0), (0, 1), (0, 2)}
+    assert F.array_equal(
+        hg["follow"].edge_ids(u, v), subg["follow"].edata[dgl.EID]
+    )
+    u, v = subg["play"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0), (0, 1), (1, 2)}
+    assert F.array_equal(hg["play"].edge_ids(u, v), subg["play"].edata[dgl.EID])
+    u, v = subg["liked-by"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0)}
+    assert F.array_equal(
+        hg["liked-by"].edge_ids(u, v), subg["liked-by"].edata[dgl.EID]
+    )
+    u, v = subg["flips"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0), (1, 0)}
+    assert F.array_equal(
+        hg["flips"].edge_ids(u, v), subg["flips"].edata[dgl.EID]
+    )
+    for ntype in subg.ntypes:
+        assert dgl.NID not in subg.nodes[ntype].data
+
+    # Test store_ids
+    subg = dgl.out_subgraph(hg, {"user": [0, 1], "game": 0}, store_ids=False)
+    for etype in subg.canonical_etypes:
+        assert dgl.EID not in subg.edges[etype].data
+    for ntype in subg.ntypes:
+        assert dgl.NID not in subg.nodes[ntype].data
+
+    # Test relabel nodes
+    subg = dgl.out_subgraph(hg, {"user": [1], "game": 0}, relabel_nodes=True)
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+
+    u, v = subg["follow"].edges()
+    old_u = F.gather_row(subg.nodes["user"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["user"].data[dgl.NID], v)
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(1, 0)}
+    assert F.array_equal(
+        hg["follow"].edge_ids(old_u, old_v), subg["follow"].edata[dgl.EID]
+    )
+
+    u, v = subg["play"].edges()
+    old_u = F.gather_row(subg.nodes["user"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["game"].data[dgl.NID], v)
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(1, 2)}
+    assert F.array_equal(
+        hg["play"].edge_ids(old_u, old_v), subg["play"].edata[dgl.EID]
+    )
+
+    u, v = subg["liked-by"].edges()
+    old_u = F.gather_row(subg.nodes["game"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["user"].data[dgl.NID], v)
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(0, 0)}
+    assert F.array_equal(
+        hg["liked-by"].edge_ids(old_u, old_v), subg["liked-by"].edata[dgl.EID]
+    )
+
+    u, v = subg["flips"].edges()
+    old_u = F.gather_row(subg.nodes["user"].data[dgl.NID], u)
+    old_v = F.gather_row(subg.nodes["coin"].data[dgl.NID], v)
+    edge_set = set(zip(list(F.asnumpy(old_u)), list(F.asnumpy(old_v))))
+    assert edge_set == {(1, 0)}
+    assert F.array_equal(
+        hg["flips"].edge_ids(old_u, old_v), subg["flips"].edata[dgl.EID]
+    )
+    assert subg.num_nodes("user") == 2
+    assert subg.num_nodes("game") == 2
+    assert subg.num_nodes("coin") == 1
+
+
+def test_subgraph_message_passing():
+    # Unit test for PR #2055
+    g = dgl.graph(([0, 1, 2], [2, 3, 4])).to(F.cpu())
+    g.ndata["x"] = F.copy_to(F.randn((5, 6)), F.cpu())
+    sg = g.subgraph([1, 2, 3]).to(F.ctx())
+    sg.update_all(
+        lambda edges: {"x": edges.src["x"]},
+        lambda nodes: {"y": F.sum(nodes.mailbox["x"], 1)},
+    )
+
+
+@parametrize_idtype
+def test_khop_in_subgraph(idtype):
+    g = dgl.graph(
+        ([1, 1, 2, 3, 4], [0, 2, 0, 4, 2]), idtype=idtype, device=F.ctx()
+    )
+    g.edata["w"] = F.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
+    sg, inv = dgl.khop_in_subgraph(g, 0, k=2)
+    assert sg.idtype == g.idtype
+    u, v = sg.edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(1, 0), (1, 2), (2, 0), (3, 2)}
+    assert F.array_equal(
+        sg.edata[dgl.EID], F.tensor([0, 1, 2, 4], dtype=idtype)
+    )
+    assert F.array_equal(
+        sg.edata["w"], F.tensor([[0, 1], [2, 3], [4, 5], [8, 9]])
+    )
+    assert F.array_equal(F.astype(inv, idtype), F.tensor([0], idtype))
+
+    # Test multiple nodes
+    sg, inv = dgl.khop_in_subgraph(g, [0, 2], k=1)
+    assert sg.num_edges() == 4
+
+    sg, inv = dgl.khop_in_subgraph(g, F.tensor([0, 2], idtype), k=1)
+    assert sg.num_edges() == 4
+
+    # Test isolated node
+    sg, inv = dgl.khop_in_subgraph(g, 1, k=2)
+    assert sg.idtype == g.idtype
+    assert sg.num_nodes() == 1
+    assert sg.num_edges() == 0
+    assert F.array_equal(F.astype(inv, idtype), F.tensor([0], idtype))
+
+    g = dgl.heterograph(
+        {
+            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 2, 1]),
+            ("user", "follows", "user"): ([0, 1, 1], [1, 2, 2]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    sg, inv = dgl.khop_in_subgraph(g, {"game": 0}, k=2)
+    assert sg.idtype == idtype
+    assert sg.num_nodes("game") == 1
+    assert sg.num_nodes("user") == 2
+    assert len(sg.ntypes) == 2
+    assert len(sg.etypes) == 2
+    u, v = sg["follows"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 1)}
+    u, v = sg["plays"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0), (1, 0)}
+    assert F.array_equal(F.astype(inv["game"], idtype), F.tensor([0], idtype))
+
+    # Test isolated node
+    sg, inv = dgl.khop_in_subgraph(g, {"user": 0}, k=2)
+    assert sg.idtype == idtype
+    assert sg.num_nodes("game") == 0
+    assert sg.num_nodes("user") == 1
+    assert sg.num_edges("follows") == 0
+    assert sg.num_edges("plays") == 0
+    assert F.array_equal(F.astype(inv["user"], idtype), F.tensor([0], idtype))
+
+    # Test multiple nodes
+    sg, inv = dgl.khop_in_subgraph(
+        g, {"user": F.tensor([0, 1], idtype), "game": 0}, k=1
+    )
+    u, v = sg["follows"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 1)}
+    u, v = sg["plays"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0), (1, 0)}
+    assert F.array_equal(
+        F.astype(inv["user"], idtype), F.tensor([0, 1], idtype)
+    )
+    assert F.array_equal(F.astype(inv["game"], idtype), F.tensor([0], idtype))
+
+
+@parametrize_idtype
+def test_khop_out_subgraph(idtype):
+    g = dgl.graph(
+        ([0, 2, 0, 4, 2], [1, 1, 2, 3, 4]), idtype=idtype, device=F.ctx()
+    )
+    g.edata["w"] = F.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
+    sg, inv = dgl.khop_out_subgraph(g, 0, k=2)
+    assert sg.idtype == g.idtype
+    u, v = sg.edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 1), (2, 1), (0, 2), (2, 3)}
+    assert F.array_equal(
+        sg.edata[dgl.EID], F.tensor([0, 2, 1, 4], dtype=idtype)
+    )
+    assert F.array_equal(
+        sg.edata["w"], F.tensor([[0, 1], [4, 5], [2, 3], [8, 9]])
+    )
+    assert F.array_equal(F.astype(inv, idtype), F.tensor([0], idtype))
+
+    # Test multiple nodes
+    sg, inv = dgl.khop_out_subgraph(g, [0, 2], k=1)
+    assert sg.num_edges() == 4
+
+    sg, inv = dgl.khop_out_subgraph(g, F.tensor([0, 2], idtype), k=1)
+    assert sg.num_edges() == 4
+
+    # Test isolated node
+    sg, inv = dgl.khop_out_subgraph(g, 1, k=2)
+    assert sg.idtype == g.idtype
+    assert sg.num_nodes() == 1
+    assert sg.num_edges() == 0
+    assert F.array_equal(F.astype(inv, idtype), F.tensor([0], idtype))
+
+    g = dgl.heterograph(
+        {
+            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 2, 1]),
+            ("user", "follows", "user"): ([0, 1], [1, 3]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    sg, inv = dgl.khop_out_subgraph(g, {"user": 0}, k=2)
+    assert sg.idtype == idtype
+    assert sg.num_nodes("game") == 2
+    assert sg.num_nodes("user") == 3
+    assert len(sg.ntypes) == 2
+    assert len(sg.etypes) == 2
+    u, v = sg["follows"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 1), (1, 2)}
+    u, v = sg["plays"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 0), (1, 0), (1, 1)}
+    assert F.array_equal(F.astype(inv["user"], idtype), F.tensor([0], idtype))
+
+    # Test isolated node
+    sg, inv = dgl.khop_out_subgraph(g, {"user": 3}, k=2)
+    assert sg.idtype == idtype
+    assert sg.num_nodes("game") == 0
+    assert sg.num_nodes("user") == 1
+    assert sg.num_edges("follows") == 0
+    assert sg.num_edges("plays") == 0
+    assert F.array_equal(F.astype(inv["user"], idtype), F.tensor([0], idtype))
+
+    # Test multiple nodes
+    sg, inv = dgl.khop_out_subgraph(
+        g, {"user": F.tensor([2], idtype), "game": 0}, k=1
+    )
+    assert sg.num_edges("follows") == 0
+    u, v = sg["plays"].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0, 1)}
+    assert F.array_equal(F.astype(inv["user"], idtype), F.tensor([0], idtype))
+    assert F.array_equal(F.astype(inv["game"], idtype), F.tensor([0], idtype))
+
+
+@unittest.skipIf(not F.gpu_ctx(), "only necessary with GPU")
+@pytest.mark.parametrize(
+    "parent_idx_device",
+    [("cpu", F.cpu()), ("cuda", F.cuda()), ("uva", F.cpu()), ("uva", F.cuda())],
+)
+@pytest.mark.parametrize("child_device", [F.cpu(), F.cuda()])
+def test_subframes(parent_idx_device, child_device):
+    parent_device, idx_device = parent_idx_device
+    g = dgl.graph(
+        (F.tensor([1, 2, 3], dtype=F.int64), F.tensor([2, 3, 4], dtype=F.int64))
+    )
+    print(g.device)
+    g.ndata["x"] = F.randn((5, 4))
+    g.edata["a"] = F.randn((3, 6))
+    idx = F.tensor([1, 2], dtype=F.int64)
+    if parent_device == "cuda":
+        g = g.to(F.cuda())
+    elif parent_device == "uva":
+        if F.backend_name != "pytorch":
+            pytest.skip("UVA only supported for PyTorch")
+        g = g.to(F.cpu())
+        g.create_formats_()
+        g.pin_memory_()
+    elif parent_device == "cpu":
+        g = g.to(F.cpu())
+    idx = F.copy_to(idx, idx_device)
+    sg = g.sample_neighbors(idx, 2).to(child_device)
+    assert sg.device == F.context(sg.ndata["x"])
+    assert sg.device == F.context(sg.edata["a"])
+    assert sg.device == child_device
+    if parent_device != "uva":
+        sg = g.to(child_device).sample_neighbors(
+            F.copy_to(idx, child_device), 2
+        )
+        assert sg.device == F.context(sg.ndata["x"])
+        assert sg.device == F.context(sg.edata["a"])
+        assert sg.device == child_device
+    if parent_device == "uva":
+        g.unpin_memory_()
+
+
+@unittest.skipIf(
+    F._default_context_str != "gpu", reason="UVA only available on GPU"
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="UVA only supported for PyTorch",
+)
+@pytest.mark.parametrize("device", [F.cpu(), F.cuda()])
+@parametrize_idtype
+def test_uva_subgraph(idtype, device):
+    g = create_test_heterograph2(idtype)
+    g = g.to(F.cpu())
+    g.create_formats_()
+    g.pin_memory_()
+    indices = {"user": F.copy_to(F.tensor([0], idtype), device)}
+    edge_indices = {"follows": F.copy_to(F.tensor([0], idtype), device)}
+    assert g.subgraph(indices).device == device
+    assert g.edge_subgraph(edge_indices).device == device
+    assert g.in_subgraph(indices).device == device
+    assert g.out_subgraph(indices).device == device
+    assert g.khop_in_subgraph(indices, 1)[0].device == device
+    assert g.khop_out_subgraph(indices, 1)[0].device == device
+    assert g.sample_neighbors(indices, 1).device == device
+    g.unpin_memory_()
+
+
+if __name__ == "__main__":
+    test_edge_subgraph()
+    test_uva_subgraph(F.int64, F.cpu())
+    test_uva_subgraph(F.int64, F.cuda())

tests/python/common/test_traversal.py

@@ -0,0 +1,140 @@
+import itertools
+import random
+import sys
+import time
+import unittest
+
+import backend as F
+
+import dgl
+import networkx as nx
+import numpy as np
+import scipy.sparse as sp
+from utils import parametrize_idtype
+
+np.random.seed(42)
+
+
+def toset(x):
+    # F.zerocopy_to_numpy may return a int
+    return set(F.zerocopy_to_numpy(x).tolist())
+
+
+@parametrize_idtype
+def test_bfs(idtype, n=100):
+    def _bfs_nx(g_nx, src):
+        edges = nx.bfs_edges(g_nx, src)
+        layers_nx = [set([src])]
+        edges_nx = []
+        frontier = set()
+        edge_frontier = set()
+        for u, v in edges:
+            if u in layers_nx[-1]:
+                frontier.add(v)
+                edge_frontier.add(g.edge_ids(int(u), int(v)))
+            else:
+                layers_nx.append(frontier)
+                edges_nx.append(edge_frontier)
+                frontier = set([v])
+                edge_frontier = set([g.edge_ids(u, v)])
+        # avoids empty successors
+        if len(frontier) > 0 and len(edge_frontier) > 0:
+            layers_nx.append(frontier)
+            edges_nx.append(edge_frontier)
+        return layers_nx, edges_nx
+
+    a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
+    g = dgl.from_scipy(a).astype(idtype)
+
+    g_nx = g.to_networkx()
+    src = random.choice(range(n))
+    layers_nx, _ = _bfs_nx(g_nx, src)
+    layers_dgl = dgl.bfs_nodes_generator(g, src)
+    assert len(layers_dgl) == len(layers_nx)
+    assert all(toset(x) == y for x, y in zip(layers_dgl, layers_nx))
+
+    g_nx = nx.random_labeled_tree(n, seed=42)
+    g = dgl.from_networkx(g_nx).astype(idtype)
+    src = 0
+    _, edges_nx = _bfs_nx(g_nx, src)
+    edges_dgl = dgl.bfs_edges_generator(g, src)
+    assert len(edges_dgl) == len(edges_nx)
+    assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))
+
+
+@parametrize_idtype
+def test_topological_nodes(idtype, n=100):
+    a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
+    b = sp.tril(a, -1).tocoo()
+    g = dgl.from_scipy(b).astype(idtype)
+
+    layers_dgl = dgl.topological_nodes_generator(g)
+
+    adjmat = g.adj_external(transpose=True)
+
+    def tensor_topo_traverse():
+        n = g.num_nodes()
+        mask = F.copy_to(F.ones((n, 1)), F.cpu())
+        degree = F.spmm(adjmat, mask)
+        while F.reduce_sum(mask) != 0.0:
+            v = F.astype((degree == 0.0), F.float32)
+            v = v * mask
+            mask = mask - v
+            frontier = F.copy_to(F.nonzero_1d(F.squeeze(v, 1)), F.cpu())
+            yield frontier
+            degree -= F.spmm(adjmat, v)
+
+    layers_spmv = list(tensor_topo_traverse())
+
+    assert len(layers_dgl) == len(layers_spmv)
+    assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))
+
+
+DFS_LABEL_NAMES = ["forward", "reverse", "nontree"]
+
+
+@parametrize_idtype
+def test_dfs_labeled_edges(idtype, example=False):
+    dgl_g = dgl.graph([]).astype(idtype)
+    dgl_g.add_nodes(6)
+    dgl_g.add_edges([0, 1, 0, 3, 3], [1, 2, 2, 4, 5])
+    dgl_edges, dgl_labels = dgl.dfs_labeled_edges_generator(
+        dgl_g, [0, 3], has_reverse_edge=True, has_nontree_edge=True
+    )
+    dgl_edges = [toset(t) for t in dgl_edges]
+    dgl_labels = [toset(t) for t in dgl_labels]
+    g1_solutions = [
+        # edges           labels
+        [[0, 1, 1, 0, 2], [0, 0, 1, 1, 2]],
+        [[2, 2, 0, 1, 0], [0, 1, 0, 2, 1]],
+    ]
+    g2_solutions = [
+        # edges        labels
+        [[3, 3, 4, 4], [0, 1, 0, 1]],
+        [[4, 4, 3, 3], [0, 1, 0, 1]],
+    ]
+
+    def combine_frontiers(sol):
+        es, ls = zip(*sol)
+        es = [
+            set(i for i in t if i is not None)
+            for t in itertools.zip_longest(*es)
+        ]
+        ls = [
+            set(i for i in t if i is not None)
+            for t in itertools.zip_longest(*ls)
+        ]
+        return es, ls
+
+    for sol_set in itertools.product(g1_solutions, g2_solutions):
+        es, ls = combine_frontiers(sol_set)
+        if es == dgl_edges and ls == dgl_labels:
+            break
+    else:
+        assert False
+
+
+if __name__ == "__main__":
+    test_bfs(idtype="int32")
+    test_topological_nodes(idtype="int32")
+    test_dfs_labeled_edges(idtype="int32")

tests/python/common/transforms/test_functional-sort.py

@@ -0,0 +1,120 @@
+import itertools
+import unittest
+from collections import Counter
+
+import backend as F
+
+import dgl
+import dgl.function as fn
+import networkx as nx
+import numpy as np
+import pytest
+import scipy.sparse as ssp
+from dgl import DGLError
+from utils import parametrize_idtype
+
+
+def create_test_heterograph(num_nodes, num_adj, idtype):
+    if isinstance(num_adj, int):
+        num_adj = [num_adj, num_adj + 1]
+    num_adj_list = list(
+        np.random.choice(np.arange(num_adj[0], num_adj[1]), num_nodes)
+    )
+    src = np.concatenate([[i] * num_adj_list[i] for i in range(num_nodes)])
+    dst = [
+        np.random.choice(num_nodes, nadj, replace=False)
+        for nadj in num_adj_list
+    ]
+    dst = np.concatenate(dst)
+    return dgl.graph((src, dst), idtype=idtype)
+
+
+def check_sort(spm, tag_arr=None, tag_pos=None):
+    if tag_arr is None:
+        tag_arr = np.arange(spm.shape[0])
+    else:
+        tag_arr = F.asnumpy(tag_arr)
+    if tag_pos is not None:
+        tag_pos = F.asnumpy(tag_pos)
+    for i in range(spm.shape[0]):
+        row = spm.getrow(i)
+        dst = row.nonzero()[1]
+        if tag_pos is not None:
+            tag_pos_row = tag_pos[i]
+            tag_pos_ptr = tag_arr[dst[0]] if len(dst) > 0 else 0
+        for j in range(len(dst) - 1):
+            if tag_pos is not None and tag_arr[dst[j]] != tag_pos_ptr:
+                # `tag_pos_ptr` is the expected tag value. Here we check whether the
+                # tag value is equal to `tag_pos_ptr`
+                return False
+            if tag_arr[dst[j]] > tag_arr[dst[j + 1]]:
+                # The tag should be in ascending order after sorting
+                return False
+            if tag_pos is not None and tag_arr[dst[j]] < tag_arr[dst[j + 1]]:
+                if j + 1 != int(tag_pos_row[tag_pos_ptr + 1]):
+                    # The boundary of tag should be consistent with `tag_pos`
+                    return False
+                tag_pos_ptr = tag_arr[dst[j + 1]]
+    return True
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu", reason="GPU sorting by tag not implemented"
+)
+@parametrize_idtype
+def test_sort_with_tag(idtype):
+    num_nodes, num_adj, num_tags = 200, [20, 50], 5
+    g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
+    tag = F.tensor(np.random.choice(num_tags, g.num_nodes()))
+    src, dst = g.edges()
+    edge_tag_dst = F.gather_row(tag, F.tensor(dst))
+    edge_tag_src = F.gather_row(tag, F.tensor(src))
+
+    for tag_type in ["node", "edge"]:
+        new_g = dgl.sort_csr_by_tag(
+            g, tag if tag_type == "node" else edge_tag_dst, tag_type=tag_type
+        )
+        old_csr = g.adj_external(scipy_fmt="csr")
+        new_csr = new_g.adj_external(scipy_fmt="csr")
+        assert check_sort(new_csr, tag, new_g.dstdata["_TAG_OFFSET"])
+        assert not check_sort(
+            old_csr, tag
+        )  # Check the original csr is not modified.
+
+    for tag_type in ["node", "edge"]:
+        new_g = dgl.sort_csc_by_tag(
+            g, tag if tag_type == "node" else edge_tag_src, tag_type=tag_type
+        )
+        old_csc = g.adj_external(transpose=True, scipy_fmt="csr")
+        new_csc = new_g.adj_external(transpose=True, scipy_fmt="csr")
+        assert check_sort(new_csc, tag, new_g.srcdata["_TAG_OFFSET"])
+        assert not check_sort(old_csc, tag)
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu", reason="GPU sorting by tag not implemented"
+)
+@parametrize_idtype
+def test_sort_with_tag_bipartite(idtype):
+    num_nodes, num_adj, num_tags = 200, [20, 50], 5
+    g = create_test_heterograph(num_nodes, num_adj, idtype=idtype)
+    g = dgl.heterograph({("_U", "_E", "_V"): g.edges()})
+    utag = F.tensor(np.random.choice(num_tags, g.num_nodes("_U")))
+    vtag = F.tensor(np.random.choice(num_tags, g.num_nodes("_V")))
+
+    new_g = dgl.sort_csr_by_tag(g, vtag)
+    old_csr = g.adj_external(scipy_fmt="csr")
+    new_csr = new_g.adj_external(scipy_fmt="csr")
+    assert check_sort(new_csr, vtag, new_g.nodes["_U"].data["_TAG_OFFSET"])
+    assert not check_sort(old_csr, vtag)
+
+    new_g = dgl.sort_csc_by_tag(g, utag)
+    old_csc = g.adj_external(transpose=True, scipy_fmt="csr")
+    new_csc = new_g.adj_external(transpose=True, scipy_fmt="csr")
+    assert check_sort(new_csc, utag, new_g.nodes["_V"].data["_TAG_OFFSET"])
+    assert not check_sort(old_csc, utag)
+
+
+if __name__ == "__main__":
+    test_sort_with_tag(F.int32)
+    test_sort_with_tag_bipartite(F.int32)

tests/python/common/transforms/test_to_block.py

@@ -0,0 +1,192 @@
+##
+#   Copyright 2019-2021 Contributors
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+
+import backend as F
+
+import dgl
+import dgl.partition
+from utils import parametrize_idtype
+
+
+@parametrize_idtype
+def test_to_block(idtype):
+    def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
+        if dst_nodes is not None:
+            assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
+        n_dst_nodes = bg.num_nodes("DST/" + ntype)
+        if include_dst_in_src:
+            assert F.array_equal(
+                bg.srcnodes[ntype].data[dgl.NID][:n_dst_nodes],
+                bg.dstnodes[ntype].data[dgl.NID],
+            )
+
+        g = g[etype]
+        bg = bg[etype]
+        induced_src = bg.srcdata[dgl.NID]
+        induced_dst = bg.dstdata[dgl.NID]
+        induced_eid = bg.edata[dgl.EID]
+
+        bg_src, bg_dst = bg.all_edges(order="eid")
+        src_ans, dst_ans = g.all_edges(order="eid")
+
+        induced_src_bg = F.gather_row(induced_src, bg_src)
+        induced_dst_bg = F.gather_row(induced_dst, bg_dst)
+        induced_src_ans = F.gather_row(src_ans, induced_eid)
+        induced_dst_ans = F.gather_row(dst_ans, induced_eid)
+
+        assert F.array_equal(induced_src_bg, induced_src_ans)
+        assert F.array_equal(induced_dst_bg, induced_dst_ans)
+
+    def checkall(g, bg, dst_nodes, include_dst_in_src=True):
+        for etype in g.etypes:
+            ntype = g.to_canonical_etype(etype)[2]
+            if dst_nodes is not None and ntype in dst_nodes:
+                check(g, bg, ntype, etype, dst_nodes[ntype], include_dst_in_src)
+            else:
+                check(g, bg, ntype, etype, None, include_dst_in_src)
+
+    # homogeneous graph
+    g = dgl.graph(
+        (F.tensor([1, 2], dtype=idtype), F.tensor([2, 3], dtype=idtype))
+    )
+    dst_nodes = F.tensor([3, 2], dtype=idtype)
+    bg = dgl.to_block(g, dst_nodes=dst_nodes)
+    check(g, bg, "_N", "_E", dst_nodes)
+
+    src_nodes = bg.srcnodes["_N"].data[dgl.NID]
+    bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
+    check(g, bg, "_N", "_E", dst_nodes)
+
+    # heterogeneous graph
+    g = dgl.heterograph(
+        {
+            ("A", "AA", "A"): ([0, 2, 1, 3], [1, 3, 2, 4]),
+            ("A", "AB", "B"): ([0, 1, 3, 1], [1, 3, 5, 6]),
+            ("B", "BA", "A"): ([2, 3], [3, 2]),
+        },
+        idtype=idtype,
+        device=F.ctx(),
+    )
+    g.nodes["A"].data["x"] = F.randn((5, 10))
+    g.nodes["B"].data["x"] = F.randn((7, 5))
+    g.edges["AA"].data["x"] = F.randn((4, 3))
+    g.edges["AB"].data["x"] = F.randn((4, 3))
+    g.edges["BA"].data["x"] = F.randn((2, 3))
+    g_a = g["AA"]
+
+    def check_features(g, bg):
+        for ntype in bg.srctypes:
+            for key in g.nodes[ntype].data:
+                assert F.array_equal(
+                    bg.srcnodes[ntype].data[key],
+                    F.gather_row(
+                        g.nodes[ntype].data[key],
+                        bg.srcnodes[ntype].data[dgl.NID],
+                    ),
+                )
+        for ntype in bg.dsttypes:
+            for key in g.nodes[ntype].data:
+                assert F.array_equal(
+                    bg.dstnodes[ntype].data[key],
+                    F.gather_row(
+                        g.nodes[ntype].data[key],
+                        bg.dstnodes[ntype].data[dgl.NID],
+                    ),
+                )
+        for etype in bg.canonical_etypes:
+            for key in g.edges[etype].data:
+                assert F.array_equal(
+                    bg.edges[etype].data[key],
+                    F.gather_row(
+                        g.edges[etype].data[key], bg.edges[etype].data[dgl.EID]
+                    ),
+                )
+
+    bg = dgl.to_block(g_a)
+    check(g_a, bg, "A", "AA", None)
+    check_features(g_a, bg)
+    assert bg.number_of_src_nodes() == 5
+    assert bg.number_of_dst_nodes() == 4
+
+    bg = dgl.to_block(g_a, include_dst_in_src=False)
+    check(g_a, bg, "A", "AA", None, False)
+    check_features(g_a, bg)
+    assert bg.number_of_src_nodes() == 4
+    assert bg.number_of_dst_nodes() == 4
+
+    dst_nodes = F.tensor([4, 3, 2, 1], dtype=idtype)
+    bg = dgl.to_block(g_a, dst_nodes)
+    check(g_a, bg, "A", "AA", dst_nodes)
+    check_features(g_a, bg)
+
+    g_ab = g["AB"]
+
+    bg = dgl.to_block(g_ab)
+    assert bg.idtype == idtype
+    assert bg.num_nodes("SRC/B") == 4
+    assert F.array_equal(
+        bg.srcnodes["B"].data[dgl.NID], bg.dstnodes["B"].data[dgl.NID]
+    )
+    assert bg.num_nodes("DST/A") == 0
+    checkall(g_ab, bg, None)
+    check_features(g_ab, bg)
+
+    dst_nodes = {"B": F.tensor([5, 6, 3, 1], dtype=idtype)}
+    bg = dgl.to_block(g, dst_nodes)
+    assert bg.num_nodes("SRC/B") == 4
+    assert F.array_equal(
+        bg.srcnodes["B"].data[dgl.NID], bg.dstnodes["B"].data[dgl.NID]
+    )
+    assert bg.num_nodes("DST/A") == 0
+    checkall(g, bg, dst_nodes)
+    check_features(g, bg)
+
+    dst_nodes = {
+        "A": F.tensor([4, 3, 2, 1], dtype=idtype),
+        "B": F.tensor([3, 5, 6, 1], dtype=idtype),
+    }
+    bg = dgl.to_block(g, dst_nodes=dst_nodes)
+    checkall(g, bg, dst_nodes)
+    check_features(g, bg)
+
+    # test specifying lhs_nodes with include_dst_in_src
+    src_nodes = {}
+    for ntype in dst_nodes.keys():
+        # use the previous run to get the list of source nodes
+        src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
+    bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
+    checkall(g, bg, dst_nodes)
+    check_features(g, bg)
+
+    # test without include_dst_in_src
+    dst_nodes = {
+        "A": F.tensor([4, 3, 2, 1], dtype=idtype),
+        "B": F.tensor([3, 5, 6, 1], dtype=idtype),
+    }
+    bg = dgl.to_block(g, dst_nodes=dst_nodes, include_dst_in_src=False)
+    checkall(g, bg, dst_nodes, False)
+    check_features(g, bg)
+
+    # test specifying lhs_nodes without include_dst_in_src
+    src_nodes = {}
+    for ntype in dst_nodes.keys():
+        # use the previous run to get the list of source nodes
+        src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
+    bg = dgl.to_block(
+        g, dst_nodes=dst_nodes, include_dst_in_src=False, src_nodes=src_nodes
+    )
+    checkall(g, bg, dst_nodes, False)
+    check_features(g, bg)

tests/python/common/utils/test_filter.py

@@ -0,0 +1,98 @@
+import unittest
+
+import backend as F
+
+import dgl
+import numpy as np
+from dgl.utils import Filter
+from utils import parametrize_idtype
+
+
+def test_graph_filter():
+    g = dgl.graph([]).to(F.ctx())
+    g.add_nodes(4)
+    g.add_edges([0, 1, 2, 3], [1, 2, 3, 0])
+
+    n_repr = np.zeros((4, 5))
+    e_repr = np.zeros((4, 5))
+    n_repr[[1, 3]] = 1
+    e_repr[[1, 3]] = 1
+    n_repr = F.copy_to(F.zerocopy_from_numpy(n_repr), F.ctx())
+    e_repr = F.copy_to(F.zerocopy_from_numpy(e_repr), F.ctx())
+
+    g.ndata["a"] = n_repr
+    g.edata["a"] = e_repr
+
+    def predicate(r):
+        return F.max(r.data["a"], 1) > 0
+
+    # full node filter
+    n_idx = g.filter_nodes(predicate)
+    assert set(F.zerocopy_to_numpy(n_idx)) == {1, 3}
+
+    # partial node filter
+    n_idx = g.filter_nodes(predicate, [0, 1])
+    assert set(F.zerocopy_to_numpy(n_idx)) == {1}
+
+    # full edge filter
+    e_idx = g.filter_edges(predicate)
+    assert set(F.zerocopy_to_numpy(e_idx)) == {1, 3}
+
+    # partial edge filter
+    e_idx = g.filter_edges(predicate, [0, 1])
+    assert set(F.zerocopy_to_numpy(e_idx)) == {1}
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="CPU not yet supported"
+)
+@parametrize_idtype
+def test_array_filter(idtype):
+    f = Filter(
+        F.copy_to(F.tensor([0, 1, 9, 4, 6, 5, 7], dtype=idtype), F.ctx())
+    )
+    x = F.copy_to(F.tensor([0, 3, 9, 11], dtype=idtype), F.ctx())
+    y = F.copy_to(
+        F.tensor([0, 19, 0, 28, 3, 9, 11, 4, 5], dtype=idtype), F.ctx()
+    )
+
+    xi_act = f.find_included_indices(x)
+    xi_exp = F.copy_to(F.tensor([0, 2], dtype=idtype), F.ctx())
+    assert F.array_equal(xi_act, xi_exp)
+    xe_act = f.find_excluded_indices(x)
+    xe_exp = F.copy_to(F.tensor([1, 3], dtype=idtype), F.ctx())
+    assert F.array_equal(xe_act, xe_exp)
+
+    yi_act = f.find_included_indices(y)
+    yi_exp = F.copy_to(F.tensor([0, 2, 5, 7, 8], dtype=idtype), F.ctx())
+    assert F.array_equal(yi_act, yi_exp)
+    ye_act = f.find_excluded_indices(y)
+    ye_exp = F.copy_to(F.tensor([1, 3, 4, 6], dtype=idtype), F.ctx())
+    assert F.array_equal(ye_act, ye_exp)
+
+
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch",
+    reason="Multiple streams are only supported by pytorch backend",
+)
+@unittest.skipIf(
+    F._default_context_str == "cpu", reason="CPU not yet supported"
+)
+@parametrize_idtype
+def test_filter_multistream(idtype):
+    # this is a smoke test to ensure we do not trip any internal assertions
+    import torch
+
+    s = torch.cuda.Stream(device=F.ctx())
+    with torch.cuda.stream(s):
+        # we must do multiple runs such that the stream is busy as we launch
+        # work
+        for i in range(10):
+            f = Filter(F.arange(1000, 4000, dtype=idtype, ctx=F.ctx()))
+            x = F.randint([30000], dtype=idtype, ctx=F.ctx(), low=0, high=50000)
+            xi = f.find_included_indices(x)
+
+
+if __name__ == "__main__":
+    test_graph_filter()
+    test_array_filter()

tests/python/common/utils/test_pin_memory.py

@@ -0,0 +1,37 @@
+import backend as F
+
+import dgl
+import pytest
+
+
+@pytest.mark.skipif(
+    F._default_context_str == "cpu", reason="Need gpu for this test"
+)
+def test_pin_unpin():
+    t = F.arange(0, 100, dtype=F.int64, ctx=F.cpu())
+
+    assert not F.is_pinned(t)
+
+    if F.backend_name == "pytorch":
+        nd = dgl.utils.pin_memory_inplace(t)
+        assert F.is_pinned(t)
+        nd.unpin_memory_()
+        assert not F.is_pinned(t)
+        del nd
+
+        # tensor will be unpinned immediately if the returned ndarray is not saved
+        dgl.utils.pin_memory_inplace(t)
+        assert not F.is_pinned(t)
+
+        t_pin = t.pin_memory()
+        # cannot unpin a tensor that is pinned outside of DGL
+        with pytest.raises(dgl.DGLError):
+            F.to_dgl_nd(t_pin).unpin_memory_()
+    else:
+        with pytest.raises(dgl.DGLError):
+            # tensorflow and mxnet should throw an error
+            dgl.utils.pin_memory_inplace(t)
+
+
+if __name__ == "__main__":
+    test_pin_unpin()