test_heterograph.py

import dgl
import dgl.function as fn
from collections import Counter
import numpy as np
import scipy.sparse as ssp
import itertools
import backend as F
import networkx as nx
import unittest, pytest
from dgl import DGLError
from utils import parametrize_dtype

def create_test_heterograph(index_dtype):
    # 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')])

    plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
    wishes_nx = nx.DiGraph()
    wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
    wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
    wishes_nx.add_edge('u0', 'g1', id=0)
    wishes_nx.add_edge('u2', 'g0', id=1)

    follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', index_dtype=index_dtype)
    plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game', index_dtype=index_dtype)
    wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game', index_dtype=index_dtype)
    develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game', index_dtype=index_dtype)
    assert follows_g._idtype_str == index_dtype
    assert plays_g._idtype_str == index_dtype
    assert wishes_g._idtype_str == index_dtype
    assert develops_g._idtype_str == index_dtype
    g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
    assert g._idtype_str == index_dtype
    return g

def create_test_heterograph1(index_dtype):
    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
    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)
    g0.ndata[dgl.NTYPE] = ntypes
    g0.edata[dgl.ETYPE] = etypes
    return dgl.to_hetero(g0, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'])

def create_test_heterograph2(index_dtype):
    plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
    wishes_nx = nx.DiGraph()
    wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
    wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
    wishes_nx.add_edge('u0', 'g1', id=0)
    wishes_nx.add_edge('u2', 'g0', id=1)
    develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')

    g = dgl.heterograph({
        ('user', 'follows', 'user'): [(0, 1), (1, 2)],
        ('user', 'plays', 'game'): plays_spmat,
        ('user', 'wishes', 'game'): wishes_nx,
        ('developer', 'develops', 'game'): develops_g,
        })
    return g

def create_test_heterograph3(index_dtype):
    plays_spmat = ssp.coo_matrix(([1, 1, 1, 1], ([0, 1, 2, 1], [0, 0, 1, 1])))
    wishes_nx = nx.DiGraph()
    wishes_nx.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
    wishes_nx.add_nodes_from(['g0', 'g1'], bipartite=1)
    wishes_nx.add_edge('u0', 'g1', id=0)
    wishes_nx.add_edge('u2', 'g0', id=1)

    follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows', _restrict_format='coo')
    plays_g = dgl.bipartite(
        [(0, 0), (1, 0), (2, 1), (1, 1)], 'user', 'plays', 'game', _restrict_format='coo')
    wishes_g = dgl.bipartite([(0, 1), (2, 0)], 'user', 'wishes', 'game', _restrict_format='coo')
    develops_g = dgl.bipartite(
        [(0, 0), (1, 1)], 'developer', 'develops', 'game', _restrict_format='coo')
    g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
    return g

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

@parametrize_dtype
def test_create(index_dtype):
    g0 = create_test_heterograph(index_dtype)
    g1 = create_test_heterograph1(index_dtype)
    g2 = create_test_heterograph2(index_dtype)
    assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
    assert set(g0.canonical_etypes) == set(g1.canonical_etypes) == set(g2.canonical_etypes)

    # create from nx complete bipartite graph
    nxg = nx.complete_bipartite_graph(3, 4)
    g = dgl.bipartite(nxg, 'user', 'plays', 'game')
    assert g.ntypes == ['user', 'game']
    assert g.etypes == ['plays']
    assert g.number_of_edges() == 12

    # create from scipy
    spmat = ssp.coo_matrix(([1,1,1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
    g = dgl.graph(spmat)
    assert g.number_of_nodes() == 4
    assert g.number_of_edges() == 3

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

    # test if validate flag works
    # homo graph
    fail = False
    try:
        g = dgl.graph(
            ([0, 0, 0, 1, 1, 2], [0, 1, 2, 0, 1, 2]),
            num_nodes=2,
            validate=True
        )
    except DGLError:
        fail = True
    finally:
        assert fail, "should catch a DGLError because node ID is out of bound."
    # bipartite graph
    def _test_validate_bipartite(card):
        fail = False
        try:
            g = dgl.bipartite(
                ([0, 0, 1, 1, 2], [1, 1, 2, 2, 3]),
                num_nodes=card,
                validate=True
            )
        except DGLError:
            fail = True
        finally:
            assert fail, "should catch a DGLError because node ID is out of bound."

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

@parametrize_dtype
def test_query(index_dtype):
    g = create_test_heterograph(index_dtype)

    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.number_of_nodes(ntype) for ntype in ntypes] == [3, 2, 2]

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

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

        assert not g.is_multigraph
        assert g.is_readonly

        for etype in etypes:
            srcs, dsts = edges[etype]
            for src, dst in zip(srcs, dsts):
                assert g.has_edge_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_edge_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_degree(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_degree(0, etype) == len(succ)

            # edge_id & edge_ids
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
                assert g.edge_id(src, dst, etype=etype) == i
                assert F.asnumpy(g.edge_id(src, dst, etype=etype, return_array=True)).tolist() == [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)
            assert F.asnumpy(u).tolist() == srcs
            assert F.asnumpy(v).tolist() == dsts
            assert F.asnumpy(e).tolist() == list(range(n_edges))

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

            # all_edges.
            for order in ['eid']:
                u, v, e = g.all_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(index_dtype)
    _test(g)
    g = create_test_heterograph1(index_dtype)
    _test(g)
    g = create_test_heterograph3(index_dtype)
    _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(index_dtype)
    _test(g)
    g = create_test_heterograph1(index_dtype)
    _test(g)
    g = create_test_heterograph3(index_dtype)
    _test(g)

    # test repr
    print(g)

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

    g = dgl.heterograph({
        ('user', 'follows', 'user'): [(0, 1)],
        ('user', 'plays', 'game'): [(0, N2)]},
        {'user': N1, 'game': N1})
    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_edge_between(0, 1, 'follows')
    assert not g.has_edge_between(0, 0, 'follows')
    mask = F.asnumpy(g.has_edges_between([0, 0], [0, 1], 'follows')).tolist()
    assert mask == [0, 1]

    assert g.has_edge_between(0, N2, 'plays')
    assert not g.has_edge_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_id(0, 1, etype='follows') == 0
    assert g.edge_id(0, N2, etype='plays') == 0
    assert F.asnumpy(g.edge_ids(0, 1, etype='follows')).tolist() == [0]
    assert F.asnumpy(g.edge_ids(0, N2, etype='plays')).tolist() == [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_degree(0, 'follows') == 0
    assert g.in_degree(1, 'follows') == 1
    assert F.asnumpy(g.in_degrees([0, 1], 'follows')).tolist() == [0, 1]
    assert g.in_degree(0, 'plays') == 0
    assert g.in_degree(N2, 'plays') == 1
    assert F.asnumpy(g.in_degrees([0, N2], 'plays')).tolist() == [0, 1]
    assert g.out_degree(0, 'follows') == 1
    assert g.out_degree(1, 'follows') == 0
    assert F.asnumpy(g.out_degrees([0, 1], 'follows')).tolist() == [1, 0]
    assert g.out_degree(0, 'plays') == 1
    assert g.out_degree(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'): [(0, 1)],
        ('user', 'plays', 'game'): [(0, N2)]},
        {'user': N1, 'game': N1})
    with pytest.raises(AssertionError):
        eids = g.edge_ids(0, 0, etype='follows')

    with pytest.raises(AssertionError):
        eid = g.edge_id(0, 0, etype='follows')

    g2 = dgl.heterograph({
        ('user', 'follows', 'user'): [(0, 1), (0, 1)],
        ('user', 'plays', 'game'): [(0, N2)]},
        {'user': N1, 'game': N1})

    with pytest.raises(AssertionError):
        eids = g2.edge_ids(0, 1, etype='follows')

    with pytest.raises(AssertionError):
        eid = g2.edge_id(0, 1, etype='follows')

@parametrize_dtype
def test_adj(index_dtype):
    g = create_test_heterograph(index_dtype)
    adj = F.sparse_to_numpy(g.adj(etype='follows'))
    assert np.allclose(
            adj,
            np.array([[0., 0., 0.],
                      [1., 0., 0.],
                      [0., 1., 0.]]))
    adj = F.sparse_to_numpy(g.adj(transpose=True, etype='follows'))
    assert np.allclose(
            adj,
            np.array([[0., 1., 0.],
                      [0., 0., 1.],
                      [0., 0., 0.]]))
    adj = F.sparse_to_numpy(g.adj(etype='plays'))
    assert np.allclose(
            adj,
            np.array([[1., 1., 0.],
                      [0., 1., 1.]]))
    adj = F.sparse_to_numpy(g.adj(transpose=True, etype='plays'))
    assert np.allclose(
            adj,
            np.array([[1., 0.],
                      [1., 1.],
                      [0., 1.]]))

    adj = g.adj(scipy_fmt='csr', etype='follows')
    assert np.allclose(
            adj.todense(),
            np.array([[0., 0., 0.],
                      [1., 0., 0.],
                      [0., 1., 0.]]))
    adj = g.adj(scipy_fmt='coo', etype='follows')
    assert np.allclose(
            adj.todense(),
            np.array([[0., 0., 0.],
                      [1., 0., 0.],
                      [0., 1., 0.]]))
    adj = g.adj(scipy_fmt='csr', etype='plays')
    assert np.allclose(
            adj.todense(),
            np.array([[1., 1., 0.],
                      [0., 1., 1.]]))
    adj = g.adj(scipy_fmt='coo', etype='plays')
    assert np.allclose(
            adj.todense(),
            np.array([[1., 1., 0.],
                      [0., 1., 1.]]))
    adj = F.sparse_to_numpy(g['follows'].adj())
    assert np.allclose(
            adj,
            np.array([[0., 0., 0.],
                      [1., 0., 0.],
                      [0., 1., 0.]]))

@parametrize_dtype
def test_inc(index_dtype):
    g = create_test_heterograph(index_dtype)
    #follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
    adj = F.sparse_to_numpy(g['follows'].inc('in'))
    assert np.allclose(
            adj,
            np.array([[0., 0.],
                      [1., 0.],
                      [0., 1.]]))
    adj = F.sparse_to_numpy(g['follows'].inc('out'))
    assert np.allclose(
            adj,
            np.array([[1., 0.],
                      [0., 1.],
                      [0., 0.]]))
    adj = F.sparse_to_numpy(g['follows'].inc('both'))
    assert np.allclose(
            adj,
            np.array([[-1., 0.],
                      [1., -1.],
                      [0., 1.]]))
    adj = F.sparse_to_numpy(g.inc('in', etype='plays'))
    assert np.allclose(
            adj,
            np.array([[1., 1., 0., 0.],
                      [0., 0., 1., 1.]]))
    adj = F.sparse_to_numpy(g.inc('out', etype='plays'))
    assert np.allclose(
            adj,
            np.array([[1., 0., 0., 0.],
                      [0., 1., 0., 1.],
                      [0., 0., 1., 0.]]))
    adj = F.sparse_to_numpy(g.inc('both', etype='follows'))
    assert np.allclose(
            adj,
            np.array([[-1., 0.],
                      [1., -1.],
                      [0., 1.]]))

@parametrize_dtype
def test_view(index_dtype):
    # test single node type
    g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
    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(index_dtype)

    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(F.tensor(g.nodes('user')), F.arange(0, 3))
    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
    g.ndata['h'] = {'user' : f1,
                    'game' : f2}
    f3 = g.nodes['user'].data['h']
    f4 = g.nodes['game'].data['h']
    assert F.array_equal(f1, f3)
    assert F.array_equal(f2, f4)
    data = g.ndata['h']
    assert F.array_equal(f1, data['user'])
    assert F.array_equal(f2, data['game'])
    # test repr
    print(g.ndata)
    g.ndata.pop('h')
    # test repr
    print(g.ndata)

    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(F.tensor(g.edges(etype='follows', form='eid')), F.arange(0, 2))
    g.edges['follows'].data.pop('h')

    f3 = F.randn((2, 4))
    fail = False
    try:
        g.edata['h'] = f3
    except Exception:
        fail = True
    assert fail
    g.edata['h'] = {('user', 'follows', 'user') : 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)
    data = g.edata['h']
    assert F.array_equal(f3, data[('user', 'follows', 'user')])
    # test repr
    print(g.edata)
    g.edata.pop('h')
    # test repr
    print(g.edata)

    # 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(F.tensor(g.srcnodes('user')), F.arange(0, 3))
    g.srcnodes['user'].data.pop('h')

    # multi type ndata
    f1 = F.randn((3, 6))
    f2 = F.randn((2, 6))
    fail = False
    try:
        g.srcdata['h'] = f1
    except Exception:
        fail = True
    assert fail
    g.srcdata['h'] = {'user' : f1,
                      'developer' : f2}
    f3 = g.srcnodes['user'].data['h']
    f4 = g.srcnodes['developer'].data['h']
    assert F.array_equal(f1, f3)
    assert F.array_equal(f2, f4)
    data = g.srcdata['h']
    assert F.array_equal(f1, data['user'])
    assert F.array_equal(f2, data['developer'])
    # test repr
    print(g.srcdata)
    g.srcdata.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(F.tensor(g.dstnodes('user')), F.arange(0, 3))
    g.dstnodes['user'].data.pop('h')

    # multi type ndata
    f1 = F.randn((3, 6))
    f2 = F.randn((2, 6))
    fail = False
    try:
        g.dstdata['h'] = f1
    except Exception:
        fail = True
    assert fail
    g.dstdata['h'] = {'user' : f1,
                      'game' : f2}
    f3 = g.dstnodes['user'].data['h']
    f4 = g.dstnodes['game'].data['h']
    assert F.array_equal(f1, f3)
    assert F.array_equal(f2, f4)
    data = g.dstdata['h']
    assert F.array_equal(f1, data['user'])
    assert F.array_equal(f2, data['game'])
    # test repr
    print(g.dstdata)
    g.dstdata.pop('h')

@parametrize_dtype
def test_view1(index_dtype):
    # test relation view
    HG = create_test_heterograph(index_dtype)
    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_edge_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_edge_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_degree(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_degree(0) == len(succ)

            # edge_id & edge_ids
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
                assert g.edge_id(src, dst) == i
                assert F.asnumpy(g.edge_id(src, dst, return_array=True)).tolist() == [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)
            assert F.asnumpy(u).tolist() == srcs
            assert F.asnumpy(v).tolist() == dsts
            assert F.asnumpy(e).tolist() == list(range(n_edges))

            # 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(F.tensor(g.nodes()), F.arange(0, 3))

    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(F.tensor(g.edges(form='eid')), F.arange(0, 2))

    # multiple types
    ndata = HG.ndata['h']
    assert isinstance(ndata, dict)
    assert F.array_equal(ndata['user'], f2)
    
    edata = HG.edata['h']
    assert isinstance(edata, dict)
    assert F.array_equal(edata[('user', 'follows', 'user')], f4)

@parametrize_dtype
def test_flatten(index_dtype):
    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(index_dtype)
    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 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([(1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1)])

    check_mapping(g, fg)

    fg = g['user', :, 'user']
    # 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.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.ntypes == ['developer+user', 'game+user']
    assert fg.etypes == ['develops+follows+plays+wishes']
    check_mapping(g, fg)

    # Test another heterograph
    g_x = dgl.graph(([0, 1, 2], [1, 2, 3]), 'user', 'follows', index_dtype=index_dtype)
    g_y = dgl.graph(([0, 2], [2, 3]), 'user', 'knows', index_dtype=index_dtype)
    g_x.nodes['user'].data['h'] = F.randn((4, 3))
    g_x.edges['follows'].data['w'] = F.randn((3, 2))
    g_y.nodes['user'].data['hh'] = F.randn((4, 5))
    g_y.edges['knows'].data['ww'] = F.randn((2, 10))
    g = dgl.hetero_from_relations([g_x, g_y])

    assert F.array_equal(g.ndata['h'], g_x.ndata['h'])
    assert F.array_equal(g.ndata['hh'], g_y.ndata['hh'])
    assert F.array_equal(g.edges['follows'].data['w'], g_x.edata['w'])
    assert F.array_equal(g.edges['knows'].data['ww'], g_y.edata['ww'])

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

    fg = g['user', :, :]
    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_dtype
def test_to_device(index_dtype):
    g = create_test_heterograph(index_dtype)
    g.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cpu())
    g.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cpu())
    g.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cpu())
    if F.is_cuda_available():
        g1 = g.to(F.cuda())
        assert g1 is not None

    # set feature after g.to
    g = create_test_heterograph(index_dtype)
    if F.is_cuda_available():
        g1 = g.to(F.cuda())
        assert g1 is not None
        g1.nodes['user'].data['h'] = F.copy_to(F.ones((3, 5)), F.cuda())
        g1.nodes['game'].data['i'] = F.copy_to(F.ones((2, 5)), F.cuda())
        g1.edges['plays'].data['e'] = F.copy_to(F.ones((4, 4)), F.cuda())

@parametrize_dtype
def test_convert_bound(index_dtype):
    def _test_bipartite_bound(data, card):
        try:
            dgl.bipartite(data, num_nodes=card, index_dtype=index_dtype)
        except dgl.DGLError:
            return
        assert False, 'bipartite bound test with wrong uid failed'

    def _test_graph_bound(data, card):
        try:
            dgl.graph(data, num_nodes=card, index_dtype=index_dtype)
        except dgl.DGLError:
            return
        assert False, 'graph bound test with wrong uid failed'

    _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_dtype
def test_convert(index_dtype):
    hg = create_test_heterograph(index_dtype)
    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_homo(hg)
    assert g._idtype_str == index_dtype
    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_hetero(
                g, hg.ntypes, hg.etypes,
                ntype_field=dgl.NTYPE, etype_field=dgl.ETYPE, metagraph=_mg)
        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, 2), (1, 2), (2, 3), (0, 3)], index_dtype=index_dtype)
    g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
    g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
    hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2'])
    assert hg._idtype_str == index_dtype
    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

    # hetero_from_homo test case 3
    mg = nx.MultiDiGraph([
        ('user', 'movie', 'watches'),
        ('user', 'TV', 'watches')])
    g = dgl.graph([(0, 1), (0, 2)], index_dtype=index_dtype)
    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_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg)
        assert hg._idtype_str == index_dtype
        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.bipartite([(0, 0), (1, 1)], num_nodes=(2, 3))
    g = dgl.to_homo(hg)
    assert g.number_of_nodes() == 5

@parametrize_dtype
def test_transform(index_dtype):
    g = create_test_heterograph(index_dtype)
    x = F.randn((3, 5))
    g.nodes['user'].data['h'] = x

    new_g = dgl.metapath_reachable_graph(g, ['follows', 'plays'])
    assert new_g._idtype_str == index_dtype
    assert new_g.ntypes == ['user', 'game']
    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_str == index_dtype
    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_dtype
def test_subgraph_mask(index_dtype):
    g = create_test_heterograph(index_dtype)
    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.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], F.int64))
        assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
                             F.tensor([0], F.int64))
        assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        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])

    # backend boo input tensor
    sg1 = g.subgraph({'user': F.tensor([False, True, True], dtype=F.data_type_dict['bool']),
                      'game': F.tensor([True, False, False, False], dtype=F.data_type_dict['bool'])})
    _check_subgraph(g, sg1)
    sg2 = g.edge_subgraph({'follows': F.tensor([False, True], dtype=F.data_type_dict['bool']),
                           'plays': F.tensor([False, True, False, False], dtype=F.data_type_dict['bool']),
                           'wishes': F.tensor([False, True], dtype=F.data_type_dict['bool'])})
    _check_subgraph(g, sg2)

@parametrize_dtype
def test_subgraph(index_dtype):
    g = create_test_heterograph(index_dtype)
    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.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], F.int64))
        assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
                             F.tensor([0], F.int64))
        assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
                             F.tensor([1], F.int64))
        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)
    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=F.data_type_dict[index_dtype]),
                      'game': F.tensor([0], dtype=F.data_type_dict[index_dtype])})
    _check_subgraph(g, sg1)
    sg2 = g.edge_subgraph({'follows': F.tensor([1], dtype=F.data_type_dict[index_dtype]),
                           'plays': F.tensor([1], dtype=F.data_type_dict[index_dtype]),
                           'wishes': F.tensor([1], dtype=F.data_type_dict[index_dtype])})
    _check_subgraph(g, sg2)

    # numpy input
    sg1 = g.subgraph({'user': np.array([1, 2]),
                      'game': np.array([0])})
    _check_subgraph(g, sg1)
    sg2 = g.edge_subgraph({'follows': np.array([1]),
                           'plays': np.array([1]),
                           'wishes': np.array([1])})
    _check_subgraph(g, sg2)

    def _check_subgraph_single_ntype(g, sg, preserve_nodes=False):
        assert sg.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], F.int64))
        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], F.int64))

        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], F.int64))
            assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
                                 F.tensor([0], F.int64))
        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], F.int64))

    sg1_graph = g_graph.subgraph([1, 2])
    _check_subgraph_single_ntype(g_graph, sg1_graph)
    sg1_graph = g_graph.edge_subgraph([1])
    _check_subgraph_single_ntype(g_graph, sg1_graph)
    sg1_graph = g_graph.edge_subgraph([1], preserve_nodes=True)
    _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], preserve_nodes=True)
    _check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)

    def _check_typed_subgraph1(g, sg):
        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)

    # Test for restricted format
    for fmt in ['csr', 'csc', 'coo']:
        g = dgl.graph([(0, 1), (1, 2)], restrict_format=fmt)
        sg = g.subgraph({g.ntypes[0]: [1, 0]})
        nids = F.asnumpy(sg.ndata[dgl.NID])
        assert np.array_equal(nids, np.array([1, 0]))
        src, dst = sg.all_edges(order='eid')
        src = F.asnumpy(src)
        dst = F.asnumpy(dst)
        assert np.array_equal(src, np.array([1]))
        assert np.array_equal(dst, np.array([0]))

@parametrize_dtype
def test_apply(index_dtype):
    def node_udf(nodes):
        return {'h': nodes.data['h'] * 2}
    def edge_udf(edges):
        return {'h': edges.data['h'] * 2 + edges.src['h']}

    g = create_test_heterograph(index_dtype)
    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 fail case
    # fail due to multiple types
    fail = False
    try:
        g.apply_nodes(node_udf)
    except dgl.DGLError:
        fail = True
    assert fail

    fail = False
    try:
        g.apply_edges(edge_udf)
    except dgl.DGLError:
        fail = True
    assert fail

@parametrize_dtype
def test_level1(index_dtype):
    #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(index_dtype)
    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}
    g.nodes['user'].data['h'] = F.ones((3, 2))
    g.send([2, 3], mfunc, etype='plays')
    g.recv([0, 1], rfunc, etype='plays')
    y = g.nodes['game'].data['y']
    assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
    g.nodes['game'].data.pop('y')

    # only one type
    play_g = g['plays']
    play_g.send([2, 3], mfunc)
    play_g.recv([0, 1], rfunc)
    y = g.nodes['game'].data['y']
    assert F.array_equal(y, F.tensor([[0., 0.], [2., 2.]]))
    # TODO(minjie): following codes will fail because messages are
    #   not shared with the base graph. However, since send and recv
    #   are rarely used, no fix at the moment.
    # g['plays'].send([2, 3], mfunc)
    # g['plays'].recv([0, 1], mfunc)

    # test fail case
    # fail due to multiple types
    fail = False
    try:
        g.send([2, 3], mfunc)
    except dgl.DGLError:
        fail = True
    assert fail

    fail = False
    try:
        g.recv([0, 1], rfunc)
    except dgl.DGLError:
        fail = True
    assert fail

    # test multi recv
    g.send(g.edges(etype='plays'), mfunc, etype='plays')
    g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
    g.multi_recv([0, 1], {'plays' : rfunc, ('user', 'wishes', 'game'): rfunc2}, 'sum')
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))

    # test multi recv with apply function
    g.send(g.edges(etype='plays'), mfunc, etype='plays')
    g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
    g.multi_recv([0, 1], {'plays' : (rfunc, afunc), ('user', 'wishes', 'game'): rfunc2}, 'sum', afunc)
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))

    # test cross reducer
    g.nodes['user'].data['h'] = F.randn((3, 2))
    for cred in ['sum', 'max', 'min', 'mean']:
        g.send(g.edges(etype='plays'), mfunc, etype='plays')
        g.send(g.edges(etype='wishes'), mfunc, etype='wishes')
        g.multi_recv([0, 1], {'plays' : (rfunc, afunc), 'wishes': rfunc2}, cred, afunc)
        y = g.nodes['game'].data['y']
        g1 = g['plays']
        g2 = g['wishes']
        g1.send(g1.edges(), mfunc)
        g1.recv(g1.nodes('game'), rfunc, afunc)
        y1 = g.nodes['game'].data['y']
        g2.send(g2.edges(), mfunc)
        g2.recv(g2.nodes('game'), rfunc2)
        y2 = g.nodes['game'].data['y']
        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
    fail = False
    try:
        g.multi_recv([0, 1], {'plays' : rfunc, 'follows': rfunc2}, 'sum')
    except dgl.DGLError:
        fail = True
    assert fail


@parametrize_dtype
def test_level2(index_dtype):
    #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(index_dtype)
    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.], [2., 2.]]))

    # 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.], [2., 2.]]))
    
    # test fail case
    # fail due to multiple types
    fail = False
    try:
        g.send_and_recv([2, 3], mfunc, rfunc)
    except dgl.DGLError:
        fail = True
    assert fail

    # test multi
    g.multi_send_and_recv(
        {'plays' : (g.edges(etype='plays'), mfunc, rfunc),
         ('user', 'wishes', 'game'): (g.edges(etype='wishes'), mfunc, rfunc2)},
        'sum')
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[3., 3.], [3., 3.]]))

    # test multi
    g.multi_send_and_recv(
        {'plays' : (g.edges(etype='plays'), mfunc, rfunc, afunc),
         ('user', 'wishes', 'game'): (g.edges(etype='wishes'), mfunc, rfunc2)},
        'sum', afunc)
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[5., 5.], [5., 5.]]))

    # test cross reducer
    g.nodes['user'].data['h'] = F.randn((3, 2))
    for cred in ['sum', 'max', 'min', 'mean']:
        g.multi_send_and_recv(
            {'plays' : (g.edges(etype='plays'), mfunc, rfunc, afunc),
             'wishes': (g.edges(etype='wishes'), mfunc, rfunc2)},
            cred, afunc)
        y = g.nodes['game'].data['y']
        g['plays'].send_and_recv(g.edges(etype='plays'), mfunc, rfunc, afunc)
        y1 = g.nodes['game'].data['y']
        g['wishes'].send_and_recv(g.edges(etype='wishes'), mfunc, rfunc2)
        y2 = g.nodes['game'].data['y']
        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
    fail = False
    try:
        g.multi_send_and_recv(
            {'plays' : (g.edges(etype='plays'), mfunc, rfunc),
             'follows': (g.edges(etype='follows'), mfunc, rfunc2)},
            'sum')
    except dgl.DGLError:
        fail = True
    assert fail

    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.], [2., 2.]]))

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

    # test fail case
    fail = False
    try:
        g.pull(1, mfunc, rfunc)
    except dgl.DGLError:
        fail = True
    assert fail

    # test multi
    g.multi_pull(
        1,
        {'plays' : (mfunc, rfunc),
         ('user', 'wishes', 'game'): (mfunc, rfunc2)},
        'sum')
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[0., 0.], [3., 3.]]))

    # test multi
    g.multi_pull(
        1,
        {'plays' : (mfunc, rfunc, afunc),
         ('user', 'wishes', 'game'): (mfunc, rfunc2)},
        'sum', afunc)
    assert F.array_equal(g.nodes['game'].data['y'], F.tensor([[0., 0.], [5., 5.]]))

    # test cross reducer
    g.nodes['user'].data['h'] = F.randn((3, 2))
    for cred in ['sum', 'max', 'min', 'mean']:
        g.multi_pull(
            1,
            {'plays' : (mfunc, rfunc, afunc),
             'wishes': (mfunc, rfunc2)},
            cred, afunc)
        y = g.nodes['game'].data['y']
        g['plays'].pull(1, mfunc, rfunc, afunc)
        y1 = g.nodes['game'].data['y']
        g['wishes'].pull(1, mfunc, rfunc2)
        y2 = g.nodes['game'].data['y']
        g.nodes['game'].data['y'] = get_redfn(cred)(F.stack([y1, y2], 0), 0)
        g.apply_nodes(afunc, 1, ntype='game')
        yy = g.nodes['game'].data['y']
        assert F.array_equal(y, yy)

    # test fail case
    # fail because cannot infer ntype
    fail = False
    try:
        g.multi_pull(
            1,
            {'plays' : (mfunc, rfunc),
             'follows': (mfunc, rfunc2)},
            'sum')
    except dgl.DGLError:
        fail = True
    assert fail

    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., 2.], [2., 2.]]))

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

    # test fail case
    # fail due to multiple types
    fail = False
    try:
        g.update_all(mfunc, rfunc)
    except dgl.DGLError:
        fail = True
    assert fail

    # 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., 3.], [3., 3.]]))

    # 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., 5.], [5., 5.]]))

    # 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
    fail = False
    try:
        g.update_all(
            {'plays' : (mfunc, rfunc),
             'follows': (mfunc, rfunc2)},
            'sum')
    except dgl.DGLError:
        fail = True
    assert fail

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

@parametrize_dtype
def test_updates(index_dtype):
    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(index_dtype)
    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']

        plays_g = g['user', 'plays', 'game']
        plays_g.send(([0, 1, 2], [0, 1, 1]), msg)
        plays_g.recv([0, 1], 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_dtype
def test_backward(index_dtype):
    g = create_test_heterograph(index_dtype)
    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., 2., 2., 2., 2.],
                                              [2., 2., 2., 2., 2.],
                                              [2., 2., 2., 2., 2.]]))


@parametrize_dtype
def test_empty_heterograph(index_dtype):
    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 edge list
    assert_empty(dgl.heterograph({('user', 'plays', 'game'): []}))
    # empty src-dst pair
    assert_empty(dgl.heterograph({('user', 'plays', 'game'): ([], [])}))
    # empty sparse matrix
    assert_empty(dgl.heterograph({('user', 'plays', 'game'): ssp.coo_matrix((0, 0))}))
    # empty networkx graph
    assert_empty(dgl.heterograph({('user', 'plays', 'game'): nx.DiGraph()}))

    g = dgl.heterograph({('user', 'follows', 'user'): []}, index_dtype=index_dtype)
    assert g._idtype_str == index_dtype
    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, 0), (1, 1)]}, index_dtype=index_dtype)
    assert g._idtype_str == index_dtype
    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


def test_types_in_function():
    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.graph([(0, 1), (1, 2)], 'user', 'follow')
    g.apply_nodes(rfunc1)
    g.apply_edges(mfunc1)
    g.update_all(mfunc1, rfunc1)
    g.send_and_recv([0, 1], mfunc1, rfunc1)
    g.send([0, 1], mfunc1)
    g.recv([1, 2], rfunc1)
    g.push([0], mfunc1, rfunc1)
    g.pull([1], mfunc1, rfunc1)
    g.filter_nodes(filter_nodes1)
    g.filter_edges(filter_edges1)

    g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
    g.apply_nodes(rfunc2, ntype='game')
    g.apply_edges(mfunc2)
    g.update_all(mfunc2, rfunc2)
    g.send_and_recv([0, 1], mfunc2, rfunc2)
    g.send([0, 1], mfunc2)
    g.recv([1, 2], rfunc2)
    g.push([0], mfunc2, rfunc2)
    g.pull([1], mfunc2, rfunc2)
    g.filter_nodes(filter_nodes2, ntype='game')
    g.filter_edges(filter_edges2)

@parametrize_dtype
def test_stack_reduce(index_dtype):
    #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(index_dtype)
    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_dtype
def test_isolated_ntype(index_dtype):
    g = dgl.heterograph({
        ('A', 'AB', 'B'): [(0, 1), (1, 2), (2, 3)]},
        num_nodes_dict={'A': 3, 'B': 4, 'C': 4})
    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), (1, 2), (2, 3)]},
        num_nodes_dict={'A': 3, 'B': 4, 'C': 4})
    assert g.number_of_nodes('A') == 3
    assert g.number_of_nodes('B') == 4
    assert g.number_of_nodes('C') == 4

    G = dgl.DGLGraph()
    G.add_nodes(11)
    G.add_edges([0, 1, 2], [4, 5, 6])
    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_hetero(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_dtype
def test_ismultigraph(index_dtype):
    g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5), (2, 5)], 'A',
                       'AB', 'B', num_nodes=(6, 6), index_dtype=index_dtype)
    assert g1.is_multigraph == False
    g2 = dgl.bipartite([(0, 1), (0, 1), (0, 2), (1, 5)], 'A',
                       'AC', 'C', num_nodes=(6, 6), index_dtype=index_dtype)
    assert g2.is_multigraph == True
    g3 = dgl.graph([(0, 1), (1, 2)], 'A', 'AA',
                   num_nodes=6, index_dtype=index_dtype)
    assert g3.is_multigraph == False
    g4 = dgl.graph([(0, 1), (0, 1), (1, 2)], 'A', 'AA',
                   num_nodes=6, index_dtype=index_dtype)
    assert g4.is_multigraph == True
    g = dgl.hetero_from_relations([g1, g3])
    assert g.is_multigraph == False
    g = dgl.hetero_from_relations([g1, g2])
    assert g.is_multigraph == True
    g = dgl.hetero_from_relations([g1, g4])
    assert g.is_multigraph == True
    g = dgl.hetero_from_relations([g2, g4])
    assert g.is_multigraph == True

@parametrize_dtype
def test_bipartite(index_dtype):
    g1 = dgl.bipartite([(0, 1), (0, 2), (1, 5)], 'A', 'AB', 'B', index_dtype=index_dtype)
    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.bipartite([(1, 0), (0, 0)], 'A', 'AC', 'C', index_dtype=index_dtype)
    g3 = dgl.hetero_from_relations([g1, g2])
    assert g3.is_unibipartite
    assert g3.srctypes == ['A']
    assert set(g3.dsttypes) == {'B', 'C'}
    assert g3.number_of_nodes('A') == 2
    assert g3.number_of_nodes('B') == 6
    assert g3.number_of_nodes('C') == 1
    assert g3.number_of_src_nodes('A') == 2
    assert g3.number_of_src_nodes() == 2
    assert g3.number_of_dst_nodes('B') == 6
    assert g3.number_of_dst_nodes('C') == 1
    g3.srcdata['h'] = F.randn((2, 5))
    assert F.array_equal(g3.srcnodes['A'].data['h'], g3.srcdata['h'])
    assert F.array_equal(g3.nodes['A'].data['h'], g3.srcdata['h'])
    assert F.array_equal(g3.nodes['SRC/A'].data['h'], g3.srcdata['h'])

    g4 = dgl.graph([(0, 0), (1, 1)], 'A', 'AA', index_dtype=index_dtype)
    g5 = dgl.hetero_from_relations([g1, g2, g4])
    assert not g5.is_unibipartite

@parametrize_dtype
def test_dtype_cast(index_dtype):
    g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], index_dtype=index_dtype)
    assert g._idtype_str == index_dtype
    g.ndata["feat"] = F.tensor([3, 4, 5])
    g.edata["h"] = F.tensor([3, 4, 5, 6])
    if index_dtype == "int32":
        g_cast = g.long()
        assert g_cast._idtype_str == 'int64'
    else:
        g_cast = g.int()
        assert g_cast._idtype_str == 'int32'
    assert "feat" in g_cast.ndata
    assert "h" in g_cast.edata
    assert F.array_equal(g.ndata["feat"], g_cast.ndata["feat"])
    assert F.array_equal(g.edata["h"], g_cast.edata["h"])

def test_format():
    # single relation
    g = dgl.graph([(0, 0), (1, 1), (0, 1), (2, 0)], restrict_format='coo')
    assert g.restrict_format() == 'coo'
    assert g.format_in_use() == ['coo']
    try:
        spmat = g.adjacency_matrix(scipy_fmt="csr")
    except:
        print('test passed, graph with restrict_format coo should not create csr matrix.')
    else:
        assert False, 'cannot create csr when restrict_format is coo'
    g1 = g.to_format('any')
    assert g1.restrict_format() == 'any'
    g1.request_format('coo')
    g1.request_format('csr')
    g1.request_format('csc')
    assert len(g1.format_in_use()) == 3
    assert g.restrict_format() == 'coo'
    assert g.format_in_use() == ['coo']

    # multiple relation
    g = dgl.heterograph({
        ('user', 'follows', 'user'): [(0, 1), (1, 2)],
        ('user', 'plays', 'game'): [(0, 0), (1, 0), (1, 1), (2, 1)],
        ('developer', 'develops', 'game'): [(0, 0), (1, 1)],
        }, restrict_format='csr')
    user_feat = F.randn((g['follows'].number_of_src_nodes(), 5))
    g['follows'].srcdata['h'] = user_feat
    for rel_type in ['follows', 'plays', 'develops']:
        assert g.restrict_format(rel_type) == 'csr'
        assert g.format_in_use(rel_type) == ['csr']
        try:
            g[rel_type].request_format('coo')
        except:
            print('test passed, graph with restrict_format csr should not create coo matrix')
        else:
            assert False, 'cannot create coo when restrict_format is csr'

    g1 = g.to_format('csc')
    # test frame
    assert F.array_equal(g1['follows'].srcdata['h'], user_feat)
    # test each relation graph
    for rel_type in ['follows', 'plays', 'develops']:
        assert g1.restrict_format(rel_type) == 'csc'
        assert g1.format_in_use(rel_type) == ['csc']
        assert g.restrict_format(rel_type) == 'csr'
        assert g.format_in_use(rel_type) == ['csr']

def test_edges_order():
    # (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])
    ))

    src, dst = g.all_edges(order='srcdst')
    assert F.array_equal(F.copy_to(src, F.cpu()),
                         F.copy_to(F.tensor([0, 0, 0, 1, 2]), F.cpu()))
    assert F.array_equal(F.copy_to(dst, F.cpu()),
                         F.copy_to(F.tensor([1, 1, 2, 2, 1]), F.cpu()))

@parametrize_dtype
def test_reverse(index_dtype):
    g = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 1, 2, 4, 3 ,1, 3], [1, 2, 3, 2, 0, 0, 1]),
    }, index_dtype=index_dtype)
    gidx = g._graph
    r_gidx = gidx.reverse(gidx.metagraph)

    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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())

    # force to start with 'csr'
    gidx = gidx.to_format('csr')
    gidx = gidx.to_format('any')
    r_gidx = gidx.reverse(gidx.metagraph)
    assert gidx.format_in_use(0)[0] == 'csr'
    assert r_gidx.format_in_use(0)[0] == 'csc'
    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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())

    # force to start with 'csc'
    gidx = gidx.to_format('csc')
    gidx = gidx.to_format('any')
    r_gidx = gidx.reverse(gidx.metagraph)
    assert gidx.format_in_use(0)[0] == 'csc'
    assert r_gidx.format_in_use(0)[0] == 'csr'
    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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())

    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]),
        }, index_dtype=index_dtype)
    gidx = g._graph
    r_gidx = gidx.reverse(gidx.metagraph)
    # 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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())

    # force to start with 'csr'
    gidx = gidx.to_format('csr')
    gidx = gidx.to_format('any')
    r_gidx = gidx.reverse(gidx.metagraph)
    # three node types and three edge types
    assert gidx.format_in_use(0)[0] == 'csr'
    assert r_gidx.format_in_use(0)[0] == 'csc'
    assert gidx.format_in_use(1)[0] == 'csr'
    assert r_gidx.format_in_use(1)[0] == 'csc'
    assert gidx.format_in_use(2)[0] == 'csr'
    assert r_gidx.format_in_use(2)[0] == 'csc'
    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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())

    # force to start with 'csc'
    gidx = gidx.to_format('csc')
    gidx = gidx.to_format('any')
    r_gidx = gidx.reverse(gidx.metagraph)
    # three node types and three edge types
    assert gidx.format_in_use(0)[0] == 'csc'
    assert r_gidx.format_in_use(0)[0] == 'csr'
    assert gidx.format_in_use(1)[0] == 'csc'
    assert r_gidx.format_in_use(1)[0] == 'csr'
    assert gidx.format_in_use(2)[0] == 'csc'
    assert r_gidx.format_in_use(2)[0] == 'csr'
    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.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
    assert F.array_equal(g_s.tousertensor(), rg_d.tousertensor())
    assert F.array_equal(g_d.tousertensor(), rg_s.tousertensor())


if __name__ == '__main__':
    # test_create()
    # test_query()
    # test_hypersparse()
    # test_adj("int32")
    # test_inc()
    # test_view("int32")
    # test_view1("int32")
    # test_flatten()
    # 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_reverse("int32")
    test_format()
    pass