import itertools
import multiprocessing as mp
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
import test_utils
from dgl import DGLError
from scipy.sparse import rand
from test_utils import get_cases, parametrize_idtype
from utils import assert_is_identical_hetero


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(),
    )
    assert g.idtype == idtype
    assert g.device == F.ctx()
    return g


def create_test_heterograph1(idtype):
    edges = []
    edges.extend([(0, 1), (1, 2)])  # follows
    edges.extend([(0, 3), (1, 3), (2, 4), (1, 4)])  # plays
    edges.extend([(0, 4), (2, 3)])  # wishes
    edges.extend([(5, 3), (6, 4)])  # develops
    edges = tuple(zip(*edges))
    ntypes = F.tensor([0, 0, 0, 1, 1, 2, 2])
    etypes = F.tensor([0, 0, 1, 1, 1, 1, 2, 2, 3, 3])
    g0 = dgl.graph(edges, idtype=idtype, device=F.ctx())
    g0.ndata[dgl.NTYPE] = ntypes
    g0.edata[dgl.ETYPE] = etypes
    return dgl.to_heterogeneous(
        g0,
        ["user", "game", "developer"],
        ["follows", "plays", "wishes", "develops"],
    )


def create_test_heterograph2(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"): ("csr", ([0, 1, 1, 2], [1, 0], [])),
            ("developer", "develops", "game"): (
                "csc",
                ([0, 1, 2], [0, 1], [0, 1]),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.idtype == idtype
    assert g.device == F.ctx()
    return g


def create_test_heterograph3(idtype):
    g = dgl.heterograph(
        {
            ("user", "plays", "game"): (
                F.tensor([0, 1, 1, 2], dtype=idtype),
                F.tensor([0, 0, 1, 1], dtype=idtype),
            ),
            ("developer", "develops", "game"): (
                F.tensor([0, 1], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )

    g.nodes["user"].data["h"] = F.copy_to(
        F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["developer"].data["h"] = F.copy_to(
        F.tensor([3, 3], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 1, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    return g


def create_test_heterograph4(idtype):
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0, 1, 1, 2, 2, 2], dtype=idtype),
                F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype),
            ),
            ("user", "plays", "game"): (
                F.tensor([0, 1], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.copy_to(
        F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edges["follows"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4, 5, 6], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
    return g


def create_test_heterograph5(idtype):
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([1, 2], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
            ("user", "plays", "game"): (
                F.tensor([0, 1], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.copy_to(
        F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edges["follows"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
    return g


def get_redfn(name):
    return getattr(F, name)


@parametrize_idtype
def test_create(idtype):
    device = F.ctx()
    g0 = create_test_heterograph(idtype)
    g1 = create_test_heterograph1(idtype)
    g2 = create_test_heterograph2(idtype)
    assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
    assert (
        set(g0.canonical_etypes)
        == set(g1.canonical_etypes)
        == set(g2.canonical_etypes)
    )

    # Create a bipartite graph from a SciPy matrix
    src_ids = np.array([2, 3, 4])
    dst_ids = np.array([1, 2, 3])
    eweight = np.array([0.2, 0.3, 0.5])
    sp_mat = ssp.coo_matrix((eweight, (src_ids, dst_ids)))
    g = dgl.bipartite_from_scipy(
        sp_mat,
        utype="user",
        etype="plays",
        vtype="game",
        idtype=idtype,
        device=device,
    )
    assert g.idtype == idtype
    assert g.device == device
    assert g.num_src_nodes() == 5
    assert g.num_dst_nodes() == 4
    assert g.num_edges() == 3
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([2, 3, 4], dtype=idtype))
    assert F.allclose(dst, F.tensor([1, 2, 3], dtype=idtype))
    g = dgl.bipartite_from_scipy(
        sp_mat,
        utype="_U",
        etype="_E",
        vtype="_V",
        eweight_name="w",
        idtype=idtype,
        device=device,
    )
    assert F.allclose(g.edata["w"], F.tensor(eweight))

    # Create a bipartite graph from a NetworkX graph
    nx_g = nx.DiGraph()
    nx_g.add_nodes_from(
        [1, 3], bipartite=0, feat1=np.zeros((2)), feat2=np.ones((2))
    )
    nx_g.add_nodes_from([2, 4, 5], bipartite=1, feat3=np.zeros((3)))
    nx_g.add_edge(1, 4, weight=np.ones((1)), eid=np.array([1]))
    nx_g.add_edge(3, 5, weight=np.ones((1)), eid=np.array([0]))
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="user",
        etype="plays",
        vtype="game",
        idtype=idtype,
        device=device,
    )
    assert g.idtype == idtype
    assert g.device == device
    assert g.num_src_nodes() == 2
    assert g.num_dst_nodes() == 3
    assert g.num_edges() == 2
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([0, 1], dtype=idtype))
    assert F.allclose(dst, F.tensor([1, 2], dtype=idtype))
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="_U",
        etype="_E",
        vtype="V",
        u_attrs=["feat1", "feat2"],
        e_attrs=["weight"],
        v_attrs=["feat3"],
    )
    assert F.allclose(g.srcdata["feat1"], F.tensor(np.zeros((2, 2))))
    assert F.allclose(g.srcdata["feat2"], F.tensor(np.ones((2, 2))))
    assert F.allclose(g.dstdata["feat3"], F.tensor(np.zeros((3, 3))))
    assert F.allclose(g.edata["weight"], F.tensor(np.ones((2, 1))))
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="_U",
        etype="_E",
        vtype="V",
        edge_id_attr_name="eid",
        idtype=idtype,
        device=device,
    )
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([1, 0], dtype=idtype))
    assert F.allclose(dst, F.tensor([2, 1], dtype=idtype))

    # create from scipy
    spmat = ssp.coo_matrix(([1, 1, 1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
    g = dgl.from_scipy(spmat, idtype=idtype, device=device)
    assert g.num_nodes() == 4
    assert g.num_edges() == 3
    assert g.idtype == idtype
    assert g.device == device

    # test inferring number of nodes for heterograph
    g = dgl.heterograph(
        {
            ("l0", "e0", "l1"): ([0, 0], [1, 2]),
            ("l0", "e1", "l2"): ([2], [2]),
            ("l2", "e2", "l2"): ([1, 3], [1, 3]),
        },
        idtype=idtype,
        device=device,
    )
    assert g.num_nodes("l0") == 3
    assert g.num_nodes("l1") == 3
    assert g.num_nodes("l2") == 4
    assert g.idtype == idtype
    assert g.device == device

    # test if validate flag works
    # homo graph
    with pytest.raises(DGLError):
        g = dgl.graph(
            ([0, 0, 0, 1, 1, 2], [0, 1, 2, 0, 1, 2]),
            num_nodes=2,
            idtype=idtype,
            device=device,
        )

    # bipartite graph
    def _test_validate_bipartite(card):
        with pytest.raises(DGLError):
            g = dgl.heterograph(
                {("_U", "_E", "_V"): ([0, 0, 1, 1, 2], [1, 1, 2, 2, 3])},
                {"_U": card[0], "_V": card[1]},
                idtype=idtype,
                device=device,
            )

    _test_validate_bipartite((3, 3))
    _test_validate_bipartite((2, 4))

    # test from_scipy
    num_nodes = 10
    density = 0.25
    for fmt in ["csr", "coo", "csc"]:
        adj = rand(num_nodes, num_nodes, density=density, format=fmt)
        g = dgl.from_scipy(adj, eweight_name="w", idtype=idtype)
        assert g.idtype == idtype
        assert g.device == F.cpu()
        assert F.array_equal(
            g.edata["w"], F.copy_to(F.tensor(adj.data), F.cpu())
        )


def test_create2():
    mat = ssp.random(20, 30, 0.1)

    # coo
    mat = mat.tocoo()
    row = F.tensor(mat.row, dtype=F.int64)
    col = F.tensor(mat.col, dtype=F.int64)
    g = dgl.heterograph(
        {("A", "AB", "B"): ("coo", (row, col))},
        num_nodes_dict={"A": 20, "B": 30},
    )

    # csr
    mat = mat.tocsr()
    indptr = F.tensor(mat.indptr, dtype=F.int64)
    indices = F.tensor(mat.indices, dtype=F.int64)
    data = F.tensor([], dtype=F.int64)
    g = dgl.heterograph(
        {("A", "AB", "B"): ("csr", (indptr, indices, data))},
        num_nodes_dict={"A": 20, "B": 30},
    )

    # csc
    mat = mat.tocsc()
    indptr = F.tensor(mat.indptr, dtype=F.int64)
    indices = F.tensor(mat.indices, dtype=F.int64)
    data = F.tensor([], dtype=F.int64)
    g = dgl.heterograph(
        {("A", "AB", "B"): ("csc", (indptr, indices, data))},
        num_nodes_dict={"A": 20, "B": 30},
    )


@parametrize_idtype
def test_query(idtype):
    g = create_test_heterograph(idtype)

    ntypes = ["user", "game", "developer"]
    canonical_etypes = [
        ("user", "follows", "user"),
        ("user", "plays", "game"),
        ("user", "wishes", "game"),
        ("developer", "develops", "game"),
    ]
    etypes = ["follows", "plays", "wishes", "develops"]

    # node & edge types
    assert set(ntypes) == set(g.ntypes)
    assert set(etypes) == set(g.etypes)
    assert set(canonical_etypes) == set(g.canonical_etypes)

    # metagraph
    mg = g.metagraph()
    assert set(g.ntypes) == set(mg.nodes)
    etype_triplets = [(u, v, e) for u, v, e in mg.edges(keys=True)]
    assert set(
        [
            ("user", "user", "follows"),
            ("user", "game", "plays"),
            ("user", "game", "wishes"),
            ("developer", "game", "develops"),
        ]
    ) == set(etype_triplets)
    for i in range(len(etypes)):
        assert g.to_canonical_etype(etypes[i]) == canonical_etypes[i]

    def _test(g):
        # number of nodes
        assert [g.num_nodes(ntype) for ntype in ntypes] == [3, 2, 2]

        # number of edges
        assert [g.num_edges(etype) for etype in etypes] == [2, 4, 2, 2]

        # has_nodes
        for ntype in ntypes:
            n = g.number_of_nodes(ntype)
            for i in range(n):
                assert g.has_nodes(i, ntype)
            assert not g.has_nodes(n, ntype)
            assert np.array_equal(
                F.asnumpy(g.has_nodes([0, n], ntype)).astype("int32"), [1, 0]
            )

        assert not g.is_multigraph

        for etype in etypes:
            srcs, dsts = edges[etype]
            for src, dst in zip(srcs, dsts):
                assert g.has_edges_between(src, dst, etype)
            assert F.asnumpy(g.has_edges_between(srcs, dsts, etype)).all()

            srcs, dsts = negative_edges[etype]
            for src, dst in zip(srcs, dsts):
                assert not g.has_edges_between(src, dst, etype)
            assert not F.asnumpy(g.has_edges_between(srcs, dsts, etype)).any()

            srcs, dsts = edges[etype]
            n_edges = len(srcs)

            # predecessors & in_edges & in_degree
            pred = [s for s, d in zip(srcs, dsts) if d == 0]
            assert set(F.asnumpy(g.predecessors(0, etype)).tolist()) == set(
                pred
            )
            u, v = g.in_edges([0], etype=etype)
            assert F.asnumpy(v).tolist() == [0] * len(pred)
            assert set(F.asnumpy(u).tolist()) == set(pred)
            assert g.in_degrees(0, etype) == len(pred)

            # successors & out_edges & out_degree
            succ = [d for s, d in zip(srcs, dsts) if s == 0]
            assert set(F.asnumpy(g.successors(0, etype)).tolist()) == set(succ)
            u, v = g.out_edges([0], etype=etype)
            assert F.asnumpy(u).tolist() == [0] * len(succ)
            assert set(F.asnumpy(v).tolist()) == set(succ)
            assert g.out_degrees(0, etype) == len(succ)

            # edge_ids
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
                assert g.edge_ids(src, dst, etype=etype) == i
                _, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
                assert eid == i
            assert F.asnumpy(
                g.edge_ids(srcs, dsts, etype=etype)
            ).tolist() == list(range(n_edges))
            u, v, e = g.edge_ids(srcs, dsts, etype=etype, return_uv=True)
            u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
            assert u[e].tolist() == srcs
            assert v[e].tolist() == dsts

            # find_edges
            for eid in [
                list(range(n_edges)),
                np.arange(n_edges),
                F.astype(F.arange(0, n_edges), g.idtype),
            ]:
                u, v = g.find_edges(eid, etype)
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts

            # all_edges.
            for order in ["eid"]:
                u, v, e = g.edges("all", order, etype)
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts
                assert F.asnumpy(e).tolist() == list(range(n_edges))

            # in_degrees & out_degrees
            in_degrees = F.asnumpy(g.in_degrees(etype=etype))
            out_degrees = F.asnumpy(g.out_degrees(etype=etype))
            src_count = Counter(srcs)
            dst_count = Counter(dsts)
            utype, _, vtype = g.to_canonical_etype(etype)
            for i in range(g.number_of_nodes(utype)):
                assert out_degrees[i] == src_count[i]
            for i in range(g.number_of_nodes(vtype)):
                assert in_degrees[i] == dst_count[i]

    edges = {
        "follows": ([0, 1], [1, 2]),
        "plays": ([0, 1, 2, 1], [0, 0, 1, 1]),
        "wishes": ([0, 2], [1, 0]),
        "develops": ([0, 1], [0, 1]),
    }
    # edges that does not exist in the graph
    negative_edges = {
        "follows": ([0, 1], [0, 1]),
        "plays": ([0, 2], [1, 0]),
        "wishes": ([0, 1], [0, 1]),
        "develops": ([0, 1], [1, 0]),
    }
    g = create_test_heterograph(idtype)
    _test(g)
    g = create_test_heterograph1(idtype)
    _test(g)
    if F._default_context_str != "gpu":
        # XXX: CUDA COO operators have not been live yet.
        g = create_test_heterograph2(idtype)
        _test(g)

    etypes = canonical_etypes
    edges = {
        ("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]),
    }
    # edges that does not exist in the graph
    negative_edges = {
        ("user", "follows", "user"): ([0, 1], [0, 1]),
        ("user", "plays", "game"): ([0, 2], [1, 0]),
        ("user", "wishes", "game"): ([0, 1], [0, 1]),
        ("developer", "develops", "game"): ([0, 1], [1, 0]),
    }
    g = create_test_heterograph(idtype)
    _test(g)
    g = create_test_heterograph1(idtype)
    _test(g)
    if F._default_context_str != "gpu":
        # XXX: CUDA COO operators have not been live yet.
        g = create_test_heterograph2(idtype)
        _test(g)

    # test repr
    print(g)


@parametrize_idtype
def test_empty_query(idtype):
    g = dgl.graph(([1, 2, 3], [0, 4, 5]), idtype=idtype, device=F.ctx())
    g.add_nodes(0)
    g.add_edges([], [])
    g.remove_edges([])
    g.remove_nodes([])
    assert F.shape(g.has_nodes([])) == (0,)
    assert F.shape(g.has_edges_between([], [])) == (0,)
    g.edge_ids([], [])
    g.edge_ids([], [], return_uv=True)
    g.find_edges([])

    assert F.shape(g.in_edges([], form="eid")) == (0,)
    u, v = g.in_edges([], form="uv")
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    u, v, e = g.in_edges([], form="all")
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    assert F.shape(e) == (0,)

    assert F.shape(g.out_edges([], form="eid")) == (0,)
    u, v = g.out_edges([], form="uv")
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    u, v, e = g.out_edges([], form="all")
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    assert F.shape(e) == (0,)

    assert F.shape(g.in_degrees([])) == (0,)
    assert F.shape(g.out_degrees([])) == (0,)

    g = dgl.graph(([], []), idtype=idtype, device=F.ctx())
    error_thrown = True
    try:
        g.in_degrees([0])
        fail = False
    except:
        pass
    assert error_thrown
    error_thrown = True
    try:
        g.out_degrees([0])
        fail = False
    except:
        pass
    assert error_thrown


@unittest.skipIf(
    F._default_context_str == "gpu", reason="GPU does not have COO impl."
)
def _test_hypersparse():
    N1 = 1 << 50  # should crash if allocated a CSR
    N2 = 1 << 48

    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0], F.int64),
                F.tensor([1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
        device=F.ctx(),
    )
    assert g.number_of_nodes("user") == N1
    assert g.number_of_nodes("game") == N1
    assert g.number_of_edges("follows") == 1
    assert g.number_of_edges("plays") == 1

    assert g.has_edges_between(0, 1, "follows")
    assert not g.has_edges_between(0, 0, "follows")
    mask = F.asnumpy(g.has_edges_between([0, 0], [0, 1], "follows")).tolist()
    assert mask == [0, 1]

    assert g.has_edges_between(0, N2, "plays")
    assert not g.has_edges_between(0, 0, "plays")
    mask = F.asnumpy(g.has_edges_between([0, 0], [0, N2], "plays")).tolist()
    assert mask == [0, 1]

    assert F.asnumpy(g.predecessors(0, "follows")).tolist() == []
    assert F.asnumpy(g.successors(0, "follows")).tolist() == [1]
    assert F.asnumpy(g.predecessors(1, "follows")).tolist() == [0]
    assert F.asnumpy(g.successors(1, "follows")).tolist() == []

    assert F.asnumpy(g.predecessors(0, "plays")).tolist() == []
    assert F.asnumpy(g.successors(0, "plays")).tolist() == [N2]
    assert F.asnumpy(g.predecessors(N2, "plays")).tolist() == [0]
    assert F.asnumpy(g.successors(N2, "plays")).tolist() == []

    assert g.edge_ids(0, 1, etype="follows") == 0
    assert g.edge_ids(0, N2, etype="plays") == 0

    u, v = g.find_edges([0], "follows")
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [1]
    u, v = g.find_edges([0], "plays")
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [N2]
    u, v, e = g.all_edges("all", "eid", "follows")
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [1]
    assert F.asnumpy(e).tolist() == [0]
    u, v, e = g.all_edges("all", "eid", "plays")
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [N2]
    assert F.asnumpy(e).tolist() == [0]

    assert g.in_degrees(0, "follows") == 0
    assert g.in_degrees(1, "follows") == 1
    assert F.asnumpy(g.in_degrees([0, 1], "follows")).tolist() == [0, 1]
    assert g.in_degrees(0, "plays") == 0
    assert g.in_degrees(N2, "plays") == 1
    assert F.asnumpy(g.in_degrees([0, N2], "plays")).tolist() == [0, 1]
    assert g.out_degrees(0, "follows") == 1
    assert g.out_degrees(1, "follows") == 0
    assert F.asnumpy(g.out_degrees([0, 1], "follows")).tolist() == [1, 0]
    assert g.out_degrees(0, "plays") == 1
    assert g.out_degrees(N2, "plays") == 0
    assert F.asnumpy(g.out_degrees([0, N2], "plays")).tolist() == [1, 0]


def _test_edge_ids():
    N1 = 1 << 50  # should crash if allocated a CSR
    N2 = 1 << 48

    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0], F.int64),
                F.tensor([1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
    )
    with pytest.raises(DGLError):
        eid = g.edge_ids(0, 0, etype="follows")

    g2 = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0, 0], F.int64),
                F.tensor([1, 1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
        device=F.cpu(),
    )

    eid = g2.edge_ids(0, 1, etype="follows")
    assert eid == 0


@parametrize_idtype
def test_adj(idtype):
    g = create_test_heterograph(idtype)
    adj = F.sparse_to_numpy(g.adj(transpose=True, etype="follows"))
    assert np.allclose(
        adj, np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    )
    adj = F.sparse_to_numpy(g.adj(transpose=False, etype="follows"))
    assert np.allclose(
        adj, np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]])
    )
    adj = F.sparse_to_numpy(g.adj(transpose=True, etype="plays"))
    assert np.allclose(adj, np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]]))
    adj = F.sparse_to_numpy(g.adj(transpose=False, etype="plays"))
    assert np.allclose(adj, np.array([[1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]))

    adj = g.adj(transpose=True, scipy_fmt="csr", etype="follows")
    assert np.allclose(
        adj.todense(),
        np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
    )
    adj = g.adj(transpose=True, scipy_fmt="coo", etype="follows")
    assert np.allclose(
        adj.todense(),
        np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
    )
    adj = g.adj(transpose=True, scipy_fmt="csr", etype="plays")
    assert np.allclose(
        adj.todense(), np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]])
    )
    adj = g.adj(transpose=True, scipy_fmt="coo", etype="plays")
    assert np.allclose(
        adj.todense(), np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]])
    )
    adj = F.sparse_to_numpy(g["follows"].adj(transpose=True))
    assert np.allclose(
        adj, np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    )


@parametrize_idtype
def test_inc(idtype):
    g = create_test_heterograph(idtype)
    adj = F.sparse_to_numpy(g["follows"].inc("in"))
    assert np.allclose(adj, np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]))
    adj = F.sparse_to_numpy(g["follows"].inc("out"))
    assert np.allclose(adj, np.array([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]]))
    adj = F.sparse_to_numpy(g["follows"].inc("both"))
    assert np.allclose(adj, np.array([[-1.0, 0.0], [1.0, -1.0], [0.0, 1.0]]))
    adj = F.sparse_to_numpy(g.inc("in", etype="plays"))
    assert np.allclose(
        adj, np.array([[1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 1.0]])
    )
    adj = F.sparse_to_numpy(g.inc("out", etype="plays"))
    assert np.allclose(
        adj,
        np.array(
            [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]
        ),
    )
    adj = F.sparse_to_numpy(g.inc("both", etype="follows"))
    assert np.allclose(adj, np.array([[-1.0, 0.0], [1.0, -1.0], [0.0, 1.0]]))


@parametrize_idtype
def test_view(idtype):
    # test single node type
    g = dgl.heterograph(
        {("user", "follows", "user"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    f1 = F.randn((3, 6))
    g.ndata["h"] = f1
    f2 = g.nodes["user"].data["h"]
    assert F.array_equal(f1, f2)
    fail = False
    try:
        g.ndata["h"] = {"user": f1}
    except Exception:
        fail = True
    assert fail

    # test single edge type
    f3 = F.randn((2, 4))
    g.edata["h"] = f3
    f4 = g.edges["follows"].data["h"]
    assert F.array_equal(f3, f4)
    fail = False
    try:
        g.edata["h"] = {"follows": f3}
    except Exception:
        fail = True
    assert fail

    # test data view
    g = create_test_heterograph(idtype)

    f1 = F.randn((3, 6))
    g.nodes["user"].data["h"] = f1  # ok
    f2 = g.nodes["user"].data["h"]
    assert F.array_equal(f1, f2)
    assert F.array_equal(g.nodes("user"), F.arange(0, 3, idtype))
    g.nodes["user"].data.pop("h")

    # multi type ndata
    f1 = F.randn((3, 6))
    f2 = F.randn((2, 6))
    fail = False
    try:
        g.ndata["h"] = f1
    except Exception:
        fail = True
    assert fail

    f3 = F.randn((2, 4))
    g.edges["user", "follows", "user"].data["h"] = f3
    f4 = g.edges["user", "follows", "user"].data["h"]
    f5 = g.edges["follows"].data["h"]
    assert F.array_equal(f3, f4)
    assert F.array_equal(f3, f5)
    assert F.array_equal(
        g.edges(etype="follows", form="eid"), F.arange(0, 2, idtype)
    )
    g.edges["follows"].data.pop("h")

    f3 = F.randn((2, 4))
    fail = False
    try:
        g.edata["h"] = f3
    except Exception:
        fail = True
    assert fail

    # test srcdata
    f1 = F.randn((3, 6))
    g.srcnodes["user"].data["h"] = f1  # ok
    f2 = g.srcnodes["user"].data["h"]
    assert F.array_equal(f1, f2)
    assert F.array_equal(g.srcnodes("user"), F.arange(0, 3, idtype))
    g.srcnodes["user"].data.pop("h")

    # test dstdata
    f1 = F.randn((3, 6))
    g.dstnodes["user"].data["h"] = f1  # ok
    f2 = g.dstnodes["user"].data["h"]
    assert F.array_equal(f1, f2)
    assert F.array_equal(g.dstnodes("user"), F.arange(0, 3, idtype))
    g.dstnodes["user"].data.pop("h")


@parametrize_idtype
def test_view1(idtype):
    # test relation view
    HG = create_test_heterograph(idtype)
    ntypes = ["user", "game", "developer"]
    canonical_etypes = [
        ("user", "follows", "user"),
        ("user", "plays", "game"),
        ("user", "wishes", "game"),
        ("developer", "develops", "game"),
    ]
    etypes = ["follows", "plays", "wishes", "develops"]

    def _test_query():
        for etype in etypes:
            utype, _, vtype = HG.to_canonical_etype(etype)
            g = HG[etype]
            srcs, dsts = edges[etype]
            for src, dst in zip(srcs, dsts):
                assert g.has_edges_between(src, dst)
            assert F.asnumpy(g.has_edges_between(srcs, dsts)).all()

            srcs, dsts = negative_edges[etype]
            for src, dst in zip(srcs, dsts):
                assert not g.has_edges_between(src, dst)
            assert not F.asnumpy(g.has_edges_between(srcs, dsts)).any()

            srcs, dsts = edges[etype]
            n_edges = len(srcs)

            # predecessors & in_edges & in_degree
            pred = [s for s, d in zip(srcs, dsts) if d == 0]
            assert set(F.asnumpy(g.predecessors(0)).tolist()) == set(pred)
            u, v = g.in_edges([0])
            assert F.asnumpy(v).tolist() == [0] * len(pred)
            assert set(F.asnumpy(u).tolist()) == set(pred)
            assert g.in_degrees(0) == len(pred)

            # successors & out_edges & out_degree
            succ = [d for s, d in zip(srcs, dsts) if s == 0]
            assert set(F.asnumpy(g.successors(0)).tolist()) == set(succ)
            u, v = g.out_edges([0])
            assert F.asnumpy(u).tolist() == [0] * len(succ)
            assert set(F.asnumpy(v).tolist()) == set(succ)
            assert g.out_degrees(0) == len(succ)

            # edge_ids
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
                assert g.edge_ids(src, dst, etype=etype) == i
                _, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
                assert eid == i
            assert F.asnumpy(g.edge_ids(srcs, dsts)).tolist() == list(
                range(n_edges)
            )
            u, v, e = g.edge_ids(srcs, dsts, return_uv=True)
            u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
            assert u[e].tolist() == srcs
            assert v[e].tolist() == dsts

            # find_edges
            u, v = g.find_edges(list(range(n_edges)))
            assert F.asnumpy(u).tolist() == srcs
            assert F.asnumpy(v).tolist() == dsts

            # all_edges.
            for order in ["eid"]:
                u, v, e = g.all_edges(form="all", order=order)
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts
                assert F.asnumpy(e).tolist() == list(range(n_edges))

            # in_degrees & out_degrees
            in_degrees = F.asnumpy(g.in_degrees())
            out_degrees = F.asnumpy(g.out_degrees())
            src_count = Counter(srcs)
            dst_count = Counter(dsts)
            for i in range(g.number_of_nodes(utype)):
                assert out_degrees[i] == src_count[i]
            for i in range(g.number_of_nodes(vtype)):
                assert in_degrees[i] == dst_count[i]

    edges = {
        "follows": ([0, 1], [1, 2]),
        "plays": ([0, 1, 2, 1], [0, 0, 1, 1]),
        "wishes": ([0, 2], [1, 0]),
        "develops": ([0, 1], [0, 1]),
    }
    # edges that does not exist in the graph
    negative_edges = {
        "follows": ([0, 1], [0, 1]),
        "plays": ([0, 2], [1, 0]),
        "wishes": ([0, 1], [0, 1]),
        "develops": ([0, 1], [1, 0]),
    }
    _test_query()
    etypes = canonical_etypes
    edges = {
        ("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]),
    }
    # edges that does not exist in the graph
    negative_edges = {
        ("user", "follows", "user"): ([0, 1], [0, 1]),
        ("user", "plays", "game"): ([0, 2], [1, 0]),
        ("user", "wishes", "game"): ([0, 1], [0, 1]),
        ("developer", "develops", "game"): ([0, 1], [1, 0]),
    }
    _test_query()

    # test features
    HG.nodes["user"].data["h"] = F.ones((HG.number_of_nodes("user"), 5))
    HG.nodes["game"].data["m"] = F.ones((HG.number_of_nodes("game"), 3)) * 2

    # test only one node type
    g = HG["follows"]
    assert g.number_of_nodes() == 3

    # test ndata and edata
    f1 = F.randn((3, 6))
    g.ndata["h"] = f1  # ok
    f2 = HG.nodes["user"].data["h"]
    assert F.array_equal(f1, f2)
    assert F.array_equal(g.nodes(), F.arange(0, 3, g.idtype))

    f3 = F.randn((2, 4))
    g.edata["h"] = f3
    f4 = HG.edges["follows"].data["h"]
    assert F.array_equal(f3, f4)
    assert F.array_equal(g.edges(form="eid"), F.arange(0, 2, g.idtype))


@parametrize_idtype
def test_flatten(idtype):
    def check_mapping(g, fg):
        if len(fg.ntypes) == 1:
            SRC = DST = fg.ntypes[0]
        else:
            SRC = fg.ntypes[0]
            DST = fg.ntypes[1]

        etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
        eids = F.asnumpy(fg.edata[dgl.EID]).tolist()

        for i, (etype, eid) in enumerate(zip(etypes, eids)):
            src_g, dst_g = g.find_edges([eid], g.canonical_etypes[etype])
            src_fg, dst_fg = fg.find_edges([i])
            # TODO(gq): I feel this code is quite redundant; can we just add new members (like
            # "induced_srcid") to returned heterograph object and not store them as features?
            assert F.asnumpy(src_g) == F.asnumpy(
                F.gather_row(fg.nodes[SRC].data[dgl.NID], src_fg)[0]
            )
            tid = F.asnumpy(
                F.gather_row(fg.nodes[SRC].data[dgl.NTYPE], src_fg)
            ).item()
            assert g.canonical_etypes[etype][0] == g.ntypes[tid]
            assert F.asnumpy(dst_g) == F.asnumpy(
                F.gather_row(fg.nodes[DST].data[dgl.NID], dst_fg)[0]
            )
            tid = F.asnumpy(
                F.gather_row(fg.nodes[DST].data[dgl.NTYPE], dst_fg)
            ).item()
            assert g.canonical_etypes[etype][2] == g.ntypes[tid]

    # check for wildcard slices
    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.edges["wishes"].data["e"] = F.ones((2, 4))
    g.edges["wishes"].data["f"] = F.ones((2, 4))

    fg = g["user", :, "game"]  # user--plays->game and user--wishes->game
    assert len(fg.ntypes) == 2
    assert fg.ntypes == ["user", "game"]
    assert fg.etypes == ["plays+wishes"]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    etype = fg.etypes[0]
    assert fg[etype] is not None  # Issue #2166

    assert F.array_equal(fg.nodes["user"].data["h"], F.ones((3, 5)))
    assert F.array_equal(fg.nodes["game"].data["i"], F.ones((2, 5)))
    assert F.array_equal(fg.edata["e"], F.ones((6, 4)))
    assert "f" not in fg.edata

    etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
    eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
    assert set(zip(etypes, eids)) == set(
        [(3, 0), (3, 1), (2, 1), (2, 0), (2, 3), (2, 2)]
    )

    check_mapping(g, fg)

    fg = g["user", :, "user"]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    # NOTE(gq): The node/edge types from the parent graph is returned if there is only one
    # node/edge type.  This differs from the behavior above.
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows"]
    u1, v1 = g.edges(etype="follows", order="eid")
    u2, v2 = fg.edges(etype="follows", order="eid")
    assert F.array_equal(u1, u2)
    assert F.array_equal(v1, v2)

    fg = g["developer", :, "game"]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    assert fg.ntypes == ["developer", "game"]
    assert fg.etypes == ["develops"]
    u1, v1 = g.edges(etype="develops", order="eid")
    u2, v2 = fg.edges(etype="develops", order="eid")
    assert F.array_equal(u1, u2)
    assert F.array_equal(v1, v2)

    fg = g[:, :, :]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    assert fg.ntypes == ["developer+user", "game+user"]
    assert fg.etypes == ["develops+follows+plays+wishes"]
    check_mapping(g, fg)

    # Test another heterograph
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1, 2], [1, 2, 3]),
            ("user", "knows", "user"): ([0, 2], [2, 3]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.randn((4, 3))
    g.edges["follows"].data["w"] = F.randn((3, 2))
    g.nodes["user"].data["hh"] = F.randn((4, 5))
    g.edges["knows"].data["ww"] = F.randn((2, 10))

    fg = g["user", :, "user"]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows+knows"]
    check_mapping(g, fg)

    fg = g["user", :, :]
    assert fg.idtype == g.idtype
    assert fg.device == g.device
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows+knows"]
    check_mapping(g, fg)


@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
@parametrize_idtype
def test_to_device(idtype):
    # TODO: rewrite this test case to accept different graphs so we
    #  can test reverse graph and batched graph
    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))
    assert g.device == F.ctx()
    g = g.to(F.cpu())
    assert g.device == F.cpu()
    assert F.context(g.nodes["user"].data["h"]) == F.cpu()
    assert F.context(g.nodes["game"].data["i"]) == F.cpu()
    assert F.context(g.edges["plays"].data["e"]) == F.cpu()
    for ntype in g.ntypes:
        assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
    for etype in g.canonical_etypes:
        assert F.context(g.batch_num_edges(etype)) == F.cpu()

    if F.is_cuda_available():
        g1 = g.to(F.cuda())
        assert g1.device == F.cuda()
        assert F.context(g1.nodes["user"].data["h"]) == F.cuda()
        assert F.context(g1.nodes["game"].data["i"]) == F.cuda()
        assert F.context(g1.edges["plays"].data["e"]) == F.cuda()
        for ntype in g1.ntypes:
            assert F.context(g1.batch_num_nodes(ntype)) == F.cuda()
        for etype in g1.canonical_etypes:
            assert F.context(g1.batch_num_edges(etype)) == F.cuda()
        assert F.context(g.nodes["user"].data["h"]) == F.cpu()
        assert F.context(g.nodes["game"].data["i"]) == F.cpu()
        assert F.context(g.edges["plays"].data["e"]) == F.cpu()
        for ntype in g.ntypes:
            assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
        for etype in g.canonical_etypes:
            assert F.context(g.batch_num_edges(etype)) == F.cpu()
        with pytest.raises(DGLError):
            g1.nodes["user"].data["h"] = F.copy_to(F.ones((3, 5)), F.cpu())
        with pytest.raises(DGLError):
            g1.edges["plays"].data["e"] = F.copy_to(F.ones((4, 4)), F.cpu())


@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
@parametrize_idtype
@pytest.mark.parametrize("g", get_cases(["block"]))
def test_to_device2(g, idtype):
    g = g.astype(idtype)
    g = g.to(F.cpu())
    assert g.device == F.cpu()
    if F.is_cuda_available():
        g1 = g.to(F.cuda())
        assert g1.device == F.cuda()
        assert g1.ntypes == g.ntypes
        assert g1.etypes == g.etypes
        assert g1.canonical_etypes == g.canonical_etypes


@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
@unittest.skipIf(
    dgl.backend.backend_name != "pytorch",
    reason="Pinning graph inplace only supported for PyTorch",
)
@parametrize_idtype
def test_pin_memory_(idtype):
    # TODO: rewrite this test case to accept different graphs so we
    #  can test reverse graph and batched graph
    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()

    # unpin an unpinned CPU graph, directly return
    g.unpin_memory_()
    assert not g.is_pinned()
    assert g.device == F.cpu()

    # pin a CPU graph
    g.pin_memory_()
    assert g.is_pinned()
    assert g.device == F.cpu()
    assert g.nodes["user"].data["h"].is_pinned()
    assert g.nodes["game"].data["i"].is_pinned()
    assert g.edges["plays"].data["e"].is_pinned()
    assert F.context(g.nodes["user"].data["h"]) == F.cpu()
    assert F.context(g.nodes["game"].data["i"]) == F.cpu()
    assert F.context(g.edges["plays"].data["e"]) == F.cpu()
    for ntype in g.ntypes:
        assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
    for etype in g.canonical_etypes:
        assert F.context(g.batch_num_edges(etype)) == F.cpu()

    # it's fine to clone with new formats, but new graphs are not pinned
    # >>> g.formats()
    # {'created': ['coo'], 'not created': ['csr', 'csc']}
    assert not g.formats("csc").is_pinned()
    assert not g.formats("csr").is_pinned()
    # 'coo' formats is already created and thus not cloned
    assert g.formats("coo").is_pinned()

    # pin a pinned graph, directly return
    g.pin_memory_()
    assert g.is_pinned()
    assert g.device == F.cpu()

    # unpin a pinned graph
    g.unpin_memory_()
    assert not g.is_pinned()
    assert g.device == F.cpu()

    g1 = g.to(F.cuda())

    # unpin an unpinned GPU graph, directly return
    g1.unpin_memory_()
    assert not g1.is_pinned()
    assert g1.device == F.cuda()

    # error pinning a GPU graph
    with pytest.raises(DGLError):
        g1.pin_memory_()

    # test pin empty homograph
    g2 = dgl.graph(([], []))
    g2.pin_memory_()
    assert g2.is_pinned()
    g2.unpin_memory_()
    assert not g2.is_pinned()

    # test pin heterograph with 0 edge of one relation type
    g3 = dgl.heterograph(
        {("a", "b", "c"): ([0, 1], [1, 2]), ("c", "d", "c"): ([], [])}
    ).astype(idtype)
    g3.pin_memory_()
    assert g3.is_pinned()
    g3.unpin_memory_()
    assert not g3.is_pinned()


@parametrize_idtype
def test_convert_bound(idtype):
    def _test_bipartite_bound(data, card):
        with pytest.raises(DGLError):
            dgl.heterograph(
                {("_U", "_E", "_V"): data},
                {"_U": card[0], "_V": card[1]},
                idtype=idtype,
                device=F.ctx(),
            )

    def _test_graph_bound(data, card):
        with pytest.raises(DGLError):
            dgl.graph(data, num_nodes=card, idtype=idtype, device=F.ctx())

    _test_bipartite_bound(([1, 2], [1, 2]), (2, 3))
    _test_bipartite_bound(([0, 1], [1, 4]), (2, 3))
    _test_graph_bound(([1, 3], [1, 2]), 3)
    _test_graph_bound(([0, 1], [1, 3]), 3)


@parametrize_idtype
def test_convert(idtype):
    hg = create_test_heterograph(idtype)
    hs = []
    for ntype in hg.ntypes:
        h = F.randn((hg.number_of_nodes(ntype), 5))
        hg.nodes[ntype].data["h"] = h
        hs.append(h)
    hg.nodes["user"].data["x"] = F.randn((3, 3))
    ws = []
    for etype in hg.canonical_etypes:
        w = F.randn((hg.number_of_edges(etype), 5))
        hg.edges[etype].data["w"] = w
        ws.append(w)
    hg.edges["plays"].data["x"] = F.randn((4, 3))

    g = dgl.to_homogeneous(hg, ndata=["h"], edata=["w"])
    assert g.idtype == idtype
    assert g.device == hg.device
    assert F.array_equal(F.cat(hs, dim=0), g.ndata["h"])
    assert "x" not in g.ndata
    assert F.array_equal(F.cat(ws, dim=0), g.edata["w"])
    assert "x" not in g.edata

    src, dst = g.all_edges(order="eid")
    src = F.asnumpy(src)
    dst = F.asnumpy(dst)
    etype_id, eid = F.asnumpy(g.edata[dgl.ETYPE]), F.asnumpy(g.edata[dgl.EID])
    ntype_id, nid = F.asnumpy(g.ndata[dgl.NTYPE]), F.asnumpy(g.ndata[dgl.NID])
    for i in range(g.number_of_edges()):
        srctype = hg.ntypes[ntype_id[src[i]]]
        dsttype = hg.ntypes[ntype_id[dst[i]]]
        etype = hg.etypes[etype_id[i]]
        src_i, dst_i = hg.find_edges([eid[i]], (srctype, etype, dsttype))
        assert np.asscalar(F.asnumpy(src_i)) == nid[src[i]]
        assert np.asscalar(F.asnumpy(dst_i)) == nid[dst[i]]

    mg = nx.MultiDiGraph(
        [
            ("user", "user", "follows"),
            ("user", "game", "plays"),
            ("user", "game", "wishes"),
            ("developer", "game", "develops"),
        ]
    )

    for _mg in [None, mg]:
        hg2 = dgl.to_heterogeneous(
            g,
            hg.ntypes,
            hg.etypes,
            ntype_field=dgl.NTYPE,
            etype_field=dgl.ETYPE,
            metagraph=_mg,
        )
        assert hg2.idtype == hg.idtype
        assert hg2.device == hg.device
        assert set(hg.ntypes) == set(hg2.ntypes)
        assert set(hg.canonical_etypes) == set(hg2.canonical_etypes)
        for ntype in hg.ntypes:
            assert hg.number_of_nodes(ntype) == hg2.number_of_nodes(ntype)
            assert F.array_equal(
                hg.nodes[ntype].data["h"], hg2.nodes[ntype].data["h"]
            )
        for canonical_etype in hg.canonical_etypes:
            src, dst = hg.all_edges(etype=canonical_etype, order="eid")
            src2, dst2 = hg2.all_edges(etype=canonical_etype, order="eid")
            assert F.array_equal(src, src2)
            assert F.array_equal(dst, dst2)
            assert F.array_equal(
                hg.edges[canonical_etype].data["w"],
                hg2.edges[canonical_etype].data["w"],
            )

    # hetero_from_homo test case 2
    g = dgl.graph(([0, 1, 2, 0], [2, 2, 3, 3]), idtype=idtype, device=F.ctx())
    g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
    g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
    hg = dgl.to_heterogeneous(g, ["l0", "l1", "l2"], ["e0", "e1", "e2"])
    assert hg.idtype == idtype
    assert hg.device == g.device
    assert set(hg.canonical_etypes) == set(
        [("l0", "e0", "l1"), ("l1", "e1", "l2"), ("l0", "e2", "l2")]
    )
    assert hg.number_of_nodes("l0") == 2
    assert hg.number_of_nodes("l1") == 1
    assert hg.number_of_nodes("l2") == 1
    assert hg.number_of_edges("e0") == 2
    assert hg.number_of_edges("e1") == 1
    assert hg.number_of_edges("e2") == 1
    assert F.array_equal(hg.ndata[dgl.NID]["l0"], F.tensor([0, 1], F.int64))
    assert F.array_equal(hg.ndata[dgl.NID]["l1"], F.tensor([2], F.int64))
    assert F.array_equal(hg.ndata[dgl.NID]["l2"], F.tensor([3], F.int64))
    assert F.array_equal(
        hg.edata[dgl.EID][("l0", "e0", "l1")], F.tensor([0, 1], F.int64)
    )
    assert F.array_equal(
        hg.edata[dgl.EID][("l0", "e2", "l2")], F.tensor([3], F.int64)
    )
    assert F.array_equal(
        hg.edata[dgl.EID][("l1", "e1", "l2")], F.tensor([2], F.int64)
    )

    # hetero_from_homo test case 3
    mg = nx.MultiDiGraph(
        [("user", "movie", "watches"), ("user", "TV", "watches")]
    )
    g = dgl.graph(((0, 0), (1, 2)), idtype=idtype, device=F.ctx())
    g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
    g.edata[dgl.ETYPE] = F.tensor([0, 0])
    for _mg in [None, mg]:
        hg = dgl.to_heterogeneous(
            g, ["user", "TV", "movie"], ["watches"], metagraph=_mg
        )
        assert hg.idtype == g.idtype
        assert hg.device == g.device
        assert set(hg.canonical_etypes) == set(
            [("user", "watches", "movie"), ("user", "watches", "TV")]
        )
        assert hg.number_of_nodes("user") == 1
        assert hg.number_of_nodes("TV") == 1
        assert hg.number_of_nodes("movie") == 1
        assert hg.number_of_edges(("user", "watches", "TV")) == 1
        assert hg.number_of_edges(("user", "watches", "movie")) == 1
        assert len(hg.etypes) == 2

    # hetero_to_homo test case 2
    hg = dgl.heterograph(
        {("_U", "_E", "_V"): ([0, 1], [0, 1])},
        {"_U": 2, "_V": 3},
        idtype=idtype,
        device=F.ctx(),
    )
    g = dgl.to_homogeneous(hg)
    assert hg.idtype == g.idtype
    assert hg.device == g.device
    assert g.number_of_nodes() == 5

    # hetero_to_subgraph_to_homo
    hg = 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(),
    )
    hg.nodes["user"].data["h"] = F.copy_to(
        F.tensor([[1, 0], [0, 1], [1, 1]], dtype=idtype), ctx=F.ctx()
    )
    sg = dgl.node_subgraph(hg, {"user": [1, 2]})
    assert len(sg.ntypes) == 2
    assert len(sg.etypes) == 2
    assert sg.num_nodes("user") == 2
    assert sg.num_nodes("game") == 0
    g = dgl.to_homogeneous(sg, ndata=["h"])
    assert "h" in g.ndata.keys()
    assert g.num_nodes() == 2


@unittest.skipIf(
    F._default_context_str == "gpu", reason="Test on cpu is enough"
)
@parametrize_idtype
def test_to_homo_zero_nodes(idtype):
    # Fix gihub issue #2870
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): (
                np.random.randint(0, 200, (1000,)),
                np.random.randint(0, 200, (1000,)),
            ),
            ("B", "BA", "A"): (
                np.random.randint(0, 200, (1000,)),
                np.random.randint(0, 200, (1000,)),
            ),
        },
        num_nodes_dict={"A": 200, "B": 200, "C": 0},
        idtype=idtype,
    )
    g.nodes["A"].data["x"] = F.randn((200, 3))
    g.nodes["B"].data["x"] = F.randn((200, 3))
    gg = dgl.to_homogeneous(g, ["x"])
    assert "x" in gg.ndata


@parametrize_idtype
def test_to_homo2(idtype):
    # test the result homogeneous graph has nodes and edges sorted by their types
    hg = create_test_heterograph(idtype)
    g = dgl.to_homogeneous(hg)
    ntypes = F.asnumpy(g.ndata[dgl.NTYPE])
    etypes = F.asnumpy(g.edata[dgl.ETYPE])
    p = 0
    for tid, ntype in enumerate(hg.ntypes):
        num_nodes = hg.num_nodes(ntype)
        for i in range(p, p + num_nodes):
            assert ntypes[i] == tid
        p += num_nodes
    p = 0
    for tid, etype in enumerate(hg.canonical_etypes):
        num_edges = hg.num_edges(etype)
        for i in range(p, p + num_edges):
            assert etypes[i] == tid
        p += num_edges
    # test store_type=False
    g = dgl.to_homogeneous(hg, store_type=False)
    assert dgl.NTYPE not in g.ndata
    assert dgl.ETYPE not in g.edata
    # test return_count=True
    g, ntype_count, etype_count = dgl.to_homogeneous(hg, return_count=True)
    for i, count in enumerate(ntype_count):
        assert count == hg.num_nodes(hg.ntypes[i])
    for i, count in enumerate(etype_count):
        assert count == hg.num_edges(hg.canonical_etypes[i])


@parametrize_idtype
def test_invertible_conversion(idtype):
    # Test whether to_homogeneous and to_heterogeneous are invertible
    hg = create_test_heterograph(idtype)
    g = dgl.to_homogeneous(hg)
    hg2 = dgl.to_heterogeneous(g, hg.ntypes, hg.etypes)
    assert_is_identical_hetero(hg, hg2, True)


@parametrize_idtype
def test_metagraph_reachable(idtype):
    g = create_test_heterograph(idtype)
    x = F.randn((3, 5))
    g.nodes["user"].data["h"] = x

    new_g = dgl.metapath_reachable_graph(g, ["follows", "plays"])
    assert new_g.idtype == idtype
    assert new_g.ntypes == ["game", "user"]
    assert new_g.number_of_edges() == 3
    assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all()

    new_g = dgl.metapath_reachable_graph(g, ["follows"])
    assert new_g.idtype == idtype
    assert new_g.ntypes == ["user"]
    assert new_g.number_of_edges() == 2
    assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all()


@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.number_of_nodes("developer") == 0
        assert sg.number_of_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)
    if F._default_context_str != "gpu":
        # TODO(minjie): enable this later
        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_subgraph(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.number_of_nodes("developer") == 0
        assert sg.number_of_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)
    if F._default_context_str != "gpu":
        # TODO(minjie): enable this later
        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)
    if F._default_context_str != "gpu":
        # TODO(minjie): enable this later
        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)
    if F._default_context_str != "gpu":
        # TODO(minjie): enable this later
        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.number_of_nodes(ntype) == sg.number_of_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.number_of_nodes(ntype) == sg.number_of_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)
    if F._default_context_str != "gpu":
        # TODO(minjie): enable this later
        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.number_of_nodes(ntype) == g.number_of_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.number_of_nodes(ntype) == g.number_of_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)


@parametrize_idtype
def test_apply(idtype):
    def node_udf(nodes):
        return {"h": nodes.data["h"] * 2}

    def node_udf2(nodes):
        return {"h": F.sum(nodes.data["h"], dim=1, keepdims=True)}

    def edge_udf(edges):
        return {"h": edges.data["h"] * 2 + edges.src["h"]}

    g = create_test_heterograph(idtype)
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.apply_nodes(node_udf, ntype="user")
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 2)

    g["plays"].edata["h"] = F.ones((4, 5))
    g.apply_edges(edge_udf, etype=("user", "plays", "game"))
    assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 4)

    # test apply on graph with only one type
    g["follows"].apply_nodes(node_udf)
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 4)

    g["plays"].apply_edges(edge_udf)
    assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 12)

    # Test the case that feature size changes
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.apply_nodes(node_udf2, ntype="user")
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 1)) * 5)

    # test fail case
    # fail due to multiple types
    with pytest.raises(DGLError):
        g.apply_nodes(node_udf)

    with pytest.raises(DGLError):
        g.apply_edges(edge_udf)


@parametrize_idtype
def test_level2(idtype):
    # edges = {
    #    'follows': ([0, 1], [1, 2]),
    #    'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
    #    'wishes': ([0, 2], [1, 0]),
    #    'develops': ([0, 1], [0, 1]),
    # }
    g = create_test_heterograph(idtype)

    def rfunc(nodes):
        return {"y": F.sum(nodes.mailbox["m"], 1)}

    def rfunc2(nodes):
        return {"y": F.max(nodes.mailbox["m"], 1)}

    def mfunc(edges):
        return {"m": edges.src["h"]}

    def afunc(nodes):
        return {"y": nodes.data["y"] + 1}

    #############################################################
    #  send_and_recv
    #############################################################

    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.send_and_recv([2, 3], mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # only one type
    g["plays"].send_and_recv([2, 3], mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # test fail case
    # fail due to multiple types
    with pytest.raises(DGLError):
        g.send_and_recv([2, 3], mfunc, rfunc)

    g.nodes["game"].data.clear()

    #############################################################
    #  pull
    #############################################################

    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.pull(1, mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # only one type
    g["plays"].pull(1, mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))

    # test fail case
    with pytest.raises(DGLError):
        g.pull(1, mfunc, rfunc)

    g.nodes["game"].data.clear()

    #############################################################
    #  update_all
    #############################################################

    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.update_all(mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[2.0, 2.0], [2.0, 2.0]]))

    # only one type
    g["plays"].update_all(mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[2.0, 2.0], [2.0, 2.0]]))

    # test fail case
    # fail due to multiple types
    with pytest.raises(DGLError):
        g.update_all(mfunc, rfunc)

    # test multi
    g.multi_update_all(
        {"plays": (mfunc, rfunc), ("user", "wishes", "game"): (mfunc, rfunc2)},
        "sum",
    )
    assert F.array_equal(
        g.nodes["game"].data["y"], F.tensor([[3.0, 3.0], [3.0, 3.0]])
    )

    # test multi
    g.multi_update_all(
        {
            "plays": (mfunc, rfunc, afunc),
            ("user", "wishes", "game"): (mfunc, rfunc2),
        },
        "sum",
        afunc,
    )
    assert F.array_equal(
        g.nodes["game"].data["y"], F.tensor([[5.0, 5.0], [5.0, 5.0]])
    )

    # test cross reducer
    g.nodes["user"].data["h"] = F.randn((3, 2))
    for cred in ["sum", "max", "min", "mean", "stack"]:
        g.multi_update_all(
            {"plays": (mfunc, rfunc, afunc), "wishes": (mfunc, rfunc2)},
            cred,
            afunc,
        )
        y = g.nodes["game"].data["y"]
        g["plays"].update_all(mfunc, rfunc, afunc)
        y1 = g.nodes["game"].data["y"]
        g["wishes"].update_all(mfunc, rfunc2)
        y2 = g.nodes["game"].data["y"]
        if cred == "stack":
            # stack has an internal order by edge type id
            yy = F.stack([y1, y2], 1)
            yy = yy + 1  # final afunc
            assert F.array_equal(y, yy)
        else:
            yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
            yy = yy + 1  # final afunc
            assert F.array_equal(y, yy)

    # test fail case
    # fail because cannot infer ntype
    with pytest.raises(DGLError):
        g.update_all(
            {"plays": (mfunc, rfunc), "follows": (mfunc, rfunc2)}, "sum"
        )

    g.nodes["game"].data.clear()


@parametrize_idtype
@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
def test_more_nnz(idtype):
    g = dgl.graph(
        ([0, 0, 0, 0, 0], [1, 1, 1, 1, 1]), idtype=idtype, device=F.ctx()
    )
    g.ndata["x"] = F.copy_to(F.ones((2, 5)), ctx=F.ctx())
    g.update_all(fn.copy_u("x", "m"), fn.sum("m", "y"))
    y = g.ndata["y"]
    ans = np.zeros((2, 5))
    ans[1] = 5
    ans = F.copy_to(F.tensor(ans, dtype=F.dtype(y)), ctx=F.ctx())
    assert F.array_equal(y, ans)


@parametrize_idtype
def test_updates(idtype):
    def msg_func(edges):
        return {"m": edges.src["h"]}

    def reduce_func(nodes):
        return {"y": F.sum(nodes.mailbox["m"], 1)}

    def apply_func(nodes):
        return {"y": nodes.data["y"] * 2}

    g = create_test_heterograph(idtype)
    x = F.randn((3, 5))
    g.nodes["user"].data["h"] = x

    for msg, red, apply in itertools.product(
        [fn.copy_u("h", "m"), msg_func],
        [fn.sum("m", "y"), reduce_func],
        [None, apply_func],
    ):
        multiplier = 1 if apply is None else 2

        g["user", "plays", "game"].update_all(msg, red, apply)
        y = g.nodes["game"].data["y"]
        assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
        assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
        del g.nodes["game"].data["y"]

        g["user", "plays", "game"].send_and_recv(
            ([0, 1, 2], [0, 1, 1]), msg, red, apply
        )
        y = g.nodes["game"].data["y"]
        assert F.array_equal(y[0], x[0] * multiplier)
        assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
        del g.nodes["game"].data["y"]

        # pulls from destination (game) node 0
        g["user", "plays", "game"].pull(0, msg, red, apply)
        y = g.nodes["game"].data["y"]
        assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
        del g.nodes["game"].data["y"]

        # pushes from source (user) node 0
        g["user", "plays", "game"].push(0, msg, red, apply)
        y = g.nodes["game"].data["y"]
        assert F.array_equal(y[0], x[0] * multiplier)
        del g.nodes["game"].data["y"]


@parametrize_idtype
def test_backward(idtype):
    g = create_test_heterograph(idtype)
    x = F.randn((3, 5))
    F.attach_grad(x)
    g.nodes["user"].data["h"] = x
    with F.record_grad():
        g.multi_update_all(
            {
                "plays": (fn.copy_u("h", "m"), fn.sum("m", "y")),
                "wishes": (fn.copy_u("h", "m"), fn.sum("m", "y")),
            },
            "sum",
        )
        y = g.nodes["game"].data["y"]
        F.backward(y, F.ones(y.shape))
    print(F.grad(x))
    assert F.array_equal(
        F.grad(x),
        F.tensor(
            [
                [2.0, 2.0, 2.0, 2.0, 2.0],
                [2.0, 2.0, 2.0, 2.0, 2.0],
                [2.0, 2.0, 2.0, 2.0, 2.0],
            ]
        ),
    )


@parametrize_idtype
def test_empty_heterograph(idtype):
    def assert_empty(g):
        assert g.number_of_nodes("user") == 0
        assert g.number_of_edges("plays") == 0
        assert g.number_of_nodes("game") == 0

    # empty src-dst pair
    assert_empty(dgl.heterograph({("user", "plays", "game"): ([], [])}))

    g = dgl.heterograph(
        {("user", "follows", "user"): ([], [])}, idtype=idtype, device=F.ctx()
    )
    assert g.idtype == idtype
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 0
    assert g.number_of_edges("follows") == 0

    # empty relation graph with others
    g = dgl.heterograph(
        {
            ("user", "plays", "game"): ([], []),
            ("developer", "develops", "game"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.idtype == idtype
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 0
    assert g.number_of_edges("plays") == 0
    assert g.number_of_nodes("game") == 2
    assert g.number_of_edges("develops") == 2
    assert g.number_of_nodes("developer") == 2


@parametrize_idtype
def test_types_in_function(idtype):
    def mfunc1(edges):
        assert edges.canonical_etype == ("user", "follow", "user")
        return {}

    def rfunc1(nodes):
        assert nodes.ntype == "user"
        return {}

    def filter_nodes1(nodes):
        assert nodes.ntype == "user"
        return F.zeros((3,))

    def filter_edges1(edges):
        assert edges.canonical_etype == ("user", "follow", "user")
        return F.zeros((2,))

    def mfunc2(edges):
        assert edges.canonical_etype == ("user", "plays", "game")
        return {}

    def rfunc2(nodes):
        assert nodes.ntype == "game"
        return {}

    def filter_nodes2(nodes):
        assert nodes.ntype == "game"
        return F.zeros((3,))

    def filter_edges2(edges):
        assert edges.canonical_etype == ("user", "plays", "game")
        return F.zeros((2,))

    g = dgl.heterograph(
        {("user", "follow", "user"): ((0, 1), (1, 2))},
        idtype=idtype,
        device=F.ctx(),
    )
    g.apply_nodes(rfunc1)
    g.apply_edges(mfunc1)
    g.update_all(mfunc1, rfunc1)
    g.send_and_recv([0, 1], mfunc1, rfunc1)
    g.push([0], mfunc1, rfunc1)
    g.pull([1], mfunc1, rfunc1)
    g.filter_nodes(filter_nodes1)
    g.filter_edges(filter_edges1)

    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.apply_nodes(rfunc2, ntype="game")
    g.apply_edges(mfunc2)
    g.update_all(mfunc2, rfunc2)
    g.send_and_recv([0, 1], mfunc2, rfunc2)
    g.push([0], mfunc2, rfunc2)
    g.pull([1], mfunc2, rfunc2)
    g.filter_nodes(filter_nodes2, ntype="game")
    g.filter_edges(filter_edges2)


@parametrize_idtype
def test_stack_reduce(idtype):
    # edges = {
    #    'follows': ([0, 1], [1, 2]),
    #    'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
    #    'wishes': ([0, 2], [1, 0]),
    #    'develops': ([0, 1], [0, 1]),
    # }
    g = create_test_heterograph(idtype)
    g.nodes["user"].data["h"] = F.randn((3, 200))

    def rfunc(nodes):
        return {"y": F.sum(nodes.mailbox["m"], 1)}

    def rfunc2(nodes):
        return {"y": F.max(nodes.mailbox["m"], 1)}

    def mfunc(edges):
        return {"m": edges.src["h"]}

    g.multi_update_all(
        {"plays": (mfunc, rfunc), "wishes": (mfunc, rfunc2)}, "stack"
    )
    assert g.nodes["game"].data["y"].shape == (
        g.number_of_nodes("game"),
        2,
        200,
    )
    # only one type-wise update_all, stack still adds one dimension
    g.multi_update_all({"plays": (mfunc, rfunc)}, "stack")
    assert g.nodes["game"].data["y"].shape == (
        g.number_of_nodes("game"),
        1,
        200,
    )


@parametrize_idtype
def test_isolated_ntype(idtype):
    g = dgl.heterograph(
        {("A", "AB", "B"): ([0, 1, 2], [1, 2, 3])},
        num_nodes_dict={"A": 3, "B": 4, "C": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.number_of_nodes("A") == 3
    assert g.number_of_nodes("B") == 4
    assert g.number_of_nodes("C") == 4

    g = dgl.heterograph(
        {("A", "AC", "C"): ([0, 1, 2], [1, 2, 3])},
        num_nodes_dict={"A": 3, "B": 4, "C": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.number_of_nodes("A") == 3
    assert g.number_of_nodes("B") == 4
    assert g.number_of_nodes("C") == 4

    G = dgl.graph(
        ([0, 1, 2], [4, 5, 6]), num_nodes=11, idtype=idtype, device=F.ctx()
    )
    G.ndata[dgl.NTYPE] = F.tensor(
        [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=F.int64
    )
    G.edata[dgl.ETYPE] = F.tensor([0, 0, 0], dtype=F.int64)
    g = dgl.to_heterogeneous(G, ["A", "B", "C"], ["AB"])
    assert g.number_of_nodes("A") == 3
    assert g.number_of_nodes("B") == 4
    assert g.number_of_nodes("C") == 4


@parametrize_idtype
def test_ismultigraph(idtype):
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5])},
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g1.is_multigraph == False
    g2 = dgl.heterograph(
        {("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5])},
        {"A": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g2.is_multigraph == True
    g3 = dgl.graph(((0, 1), (1, 2)), num_nodes=6, idtype=idtype, device=F.ctx())
    assert g3.is_multigraph == False
    g4 = dgl.graph(
        ([0, 0, 1], [1, 1, 2]), num_nodes=6, idtype=idtype, device=F.ctx()
    )
    assert g4.is_multigraph == True
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AA", "A"): ([0, 1], [1, 2]),
        },
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.is_multigraph == False
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5]),
        },
        {"A": 6, "B": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.is_multigraph == True
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AA", "A"): ([0, 0, 1], [1, 1, 2]),
        },
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.is_multigraph == True
    g = dgl.heterograph(
        {
            ("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5]),
            ("A", "AA", "A"): ([0, 1], [1, 2]),
        },
        {"A": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g.is_multigraph == True


@parametrize_idtype
def test_graph_index_is_unibipartite(idtype):
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1], [1, 2, 5])},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g1._graph.is_metagraph_unibipartite()

    # more complicated bipartite
    g2 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert g2._graph.is_metagraph_unibipartite()

    g3 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
            ("A", "AA", "A"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert not g3._graph.is_metagraph_unibipartite()

    g4 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("C", "CA", "A"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )

    assert not g4._graph.is_metagraph_unibipartite()


@parametrize_idtype
def test_bipartite(idtype):
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1], [1, 2, 5])},
        idtype=idtype,
        device=F.ctx(),
    )
    assert g1.is_unibipartite
    assert len(g1.ntypes) == 2
    assert g1.etypes == ["AB"]
    assert g1.srctypes == ["A"]
    assert g1.dsttypes == ["B"]
    assert g1.number_of_nodes("A") == 2
    assert g1.number_of_nodes("B") == 6
    assert g1.number_of_src_nodes("A") == 2
    assert g1.number_of_src_nodes() == 2
    assert g1.number_of_dst_nodes("B") == 6
    assert g1.number_of_dst_nodes() == 6
    assert g1.number_of_edges() == 3
    g1.srcdata["h"] = F.randn((2, 5))
    assert F.array_equal(g1.srcnodes["A"].data["h"], g1.srcdata["h"])
    assert F.array_equal(g1.nodes["A"].data["h"], g1.srcdata["h"])
    assert F.array_equal(g1.nodes["SRC/A"].data["h"], g1.srcdata["h"])
    g1.dstdata["h"] = F.randn((6, 3))
    assert F.array_equal(g1.dstnodes["B"].data["h"], g1.dstdata["h"])
    assert F.array_equal(g1.nodes["B"].data["h"], g1.dstdata["h"])
    assert F.array_equal(g1.nodes["DST/B"].data["h"], g1.dstdata["h"])

    # more complicated bipartite
    g2 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )

    assert g2.is_unibipartite
    assert g2.srctypes == ["A"]
    assert set(g2.dsttypes) == {"B", "C"}
    assert g2.number_of_nodes("A") == 2
    assert g2.number_of_nodes("B") == 6
    assert g2.number_of_nodes("C") == 1
    assert g2.number_of_src_nodes("A") == 2
    assert g2.number_of_src_nodes() == 2
    assert g2.number_of_dst_nodes("B") == 6
    assert g2.number_of_dst_nodes("C") == 1
    g2.srcdata["h"] = F.randn((2, 5))
    assert F.array_equal(g2.srcnodes["A"].data["h"], g2.srcdata["h"])
    assert F.array_equal(g2.nodes["A"].data["h"], g2.srcdata["h"])
    assert F.array_equal(g2.nodes["SRC/A"].data["h"], g2.srcdata["h"])

    g3 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
            ("A", "AA", "A"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert not g3.is_unibipartite

    g4 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("C", "CA", "A"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )

    assert not g4.is_unibipartite


@parametrize_idtype
def test_dtype_cast(idtype):
    g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
    assert g.idtype == idtype
    g.ndata["feat"] = F.tensor([3, 4, 5])
    g.edata["h"] = F.tensor([3, 4, 5, 6])
    if idtype == "int32":
        g_cast = g.long()
        assert g_cast.idtype == F.int64
    else:
        g_cast = g.int()
        assert g_cast.idtype == F.int32
    test_utils.check_graph_equal(g, g_cast, check_idtype=False)


def test_float_cast():
    for t in [F.float16, F.float32, F.float64]:
        idtype = F.int32
        g = dgl.heterograph(
            {
                ("user", "follows", "user"): (
                    F.tensor([0, 1, 1, 2, 2, 3], dtype=idtype),
                    F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype),
                ),
                ("user", "plays", "game"): (
                    F.tensor([0, 1, 1], dtype=idtype),
                    F.tensor([0, 0, 1], dtype=idtype),
                ),
            },
            idtype=idtype,
            device=F.ctx(),
        )
        uvalues = [1, 2, 3, 4]
        gvalues = [5, 6]
        fvalues = [7, 8, 9, 10, 11, 12]
        pvalues = [13, 14, 15]
        dataNamesTypes = [
            ("a", F.float16),
            ("b", F.float32),
            ("c", F.float64),
            ("d", F.int32),
            ("e", F.int64),
        ]
        for name, type in dataNamesTypes:
            g.nodes["user"].data[name] = F.copy_to(
                F.tensor(uvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.nodes["game"].data[name] = F.copy_to(
                F.tensor(gvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.edges["follows"].data[name] = F.copy_to(
                F.tensor(fvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.edges["plays"].data[name] = F.copy_to(
                F.tensor(pvalues, dtype=type), ctx=F.ctx()
            )

        if t == F.float16:
            g = dgl.transforms.functional.to_half(g)
        if t == F.float32:
            g = dgl.transforms.functional.to_float(g)
        if t == F.float64:
            g = dgl.transforms.functional.to_double(g)

        for name, origType in dataNamesTypes:
            # integer tensors shouldn't be converted
            reqType = (
                t
                if (origType in [F.float16, F.float32, F.float64])
                else origType
            )

            values = g.nodes["user"].data[name]
            assert values.dtype == reqType
            assert len(values) == len(uvalues)
            assert F.allclose(values, F.tensor(uvalues), 0, 0)

            values = g.nodes["game"].data[name]
            assert values.dtype == reqType
            assert len(values) == len(gvalues)
            assert F.allclose(values, F.tensor(gvalues), 0, 0)

            values = g.edges["follows"].data[name]
            assert values.dtype == reqType
            assert len(values) == len(fvalues)
            assert F.allclose(values, F.tensor(fvalues), 0, 0)

            values = g.edges["plays"].data[name]
            assert values.dtype == reqType
            assert len(values) == len(pvalues)
            assert F.allclose(values, F.tensor(pvalues), 0, 0)


@parametrize_idtype
def test_format(idtype):
    # single relation
    g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
    assert g.formats()["created"] == ["coo"]
    g1 = g.formats(["coo", "csr", "csc"])
    assert len(g1.formats()["created"]) + len(g1.formats()["not created"]) == 3
    g1.create_formats_()
    assert len(g1.formats()["created"]) == 3
    assert g.formats()["created"] == ["coo"]

    # multiple relation
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1]),
            ("developer", "develops", "game"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    user_feat = F.randn((g["follows"].number_of_src_nodes(), 5))
    g["follows"].srcdata["h"] = user_feat
    g1 = g.formats("csc")
    # test frame
    assert F.array_equal(g1["follows"].srcdata["h"], user_feat)
    # test each relation graph
    assert g1.formats()["created"] == ["csc"]
    assert len(g1.formats()["not created"]) == 0

    # in_degrees
    g = dgl.rand_graph(100, 2340).to(F.ctx())
    ind_arr = []
    for vid in range(0, 100):
        ind_arr.append(g.in_degrees(vid))
    in_degrees = g.in_degrees()
    g = g.formats("coo")
    for vid in range(0, 100):
        assert g.in_degrees(vid) == ind_arr[vid]
    assert F.array_equal(in_degrees, g.in_degrees())


@parametrize_idtype
def test_edges_order(idtype):
    # (0, 2), (1, 2), (0, 1), (0, 1), (2, 1)
    g = dgl.graph(
        (np.array([0, 1, 0, 0, 2]), np.array([2, 2, 1, 1, 1])),
        idtype=idtype,
        device=F.ctx(),
    )

    print(g.formats())
    src, dst = g.all_edges(order="srcdst")
    assert F.array_equal(src, F.tensor([0, 0, 0, 1, 2], dtype=idtype))
    assert F.array_equal(dst, F.tensor([1, 1, 2, 2, 1], dtype=idtype))


@parametrize_idtype
def test_reverse(idtype):
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                [0, 1, 2, 4, 3, 1, 3],
                [1, 2, 3, 2, 0, 0, 1],
            )
        },
        idtype=idtype,
        device=F.ctx(),
    )
    gidx = g._graph
    r_gidx = gidx.reverse()

    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

    # force to start with 'csr'
    gidx = gidx.formats("csr")
    gidx = gidx.formats(["coo", "csr", "csc"])
    r_gidx = gidx.reverse()
    assert "csr" in gidx.formats()["created"]
    assert "csc" in r_gidx.formats()["created"]
    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

    # force to start with 'csc'
    gidx = gidx.formats("csc")
    gidx = gidx.formats(["coo", "csr", "csc"])
    r_gidx = gidx.reverse()
    assert "csc" in gidx.formats()["created"]
    assert "csr" in r_gidx.formats()["created"]
    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                [0, 1, 2, 4, 3, 1, 3],
                [1, 2, 3, 2, 0, 0, 1],
            ),
            ("user", "plays", "game"): (
                [0, 0, 2, 3, 3, 4, 1],
                [1, 0, 1, 0, 1, 0, 0],
            ),
            ("developer", "develops", "game"): ([0, 1, 1, 2], [0, 0, 1, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    gidx = g._graph
    r_gidx = gidx.reverse()

    # metagraph
    mg = gidx.metagraph
    r_mg = r_gidx.metagraph
    for etype in range(3):
        assert mg.find_edge(etype) == r_mg.find_edge(etype)[::-1]

    # three node types and three edge types
    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
    assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
    assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

    # force to start with 'csr'
    gidx = gidx.formats("csr")
    gidx = gidx.formats(["coo", "csr", "csc"])
    r_gidx = gidx.reverse()
    # three node types and three edge types
    assert "csr" in gidx.formats()["created"]
    assert "csc" in r_gidx.formats()["created"]
    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
    assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
    assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

    # force to start with 'csc'
    gidx = gidx.formats("csc")
    gidx = gidx.formats(["coo", "csr", "csc"])
    r_gidx = gidx.reverse()
    # three node types and three edge types
    assert "csc" in gidx.formats()["created"]
    assert "csr" in r_gidx.formats()["created"]
    assert gidx.number_of_nodes(0) == r_gidx.number_of_nodes(0)
    assert gidx.number_of_nodes(1) == r_gidx.number_of_nodes(1)
    assert gidx.number_of_nodes(2) == r_gidx.number_of_nodes(2)
    assert gidx.number_of_edges(0) == r_gidx.number_of_edges(0)
    assert gidx.number_of_edges(1) == r_gidx.number_of_edges(1)
    assert gidx.number_of_edges(2) == r_gidx.number_of_edges(2)
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)


@parametrize_idtype
def test_clone(idtype):
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.edata["h"] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())

    new_g = g.clone()
    assert g.number_of_nodes() == new_g.number_of_nodes()
    assert g.number_of_edges() == new_g.number_of_edges()
    assert g.device == new_g.device
    assert g.idtype == new_g.idtype
    assert F.array_equal(g.ndata["h"], new_g.ndata["h"])
    assert F.array_equal(g.edata["h"], new_g.edata["h"])
    # data change
    new_g.ndata["h"] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
    assert F.array_equal(g.ndata["h"], new_g.ndata["h"]) == False
    g.edata["h"] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
    assert F.array_equal(g.edata["h"], new_g.edata["h"]) == False
    # graph structure change
    g.add_nodes(1)
    assert g.number_of_nodes() != new_g.number_of_nodes()
    new_g.add_edges(1, 1)
    assert g.number_of_edges() != new_g.number_of_edges()

    # zero data graph
    g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
    new_g = g.clone()
    assert g.number_of_nodes() == new_g.number_of_nodes()
    assert g.number_of_edges() == new_g.number_of_edges()

    # heterograph
    g = create_test_heterograph3(idtype)
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
    new_g = g.clone()
    assert g.number_of_nodes("user") == new_g.number_of_nodes("user")
    assert g.number_of_nodes("game") == new_g.number_of_nodes("game")
    assert g.number_of_nodes("developer") == new_g.number_of_nodes("developer")
    assert g.number_of_edges("plays") == new_g.number_of_edges("plays")
    assert g.number_of_edges("develops") == new_g.number_of_edges("develops")
    assert F.array_equal(
        g.nodes["user"].data["h"], new_g.nodes["user"].data["h"]
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], new_g.nodes["game"].data["h"]
    )
    assert F.array_equal(
        g.edges["plays"].data["h"], new_g.edges["plays"].data["h"]
    )
    assert g.device == new_g.device
    assert g.idtype == new_g.idtype
    u, v = g.edges(form="uv", order="eid", etype="plays")
    nu, nv = new_g.edges(form="uv", order="eid", etype="plays")
    assert F.array_equal(u, nu)
    assert F.array_equal(v, nv)
    # graph structure change
    u = F.tensor([0, 4], dtype=idtype)
    v = F.tensor([2, 6], dtype=idtype)
    g.add_edges(u, v, etype="plays")
    u, v = g.edges(form="uv", order="eid", etype="plays")
    assert u.shape[0] != nu.shape[0]
    assert v.shape[0] != nv.shape[0]
    assert (
        g.nodes["user"].data["h"].shape[0]
        != new_g.nodes["user"].data["h"].shape[0]
    )
    assert (
        g.nodes["game"].data["h"].shape[0]
        != new_g.nodes["game"].data["h"].shape[0]
    )
    assert (
        g.edges["plays"].data["h"].shape[0]
        != new_g.edges["plays"].data["h"].shape[0]
    )


@parametrize_idtype
def test_add_edges(idtype):
    # homogeneous graph
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    u = 0
    v = 1
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 3
    u = [0]
    v = [1]
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 4
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 5
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))

    # node id larger than current max node id
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g.add_edges(u, v)
    assert g.number_of_nodes() == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))

    # has data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.edata["h"] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    e_feat = {
        "h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
        "hh": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
    }
    g.add_edges(u, v, e_feat)
    assert g.number_of_nodes() == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
    assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 0], dtype=idtype))
    assert F.array_equal(g.edata["h"], F.tensor([1, 1, 2, 2], dtype=idtype))
    assert F.array_equal(g.edata["hh"], F.tensor([0, 0, 2, 2], dtype=idtype))

    # zero data graph
    g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 2], dtype=idtype)
    e_feat = {
        "h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
        "hh": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
    }
    g.add_edges(u, v, e_feat)
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 2
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
    assert F.array_equal(g.edata["h"], F.tensor([2, 2], dtype=idtype))
    assert F.array_equal(g.edata["hh"], F.tensor([2, 2], dtype=idtype))

    # bipartite graph
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    u = 0
    v = 1
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 2
    assert g.number_of_nodes("game") == 3
    assert g.number_of_edges() == 3
    u = [0]
    v = [1]
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 2
    assert g.number_of_nodes("game") == 3
    assert g.number_of_edges() == 4
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 2
    assert g.number_of_nodes("game") == 3
    assert g.number_of_edges() == 5
    u, v = g.edges(form="uv")
    assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))

    # node id larger than current max node id
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes("user") == 3
    assert g.number_of_nodes("game") == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))

    # has data
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.copy_to(
        F.tensor([1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edata["h"] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    e_feat = {
        "h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
        "hh": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
    }
    g.add_edges(u, v, e_feat)
    assert g.number_of_nodes("user") == 3
    assert g.number_of_nodes("game") == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1, 0], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2, 2, 0], dtype=idtype)
    )
    assert F.array_equal(g.edata["h"], F.tensor([1, 1, 2, 2], dtype=idtype))
    assert F.array_equal(g.edata["hh"], F.tensor([0, 0, 2, 2], dtype=idtype))

    # heterogeneous graph
    g = create_test_heterograph3(idtype)
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g.add_edges(u, v, etype="plays")
    assert g.number_of_nodes("user") == 3
    assert g.number_of_nodes("game") == 4
    assert g.number_of_nodes("developer") == 2
    assert g.number_of_edges("plays") == 6
    assert g.number_of_edges("develops") == 2
    u, v = g.edges(form="uv", order="eid", etype="plays")
    assert F.array_equal(u, F.tensor([0, 1, 1, 2, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 0, 1, 1, 2, 3], dtype=idtype))
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1, 1], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2, 0, 0], dtype=idtype)
    )
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([1, 1, 1, 1, 0, 0], dtype=idtype)
    )

    # add with feature
    e_feat = {"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.add_edges(u, v, data=e_feat, etype="develops")
    assert g.number_of_nodes("user") == 3
    assert g.number_of_nodes("game") == 4
    assert g.number_of_nodes("developer") == 3
    assert g.number_of_edges("plays") == 6
    assert g.number_of_edges("develops") == 4
    u, v = g.edges(form="uv", order="eid", etype="develops")
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 2, 3], dtype=idtype))
    assert F.array_equal(
        g.nodes["developer"].data["h"], F.tensor([3, 3, 0], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2, 1, 1], dtype=idtype)
    )
    assert F.array_equal(
        g.edges["develops"].data["h"], F.tensor([0, 0, 2, 2], dtype=idtype)
    )


@parametrize_idtype
def test_add_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.add_nodes(1)
    assert g.number_of_nodes() == 4
    assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 0], dtype=idtype))

    # zero node graph
    g = dgl.graph(([], []), num_nodes=3, idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.add_nodes(
        1, data={"h": F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())}
    )
    assert g.number_of_nodes() == 4
    assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 2], dtype=idtype))

    # bipartite graph
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.add_nodes(
        2,
        data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
        ntype="user",
    )
    assert g.number_of_nodes("user") == 4
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([0, 0, 2, 2], dtype=idtype)
    )
    g.add_nodes(2, ntype="game")
    assert g.number_of_nodes("game") == 5

    # heterogeneous graph
    g = create_test_heterograph3(idtype)
    g.add_nodes(1, ntype="user")
    g.add_nodes(
        2,
        data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
        ntype="game",
    )
    g.add_nodes(0, ntype="developer")
    assert g.number_of_nodes("user") == 4
    assert g.number_of_nodes("game") == 4
    assert g.number_of_nodes("developer") == 2
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1, 1, 0], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2, 2, 2], dtype=idtype)
    )


@unittest.skipIf(
    dgl.backend.backend_name == "mxnet",
    reason="MXNet has error with (0,) shape tensor.",
)
@parametrize_idtype
def test_remove_edges(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    e = 0
    g.remove_edges(e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    e = [0]
    g.remove_edges(e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    e = F.tensor([0], dtype=idtype)
    g.remove_edges(e)
    assert g.number_of_edges() == 0

    # has node data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.remove_edges(1)
    assert g.number_of_edges() == 1
    assert F.array_equal(g.ndata["h"], F.tensor([1, 2, 3], dtype=idtype))

    # has edge data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.edata["h"] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
    g.remove_edges(0)
    assert g.number_of_edges() == 1
    assert F.array_equal(g.edata["h"], F.tensor([2], dtype=idtype))

    # invalid eid
    assert_fail = False
    try:
        g.remove_edges(1)
    except:
        assert_fail = True
    assert assert_fail

    # bipartite graph
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    e = 0
    g.remove_edges(e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    e = [0]
    g.remove_edges(e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    e = F.tensor([0], dtype=idtype)
    g.remove_edges(e)
    assert g.number_of_edges() == 0

    # has data
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.copy_to(
        F.tensor([1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edata["h"] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
    g.remove_edges(1)
    assert g.number_of_edges() == 1
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2, 2], dtype=idtype)
    )
    assert F.array_equal(g.edata["h"], F.tensor([1], dtype=idtype))

    # heterogeneous graph
    g = create_test_heterograph3(idtype)
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
    g.remove_edges(1, etype="plays")
    assert g.number_of_edges("plays") == 3
    u, v = g.edges(form="uv", order="eid", etype="plays")
    assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([1, 3, 4], dtype=idtype)
    )
    # remove all edges of 'develops'
    g.remove_edges([0, 1], etype="develops")
    assert g.number_of_edges("develops") == 0
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1, 1], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], F.tensor([2, 2], dtype=idtype)
    )
    assert F.array_equal(
        g.nodes["developer"].data["h"], F.tensor([3, 3], dtype=idtype)
    )


@parametrize_idtype
def test_remove_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = 0
    g.remove_nodes(n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = [1]
    g.remove_nodes(n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 0
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = F.tensor([2], dtype=idtype)
    g.remove_nodes(n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))

    # invalid nid
    assert_fail = False
    try:
        g.remove_nodes(3)
    except:
        assert_fail = True
    assert assert_fail

    # has node and edge data
    g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
    g.ndata["hv"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.edata["he"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.remove_nodes(F.tensor([0], dtype=idtype))
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
    assert F.array_equal(g.ndata["hv"], F.tensor([2, 3], dtype=idtype))
    assert F.array_equal(g.edata["he"], F.tensor([3], dtype=idtype))

    # node id larger than current max node id
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    n = 0
    g.remove_nodes(n, ntype="user")
    assert g.number_of_nodes("user") == 1
    assert g.number_of_nodes("game") == 3
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    n = [1]
    g.remove_nodes(n, ntype="user")
    assert g.number_of_nodes("user") == 1
    assert g.number_of_nodes("game") == 3
    assert g.number_of_edges() == 1
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    n = F.tensor([0], dtype=idtype)
    g.remove_nodes(n, ntype="game")
    assert g.number_of_nodes("user") == 2
    assert g.number_of_nodes("game") == 2
    assert g.number_of_edges() == 2
    u, v = g.edges(form="uv", order="eid")
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))

    # heterogeneous graph
    g = create_test_heterograph3(idtype)
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
    g.remove_nodes(0, ntype="game")
    assert g.number_of_nodes("user") == 3
    assert g.number_of_nodes("game") == 1
    assert g.number_of_nodes("developer") == 2
    assert g.number_of_edges("plays") == 2
    assert g.number_of_edges("develops") == 1
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([1, 1, 1], dtype=idtype)
    )
    assert F.array_equal(g.nodes["game"].data["h"], F.tensor([2], dtype=idtype))
    assert F.array_equal(
        g.nodes["developer"].data["h"], F.tensor([3, 3], dtype=idtype)
    )
    u, v = g.edges(form="uv", order="eid", etype="plays")
    assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([3, 4], dtype=idtype)
    )
    u, v = g.edges(form="uv", order="eid", etype="develops")
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([0], dtype=idtype))


@parametrize_idtype
def test_frame(idtype):
    g = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
    g.ndata["h"] = F.copy_to(F.tensor([0, 1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.edata["h"] = F.copy_to(F.tensor([0, 1, 2], dtype=idtype), ctx=F.ctx())

    # remove nodes
    sg = dgl.remove_nodes(g, [3])
    # check for lazy update
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])
    assert sg.ndata["h"].shape[0] == 3
    assert sg.edata["h"].shape[0] == 2
    # update after read
    assert F.array_equal(
        sg._node_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 2], dtype=idtype),
    )
    assert F.array_equal(
        sg._edge_frames[0]._columns["h"].storage, F.tensor([0, 1], dtype=idtype)
    )

    ng = dgl.add_nodes(sg, 1)
    assert ng.ndata["h"].shape[0] == 4
    assert F.array_equal(
        ng._node_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 2, 0], dtype=idtype),
    )
    ng = dgl.add_edges(ng, [3], [1])
    assert ng.edata["h"].shape[0] == 3
    assert F.array_equal(
        ng._edge_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 0], dtype=idtype),
    )

    # multi level lazy update
    sg = dgl.remove_nodes(g, [3])
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])
    ssg = dgl.remove_nodes(sg, [1])
    assert F.array_equal(
        ssg._node_frames[0]._columns["h"].storage, g.ndata["h"]
    )
    assert F.array_equal(
        ssg._edge_frames[0]._columns["h"].storage, g.edata["h"]
    )
    # ssg is changed
    assert ssg.ndata["h"].shape[0] == 2
    assert ssg.edata["h"].shape[0] == 0
    assert F.array_equal(
        ssg._node_frames[0]._columns["h"].storage,
        F.tensor([0, 2], dtype=idtype),
    )
    # sg still in lazy model
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])


@unittest.skipIf(
    dgl.backend.backend_name == "tensorflow",
    reason="TensorFlow always create a new tensor",
)
@unittest.skipIf(
    F._default_context_str == "cpu",
    reason="cpu do not have context change problem",
)
@parametrize_idtype
def test_frame_device(idtype):
    g = dgl.graph(([0, 1, 2], [2, 3, 1]))
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1, 2], dtype=idtype), ctx=F.cpu())
    g.ndata["hh"] = F.copy_to(F.ones((4, 3), dtype=idtype), ctx=F.cpu())
    g.edata["h"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.cpu())

    g = g.to(F.ctx())
    # lazy device copy
    assert F.context(g._node_frames[0]._columns["h"].storage) == F.cpu()
    assert F.context(g._node_frames[0]._columns["hh"].storage) == F.cpu()
    print(g.ndata["h"])
    assert F.context(g._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(g._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(g._edge_frames[0]._columns["h"].storage) == F.cpu()

    # lazy device copy in subgraph
    sg = dgl.node_subgraph(g, [0, 1, 2])
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
    print(sg.ndata["hh"])
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()

    # back to cpu
    sg = sg.to(F.cpu())
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
    print(sg.ndata["h"])
    print(sg.ndata["hh"])
    print(sg.edata["h"])
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.cpu()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()

    # set some field
    sg = sg.to(F.ctx())
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.cpu()
    sg.ndata["h"][0] = 5
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()

    # add nodes
    ng = dgl.add_nodes(sg, 3)
    assert F.context(ng._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(ng._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(ng._edge_frames[0]._columns["h"].storage) == F.cpu()


@parametrize_idtype
def test_create_block(idtype):
    block = dgl.create_block(
        ([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()
    )
    assert block.num_src_nodes() == 3
    assert block.num_dst_nodes() == 4
    assert block.num_edges() == 3

    block = dgl.create_block(([], []), idtype=idtype, device=F.ctx())
    assert block.num_src_nodes() == 0
    assert block.num_dst_nodes() == 0
    assert block.num_edges() == 0

    block = dgl.create_block(([], []), 3, 4, idtype=idtype, device=F.ctx())
    assert block.num_src_nodes() == 3
    assert block.num_dst_nodes() == 4
    assert block.num_edges() == 0

    block = dgl.create_block(
        ([0, 1, 2], [1, 2, 3]), 4, 5, idtype=idtype, device=F.ctx()
    )
    assert block.num_src_nodes() == 4
    assert block.num_dst_nodes() == 5
    assert block.num_edges() == 3

    sx = F.randn((4, 5))
    dx = F.randn((5, 6))
    ex = F.randn((3, 4))
    block.srcdata["x"] = sx
    block.dstdata["x"] = dx
    block.edata["x"] = ex

    g = dgl.block_to_graph(block)
    assert g.num_src_nodes() == 4
    assert g.num_dst_nodes() == 5
    assert g.num_edges() == 3
    assert g.srcdata["x"] is sx
    assert g.dstdata["x"] is dx
    assert g.edata["x"] is ex

    block = dgl.create_block(
        {
            ("A", "AB", "B"): ([1, 2, 3], [2, 1, 0]),
            ("B", "BA", "A"): ([2, 3], [3, 4]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 4
    assert block.num_src_nodes("B") == 4
    assert block.num_dst_nodes("B") == 3
    assert block.num_dst_nodes("A") == 5
    assert block.num_edges("AB") == 3
    assert block.num_edges("BA") == 2

    block = dgl.create_block(
        {("A", "AB", "B"): ([], []), ("B", "BA", "A"): ([], [])},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 0
    assert block.num_src_nodes("B") == 0
    assert block.num_dst_nodes("B") == 0
    assert block.num_dst_nodes("A") == 0
    assert block.num_edges("AB") == 0
    assert block.num_edges("BA") == 0

    block = dgl.create_block(
        {("A", "AB", "B"): ([], []), ("B", "BA", "A"): ([], [])},
        num_src_nodes={"A": 5, "B": 5},
        num_dst_nodes={"A": 6, "B": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 5
    assert block.num_src_nodes("B") == 5
    assert block.num_dst_nodes("B") == 4
    assert block.num_dst_nodes("A") == 6
    assert block.num_edges("AB") == 0
    assert block.num_edges("BA") == 0

    block = dgl.create_block(
        {
            ("A", "AB", "B"): ([1, 2, 3], [2, 1, 0]),
            ("B", "BA", "A"): ([2, 3], [3, 4]),
        },
        num_src_nodes={"A": 5, "B": 5},
        num_dst_nodes={"A": 6, "B": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 5
    assert block.num_src_nodes("B") == 5
    assert block.num_dst_nodes("B") == 4
    assert block.num_dst_nodes("A") == 6
    assert block.num_edges(("A", "AB", "B")) == 3
    assert block.num_edges(("B", "BA", "A")) == 2

    sax = F.randn((5, 3))
    sbx = F.randn((5, 4))
    dax = F.randn((6, 5))
    dbx = F.randn((4, 6))
    eabx = F.randn((3, 7))
    ebax = F.randn((2, 8))
    block.srcnodes["A"].data["x"] = sax
    block.srcnodes["B"].data["x"] = sbx
    block.dstnodes["A"].data["x"] = dax
    block.dstnodes["B"].data["x"] = dbx
    block.edges["AB"].data["x"] = eabx
    block.edges["BA"].data["x"] = ebax

    hg = dgl.block_to_graph(block)
    assert hg.num_nodes("A_src") == 5
    assert hg.num_nodes("B_src") == 5
    assert hg.num_nodes("A_dst") == 6
    assert hg.num_nodes("B_dst") == 4
    assert hg.num_edges(("A_src", "AB", "B_dst")) == 3
    assert hg.num_edges(("B_src", "BA", "A_dst")) == 2
    assert hg.nodes["A_src"].data["x"] is sax
    assert hg.nodes["B_src"].data["x"] is sbx
    assert hg.nodes["A_dst"].data["x"] is dax
    assert hg.nodes["B_dst"].data["x"] is dbx
    assert hg.edges["AB"].data["x"] is eabx
    assert hg.edges["BA"].data["x"] is ebax


@parametrize_idtype
@pytest.mark.parametrize("fmt", ["coo", "csr", "csc"])
def test_adj_sparse(idtype, fmt):
    if fmt == "coo":
        A = ssp.random(10, 10, 0.2).tocoo()
        A.data = np.arange(20)
        row = F.tensor(A.row, idtype)
        col = F.tensor(A.col, idtype)
        g = dgl.graph((row, col))
    elif fmt == "csr":
        A = ssp.random(10, 10, 0.2).tocsr()
        A.data = np.arange(20)
        indptr = F.tensor(A.indptr, idtype)
        indices = F.tensor(A.indices, idtype)
        g = dgl.graph(("csr", (indptr, indices, [])))
        with pytest.raises(DGLError):
            g2 = dgl.graph(("csr", (indptr[:-1], indices, [])), num_nodes=10)
    elif fmt == "csc":
        A = ssp.random(10, 10, 0.2).tocsc()
        A.data = np.arange(20)
        indptr = F.tensor(A.indptr, idtype)
        indices = F.tensor(A.indices, idtype)
        g = dgl.graph(("csc", (indptr, indices, [])))
        with pytest.raises(DGLError):
            g2 = dgl.graph(("csr", (indptr[:-1], indices, [])), num_nodes=10)

    A_coo = A.tocoo()
    A_csr = A.tocsr()
    A_csc = A.tocsc()
    row, col = g.adj_sparse("coo")
    assert np.array_equal(F.asnumpy(row), A_coo.row)
    assert np.array_equal(F.asnumpy(col), A_coo.col)

    indptr, indices, eids = g.adj_sparse("csr")
    assert np.array_equal(F.asnumpy(indptr), A_csr.indptr)
    if fmt == "csr":
        assert len(eids) == 0
        assert np.array_equal(F.asnumpy(indices), A_csr.indices)
    else:
        indices_sorted = F.zeros(len(indices), idtype)
        indices_sorted = F.scatter_row(indices_sorted, eids, indices)
        indices_sorted_np = np.zeros(len(indices), dtype=A_csr.indices.dtype)
        indices_sorted_np[A_csr.data] = A_csr.indices
        assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)

    indptr, indices, eids = g.adj_sparse("csc")
    assert np.array_equal(F.asnumpy(indptr), A_csc.indptr)
    if fmt == "csc":
        assert len(eids) == 0
        assert np.array_equal(F.asnumpy(indices), A_csc.indices)
    else:
        indices_sorted = F.zeros(len(indices), idtype)
        indices_sorted = F.scatter_row(indices_sorted, eids, indices)
        indices_sorted_np = np.zeros(len(indices), dtype=A_csc.indices.dtype)
        indices_sorted_np[A_csc.data] = A_csc.indices
        assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)


def _test_forking_pickler_entry(g, q):
    q.put(g.formats())


@unittest.skipIf(
    dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support spawning"
)
def test_forking_pickler():
    ctx = mp.get_context("spawn")
    g = dgl.graph(([0, 1, 2], [1, 2, 3]))
    g.create_formats_()
    q = ctx.Queue(1)
    proc = ctx.Process(target=_test_forking_pickler_entry, args=(g, q))
    proc.start()
    fmt = q.get()["created"]
    proc.join()
    assert "coo" in fmt
    assert "csr" in fmt
    assert "csc" in fmt


if __name__ == "__main__":
    # test_create()
    # test_query()
    # test_hypersparse()
    # test_adj("int32")
    # test_inc()
    # test_view("int32")
    # test_view1("int32")
    # test_flatten(F.int32)
    # test_convert_bound()
    # test_convert()
    # test_to_device("int32")
    # test_transform("int32")
    # test_subgraph("int32")
    # test_subgraph_mask("int32")
    # test_apply()
    # test_level1()
    # test_level2()
    # test_updates()
    # test_backward()
    # test_empty_heterograph('int32')
    # test_types_in_function()
    # test_stack_reduce()
    # test_isolated_ntype()
    # test_bipartite()
    # test_dtype_cast()
    # test_float_cast()
    # test_reverse("int32")
    # test_format()
    # test_add_edges(F.int32)
    # test_add_nodes(F.int32)
    # test_remove_edges(F.int32)
    # test_remove_nodes(F.int32)
    # test_clone(F.int32)
    # test_frame(F.int32)
    # test_frame_device(F.int32)
    # test_empty_query(F.int32)
    # test_create_block(F.int32)
    pass