[Refactor] Enable new kernel in all message passing APIs (#1953)

* WIP: frame refactor

* new frame

* simple update_all builtin

* move all subgraph routines into the same file

* sddmm & spmm schedule; node & edge udf

* degree bucketing

* some tricky 0deg corner cases

* bug in frame append

* merge test_hetero_basics and test_basics

* some code rearange

* fix test_heterograph

* add mean spmm

* enable all builtin combinations

* pass gpu test

* pass pytorch tests

* wip

* fix some pt debugging codes

* fix bug in mxnet backward

* pass all mxnet utests

* passed tf tests

* docstring

* lint

* lint

* fix broadcasting bugs

* add warning and clamp for mean reducer

* add test for zero-degree mean

* address comments

* lint

* small fix

python/dgl/ops/sddmm.py

@@ -3,6 +3,7 @@ from itertools import product
 import sys
 
 from ..backend import gsddmm as gsddmm_internal
+from .. import backend as F
 
 __all__ = ['gsddmm', 'copy_u', 'copy_v']
 
@@ -41,6 +42,21 @@ def gsddmm(g, op, lhs_data, rhs_data, lhs_target='u', rhs_target='v'):
     tensor
         The result tensor.
     """
+    if op not in ['copy_lhs', 'copy_rhs']:
+        # Expand dims so that there will be no broadcasting issues with different
+        # number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
+        # that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
+        # (E, 5, 3, 4)
+        lhs_shape = F.shape(lhs_data)
+        rhs_shape = F.shape(rhs_data)
+        if len(lhs_shape) != len(rhs_shape):
+            max_ndims = max(len(lhs_shape), len(rhs_shape))
+            lhs_pad_ndims = max_ndims - len(lhs_shape)
+            rhs_pad_ndims = max_ndims - len(rhs_shape)
+            new_lhs_shape = (lhs_shape[0],) + (1,) * lhs_pad_ndims + lhs_shape[1:]
+            new_rhs_shape = (rhs_shape[0],) + (1,) * rhs_pad_ndims + rhs_shape[1:]
+            lhs_data = F.reshape(lhs_data, new_lhs_shape)
+            rhs_data = F.reshape(rhs_data, new_rhs_shape)
     return gsddmm_internal(
         g._graph, op, lhs_data, rhs_data, lhs_target, rhs_target)
 

python/dgl/ops/spmm.py

 """dgl spmm operator module."""
 import sys
+
+from ..base import dgl_warning
 from ..backend import gspmm as gspmm_internal
+from .. import backend as F
 
 __all__ = ['gspmm']
 
@@ -30,7 +33,7 @@ def gspmm(g, op, reduce_op, lhs_data, rhs_data):
         The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
         ``copy_lhs``, ``copy_rhs``.
     reduce_op : str
-        Reduce operator, could be ``sum``, ``max``, ``min``.
+        Reduce operator, could be ``sum``, ``max``, ``min``, ``mean``.
     lhs_data : tensor or None
         The left operand, could be None if it's not required by the op.
     rhs_data : tensor or None
@@ -41,7 +44,32 @@ def gspmm(g, op, reduce_op, lhs_data, rhs_data):
     tensor
         The result tensor.
     """
-    return gspmm_internal(g._graph, op, reduce_op, lhs_data, rhs_data)
+    if op not in ['copy_lhs', 'copy_rhs']:
+        # Expand dims so that there will be no broadcasting issues with different
+        # number of dimensions. For example, given two shapes (N, 3, 1), (E, 5, 3, 4)
+        # that are valid broadcastable shapes, change them to (N, 1, 3, 1) and
+        # (E, 5, 3, 4)
+        lhs_shape = F.shape(lhs_data)
+        rhs_shape = F.shape(rhs_data)
+        if len(lhs_shape) != len(rhs_shape):
+            max_ndims = max(len(lhs_shape), len(rhs_shape))
+            lhs_pad_ndims = max_ndims - len(lhs_shape)
+            rhs_pad_ndims = max_ndims - len(rhs_shape)
+            new_lhs_shape = (lhs_shape[0],) + (1,) * lhs_pad_ndims + lhs_shape[1:]
+            new_rhs_shape = (rhs_shape[0],) + (1,) * rhs_pad_ndims + rhs_shape[1:]
+            lhs_data = F.reshape(lhs_data, new_lhs_shape)
+            rhs_data = F.reshape(rhs_data, new_rhs_shape)
+    if reduce_op == 'mean':
+        ret = gspmm_internal(g._graph, op, 'sum', lhs_data, rhs_data)
+        ret_shape = F.shape(ret)
+        deg = g.in_degrees()
+        if F.as_scalar(F.min(deg, dim=0)) == 0:
+            dgl_warning('Zero-degree nodes encountered in mean reducer. Setting the mean to 0.')
+        deg = F.astype(F.clamp(deg, 1, g.number_of_edges()), F.dtype(ret))
+        deg_shape = (ret_shape[0],) + (1,) * (len(ret_shape) - 1)
+        return ret / F.reshape(deg, deg_shape)
+    else:
+        return gspmm_internal(g._graph, op, reduce_op, lhs_data, rhs_data)
 
 
 def _gen_spmm_func(binary_op, reduce_op):
@@ -130,9 +158,14 @@ def _gen_copy_reduce_func(binary_op, reduce_op):
 
 
 def _register_spmm_func():
-    """Register spmm functions"""
+    """Register spmm functions
+
+    - Binary operation plus reduction between u and e: u_[]_e_[]
+    - Copy u plus reduction: copy_u_[]
+    - Copy e plus reduction: copy_e_[]
+    """
     for binary_op in ["add", "sub", "mul", "div", "copy_u", "copy_e"]:
-        for reduce_op in ["sum", "max", "min"]:
+        for reduce_op in ["sum", "max", "min", "mean"]:
             if binary_op.startswith("copy"):
                 func = _gen_copy_reduce_func(binary_op, reduce_op)
             else:

python/dgl/sampling/neighbor.py

@@ -6,7 +6,7 @@ from ..base import DGLError, EID
 from ..heterograph import DGLHeteroGraph
 from .. import ndarray as nd
 from .. import utils
-from .. import transform
+from .. import subgraph as subg
 from .dataloader import BlockSampler, assign_block_eids
 
 __all__ = [
@@ -300,7 +300,7 @@ class MultiLayerNeighborSampler(BlockSampler):
     def sample_frontier(self, block_id, g, seed_nodes, *args, **kwargs):
         fanout = self.fanouts[block_id]
         if fanout is None:
-            frontier = transform.in_subgraph(g, seed_nodes)
+            frontier = subg.in_subgraph(g, seed_nodes)
         else:
             frontier = sample_neighbors(g, seed_nodes, fanout, replace=self.replace)
         return frontier

python/dgl/subgraph.py

+"""Functions for extracting subgraphs.
+
+The module only contains functions for extracting subgraphs deterministically.
+For stochastic subgraph extraction, please see functions under :mod:`dgl.sampling`.
+"""
+from collections.abc import Mapping
+
+from ._ffi.function import _init_api
+from .base import DGLError
+from . import backend as F
+from . import graph_index
+from . import heterograph_index
+from . import ndarray as nd
+from .heterograph import DGLHeteroGraph
+from . import utils
+
+__all__ = ['node_subgraph', 'edge_subgraph', 'node_type_subgraph', 'edge_type_subgraph',
+           'in_subgraph', 'out_subgraph']
+
+def node_subgraph(graph, nodes):
+    """Return the subgraph induced on given nodes.
+
+    The metagraph of the returned subgraph is the same as the parent graph.
+    Features are copied from the original graph.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The graph to extract subgraphs from.
+    nodes : list or dict[str->list or iterable]
+        A dictionary mapping node types to node ID array for constructing
+        subgraph. All nodes must exist in the graph.
+
+        If the graph only has one node type, one can just specify a list,
+        tensor, or any iterable of node IDs intead.
+
+        The node ID array can be either an interger tensor or a bool tensor.
+        When a bool tensor is used, it is automatically converted to
+        an interger tensor using the semantic of np.where(nodes_idx == True).
+
+        Note: When using bool tensor, only backend (torch, tensorflow, mxnet)
+        tensors are supported.
+
+    Returns
+    -------
+    G : DGLHeteroGraph
+        The subgraph.
+
+        The nodes and edges in the subgraph are relabeled using consecutive
+        integers from 0.
+
+        One can retrieve the mapping from subgraph node/edge ID to parent
+        node/edge ID via ``dgl.NID`` and ``dgl.EID`` node/edge features of the
+        subgraph.
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    Instantiate a heterograph.
+
+    >>> plays_g = dgl.bipartite(([0, 1, 1, 2], [0, 0, 2, 1]), 'user', 'plays', 'game')
+    >>> follows_g = dgl.graph(([0, 1, 1], [1, 2, 2]), 'user', 'follows')
+    >>> g = dgl.hetero_from_relations([plays_g, follows_g])
+    >>> # Set node features
+    >>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
+
+    Get subgraphs.
+
+    >>> g.subgraph({'user': [4, 5]})
+    An error occurs as these nodes do not exist.
+    >>> sub_g = g.subgraph({'user': [1, 2]})
+    >>> print(sub_g)
+    Graph(num_nodes={'user': 2, 'game': 0},
+          num_edges={('user', 'plays', 'game'): 0, ('user', 'follows', 'user'): 2},
+          metagraph=[('user', 'game'), ('user', 'user')])
+
+    Get subgraphs using boolean mask tensor.
+
+    >>> sub_g = g.subgraph({'user': th.tensor([False, True, True])})
+    >>> print(sub_g)
+    Graph(num_nodes={'user': 2, 'game': 0},
+          num_edges={('user', 'plays', 'game'): 0, ('user', 'follows', 'user'): 2},
+          metagraph=[('user', 'game'), ('user', 'user')])
+
+    Get the original node/edge indices.
+
+    >>> sub_g['follows'].ndata[dgl.NID] # Get the node indices in the raw graph
+    tensor([1, 2])
+    >>> sub_g['follows'].edata[dgl.EID] # Get the edge indices in the raw graph
+    tensor([1, 2])
+
+    Get the copied node features.
+
+    >>> sub_g.nodes['user'].data['h']
+    tensor([[1.],
+            [2.]])
+    >>> sub_g.nodes['user'].data['h'] += 1
+    >>> g.nodes['user'].data['h']          # Features are not shared.
+    tensor([[0.],
+            [1.],
+            [2.]])
+
+    See Also
+    --------
+    edge_subgraph
+    """
+    if graph.is_block:
+        raise DGLError('Extracting subgraph from a block graph is not allowed.')
+    if not isinstance(nodes, Mapping):
+        assert len(graph.ntypes) == 1, \
+            'need a dict of node type and IDs for graph with multiple node types'
+        nodes = {graph.ntypes[0]: nodes}
+
+    def _process_nodes(ntype, v):
+        if F.is_tensor(v) and F.dtype(v) == F.bool:
+            return F.astype(F.nonzero_1d(F.copy_to(v, graph.device)), graph.idtype)
+        else:
+            return utils.prepare_tensor(graph, v, 'nodes["{}"]'.format(ntype))
+    induced_nodes = [_process_nodes(ntype, nodes.get(ntype, [])) for ntype in graph.ntypes]
+    sgi = graph._graph.node_subgraph(induced_nodes)
+    induced_edges = sgi.induced_edges
+    return _create_hetero_subgraph(graph, sgi, induced_nodes, induced_edges)
+
+DGLHeteroGraph.subgraph = node_subgraph
+
+def edge_subgraph(graph, edges, preserve_nodes=False):
+    """Return the subgraph induced on given edges.
+
+    The metagraph of the returned subgraph is the same as the parent graph.
+
+    Features are copied from the original graph.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The graph to extract subgraphs from.
+    edges : dict[(str, str, str), Tensor]
+        A dictionary mapping edge types to edge ID array for constructing
+        subgraph. All edges must exist in the subgraph.
+
+        The edge types are characterized by triplets of
+        ``(src type, etype, dst type)``.
+
+        If the graph only has one edge type, one can just specify a list,
+        tensor, or any iterable of edge IDs intead.
+
+        The edge ID array can be either an interger tensor or a bool tensor.
+        When a bool tensor is used, it is automatically converted to
+        an interger tensor using the semantic of np.where(edges_idx == True).
+
+        Note: When using bool tensor, only backend (torch, tensorflow, mxnet)
+        tensors are supported.
+
+    preserve_nodes : bool
+        Whether to preserve all nodes or not. If false, all nodes
+        without edges will be removed. (Default: False)
+
+    Returns
+    -------
+    G : DGLHeteroGraph
+        The subgraph.
+
+        The nodes and edges are relabeled using consecutive integers from 0.
+
+        One can retrieve the mapping from subgraph node/edge ID to parent
+        node/edge ID via ``dgl.NID`` and ``dgl.EID`` node/edge features of the
+        subgraph.
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    Instantiate a heterograph.
+
+    >>> plays_g = dgl.bipartite(([0, 1, 1, 2], [0, 0, 2, 1]), 'user', 'plays', 'game')
+    >>> follows_g = dgl.graph(([0, 1, 1], [1, 2, 2]), 'user', 'follows')
+    >>> g = dgl.hetero_from_relations([plays_g, follows_g])
+    >>> # Set edge features
+    >>> g.edges['follows'].data['h'] = torch.tensor([[0.], [1.], [2.]])
+
+    Get subgraphs.
+
+    >>> g.edge_subgraph({('user', 'follows', 'user'): [5, 6]})
+    An error occurs as these edges do not exist.
+    >>> sub_g = g.edge_subgraph({('user', 'follows', 'user'): [1, 2],
+    >>>                          ('user', 'plays', 'game'): [2]})
+    >>> print(sub_g)
+    Graph(num_nodes={'user': 2, 'game': 1},
+          num_edges={('user', 'plays', 'game'): 1, ('user', 'follows', 'user'): 2},
+          metagraph=[('user', 'game'), ('user', 'user')])
+
+    Get subgraphs using boolean mask tensor.
+    >>> sub_g = g.edge_subgraph({('user', 'follows', 'user'): th.tensor([False, True, True]),
+    >>>                   ('user', 'plays', 'game'): th.tensor([False, False, True, False])})
+    >>> sub_g
+    Graph(num_nodes={'user': 2, 'game': 1},
+        num_edges={('user', 'plays', 'game'): 1, ('user', 'follows', 'user'): 2},
+        metagraph=[('user', 'game'), ('user', 'user')])
+
+    Get the original node/edge indices.
+
+    >>> sub_g['follows'].ndata[dgl.NID] # Get the node indices in the raw graph
+    tensor([1, 2])
+    >>> sub_g['plays'].edata[dgl.EID]   # Get the edge indices in the raw graph
+    tensor([2])
+
+    Get the copied node features.
+
+    >>> sub_g.edges['follows'].data['h']
+    tensor([[1.],
+            [2.]])
+    >>> sub_g.edges['follows'].data['h'] += 1
+    >>> g.edges['follows'].data['h']          # Features are not shared.
+    tensor([[0.],
+            [1.],
+            [2.]])
+
+    See Also
+    --------
+    subgraph
+    """
+    if graph.is_block:
+        raise DGLError('Extracting subgraph from a block graph is not allowed.')
+    if not isinstance(edges, Mapping):
+        assert len(graph.canonical_etypes) == 1, \
+            'need a dict of edge type and IDs for graph with multiple edge types'
+        edges = {graph.canonical_etypes[0]: edges}
+
+    def _process_edges(etype, e):
+        if F.is_tensor(e) and F.dtype(e) == F.bool:
+            return F.astype(F.nonzero_1d(F.copy_to(e, graph.device)), graph.idtype)
+        else:
+            return utils.prepare_tensor(graph, e, 'edges["{}"]'.format(etype))
+
+    edges = {graph.to_canonical_etype(etype): e for etype, e in edges.items()}
+    induced_edges = [
+        _process_edges(cetype, edges.get(cetype, []))
+        for cetype in graph.canonical_etypes]
+    sgi = graph._graph.edge_subgraph(induced_edges, preserve_nodes)
+    induced_nodes = sgi.induced_nodes
+    return _create_hetero_subgraph(graph, sgi, induced_nodes, induced_edges)
+
+DGLHeteroGraph.edge_subgraph = edge_subgraph
+
+def in_subgraph(g, nodes):
+    """Extract the subgraph containing only the in edges of the given nodes.
+
+    The subgraph keeps the same type schema and the cardinality of the original one.
+    Node/edge features are not preserved. The original IDs
+    the extracted edges are stored as the `dgl.EID` feature in the returned graph.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        Full graph structure.
+    nodes : tensor or dict
+        Node ids to sample neighbors from. The allowed types
+        are dictionary of node types to node id tensors, or simply node id tensor if
+        the given graph g has only one type of nodes.
+
+    Returns
+    -------
+    DGLHeteroGraph
+        The subgraph.
+    """
+    if g.is_block:
+        raise DGLError('Extracting subgraph of a block graph is not allowed.')
+    if not isinstance(nodes, dict):
+        if len(g.ntypes) > 1:
+            raise DGLError("Must specify node type when the graph is not homogeneous.")
+        nodes = {g.ntypes[0] : nodes}
+    nodes = utils.prepare_tensor_dict(g, nodes, 'nodes')
+    nodes_all_types = []
+    for ntype in g.ntypes:
+        if ntype in nodes:
+            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
+        else:
+            nodes_all_types.append(nd.NULL[g._idtype_str])
+
+    sgi = _CAPI_DGLInSubgraph(g._graph, nodes_all_types)
+    induced_edges = sgi.induced_edges
+    return _create_hetero_subgraph(g, sgi, None, induced_edges)
+
+DGLHeteroGraph.in_subgraph = in_subgraph
+
+def out_subgraph(g, nodes):
+    """Extract the subgraph containing only the out edges of the given nodes.
+
+    The subgraph keeps the same type schema and the cardinality of the original one.
+    Node/edge features are not preserved. The original IDs
+    the extracted edges are stored as the `dgl.EID` feature in the returned graph.
+
+    Parameters
+    ----------
+    g : DGLHeteroGraph
+        Full graph structure.
+    nodes : tensor or dict
+        Node ids to sample neighbors from. The allowed types
+        are dictionary of node types to node id tensors, or simply node id tensor if
+        the given graph g has only one type of nodes.
+
+    Returns
+    -------
+    DGLHeteroGraph
+        The subgraph.
+    """
+    if g.is_block:
+        raise DGLError('Extracting subgraph of a block graph is not allowed.')
+    if not isinstance(nodes, dict):
+        if len(g.ntypes) > 1:
+            raise DGLError("Must specify node type when the graph is not homogeneous.")
+        nodes = {g.ntypes[0] : nodes}
+    nodes = utils.prepare_tensor_dict(g, nodes, 'nodes')
+    nodes_all_types = []
+    for ntype in g.ntypes:
+        if ntype in nodes:
+            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
+        else:
+            nodes_all_types.append(nd.NULL[g._idtype_str])
+
+    sgi = _CAPI_DGLOutSubgraph(g._graph, nodes_all_types)
+    induced_edges = sgi.induced_edges
+    return _create_hetero_subgraph(g, sgi, None, induced_edges)
+
+DGLHeteroGraph.out_subgraph = out_subgraph
+
+def node_type_subgraph(graph, ntypes):
+    """Return the subgraph induced on given node types.
+
+    The metagraph of the returned subgraph is the subgraph of the original
+    metagraph induced from the node types.
+
+    Features are shared with the original graph.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The graph to extract subgraphs from.
+    ntypes : list[str]
+        The node types
+
+    Returns
+    -------
+    G : DGLHeteroGraph
+        The subgraph.
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    Instantiate a heterograph.
+
+    >>> plays_g = dgl.bipartite(([0, 1, 1, 2], [0, 0, 2, 1]), 'user', 'plays', 'game')
+    >>> follows_g = dgl.graph(([0, 1, 1], [1, 2, 2]), 'user', 'follows')
+    >>> g = dgl.hetero_from_relations([plays_g, follows_g])
+    >>> # Set node features
+    >>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
+
+    Get subgraphs.
+
+    >>> sub_g = g.node_type_subgraph(['user'])
+    >>> print(sub_g)
+    Graph(num_nodes=3, num_edges=3,
+          ndata_schemes={'h': Scheme(shape=(1,), dtype=torch.float32)}
+          edata_schemes={})
+
+    Get the shared node features.
+
+    >>> sub_g.nodes['user'].data['h']
+    tensor([[0.],
+            [1.],
+            [2.]])
+    >>> sub_g.nodes['user'].data['h'] += 1
+    >>> g.nodes['user'].data['h']          # Features are shared.
+    tensor([[1.],
+            [2.],
+            [3.]])
+
+    See Also
+    --------
+    edge_type_subgraph
+    """
+    ntid = [graph.get_ntype_id(ntype) for ntype in ntypes]
+    stids, dtids, etids = graph._graph.metagraph.edges('eid')
+    stids, dtids, etids = stids.tonumpy(), dtids.tonumpy(), etids.tonumpy()
+    etypes = []
+    for stid, dtid, etid in zip(stids, dtids, etids):
+        if stid in ntid and dtid in ntid:
+            etypes.append(graph.canonical_etypes[etid])
+    return edge_type_subgraph(graph, etypes)
+
+DGLHeteroGraph.node_type_subgraph = node_type_subgraph
+
+def edge_type_subgraph(graph, etypes):
+    """Return the subgraph induced on given edge types.
+
+    The metagraph of the returned subgraph is the subgraph of the original metagraph
+    induced from the edge types.
+
+    Features are shared with the original graph.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The graph to extract subgraphs from.
+    etypes : list[str or tuple]
+        The edge types
+
+    Returns
+    -------
+    G : DGLHeteroGraph
+        The subgraph.
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    Instantiate a heterograph.
+
+    >>> plays_g = dgl.bipartite(([0, 1, 1, 2], [0, 0, 2, 1]), 'user', 'plays', 'game')
+    >>> follows_g = dgl.graph(([0, 1, 1], [1, 2, 2]), 'user', 'follows')
+    >>> g = dgl.hetero_from_relations([plays_g, follows_g])
+    >>> # Set edge features
+    >>> g.edges['follows'].data['h'] = torch.tensor([[0.], [1.], [2.]])
+
+    Get subgraphs.
+
+    >>> sub_g = g.edge_type_subgraph(['follows'])
+    >>> print(sub_g)
+    Graph(num_nodes=3, num_edges=3,
+          ndata_schemes={}
+          edata_schemes={'h': Scheme(shape=(1,), dtype=torch.float32)})
+
+    Get the shared edge features.
+
+    >>> sub_g.edges['follows'].data['h']
+    tensor([[0.],
+            [1.],
+            [2.]])
+    >>> sub_g.edges['follows'].data['h'] += 1
+    >>> g.edges['follows'].data['h']          # Features are shared.
+    tensor([[1.],
+            [2.],
+            [3.]])
+
+    See Also
+    --------
+    node_type_subgraph
+    """
+    etype_ids = [graph.get_etype_id(etype) for etype in etypes]
+    # meta graph is homograph, still using int64
+    meta_src, meta_dst, _ = graph._graph.metagraph.find_edges(utils.toindex(etype_ids, "int64"))
+    rel_graphs = [graph._graph.get_relation_graph(i) for i in etype_ids]
+    meta_src = meta_src.tonumpy()
+    meta_dst = meta_dst.tonumpy()
+    ntypes_invmap = {n: i for i, n in enumerate(set(meta_src) | set(meta_dst))}
+    mapped_meta_src = [ntypes_invmap[v] for v in meta_src]
+    mapped_meta_dst = [ntypes_invmap[v] for v in meta_dst]
+    node_frames = [graph._node_frames[i] for i in ntypes_invmap]
+    edge_frames = [graph._edge_frames[i] for i in etype_ids]
+    induced_ntypes = [graph._ntypes[i] for i in ntypes_invmap]
+    induced_etypes = [graph._etypes[i] for i in etype_ids]   # get the "name" of edge type
+    num_nodes_per_induced_type = [graph.number_of_nodes(ntype) for ntype in induced_ntypes]
+
+    metagraph = graph_index.from_edge_list((mapped_meta_src, mapped_meta_dst), True)
+    # num_nodes_per_type should be int64
+    hgidx = heterograph_index.create_heterograph_from_relations(
+        metagraph, rel_graphs, utils.toindex(num_nodes_per_induced_type, "int64"))
+    hg = DGLHeteroGraph(hgidx, induced_ntypes, induced_etypes, node_frames, edge_frames)
+    return hg
+
+DGLHeteroGraph.edge_type_subgraph = edge_type_subgraph
+
+#################### Internal functions ####################
+
+def _create_hetero_subgraph(parent, sgi, induced_nodes, induced_edges):
+    """Internal function to create a subgraph.
+
+    Parameters
+    ----------
+    parent : DGLGraph
+        The parent DGLGraph.
+    sgi : HeteroSubgraphIndex
+        Subgraph object returned by CAPI.
+    induced_nodes : list[Tensor] or None
+        Induced node IDs. Will store it as the dgl.NID ndata unless it
+        is None, which means the induced node IDs are the same as the parent node IDs.
+    induced_edges : list[Tensor] or None
+        Induced edge IDs. Will store it as the dgl.EID ndata unless it
+        is None, which means the induced edge IDs are the same as the parent edge IDs.
+
+    Returns
+    -------
+    DGLGraph
+        Graph
+    """
+    node_frames, edge_frames = utils.extract_subframes(parent, induced_nodes, induced_edges)
+    hsg = DGLHeteroGraph(sgi.graph, parent.ntypes, parent.etypes,
+                         node_frames, edge_frames)
+    return hsg
+
+_init_api("dgl.subgraph")

python/dgl/transform.py

@@ -43,8 +43,6 @@ __all__ = [
     'to_block',
     'to_simple',
     'to_simple_graph',
-    'in_subgraph',
-    'out_subgraph',
     'as_immutable_graph',
     'as_heterograph']
 
@@ -1908,90 +1906,6 @@ def to_block(g, dst_nodes=None, include_dst_in_src=True, copy_ndata=True, copy_e
 
     return new_graph
 
-def in_subgraph(g, nodes):
-    """Extract the subgraph containing only the in edges of the given nodes.
-
-    The subgraph keeps the same type schema and the cardinality of the original one.
-    Node/edge features are not preserved. The original IDs
-    the extracted edges are stored as the `dgl.EID` feature in the returned graph.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        Full graph structure.
-    nodes : tensor or dict
-        Node ids to sample neighbors from. The allowed types
-        are dictionary of node types to node id tensors, or simply node id tensor if
-        the given graph g has only one type of nodes.
-
-    Returns
-    -------
-    DGLHeteroGraph
-        The subgraph.
-    """
-    if g.is_block:
-        raise DGLError('Extracting subgraph of a block graph is not allowed.')
-    if not isinstance(nodes, dict):
-        if len(g.ntypes) > 1:
-            raise DGLError("Must specify node type when the graph is not homogeneous.")
-        nodes = {g.ntypes[0] : nodes}
-    nodes = utils.prepare_tensor_dict(g, nodes, 'nodes')
-    nodes_all_types = []
-    for ntype in g.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[g._idtype_str])
-
-    subgidx = _CAPI_DGLInSubgraph(g._graph, nodes_all_types)
-    induced_edges = subgidx.induced_edges
-    ret = DGLHeteroGraph(subgidx.graph, g.ntypes, g.etypes)
-    for i, etype in enumerate(ret.canonical_etypes):
-        ret.edges[etype].data[EID] = induced_edges[i]
-    return ret
-
-def out_subgraph(g, nodes):
-    """Extract the subgraph containing only the out edges of the given nodes.
-
-    The subgraph keeps the same type schema and the cardinality of the original one.
-    Node/edge features are not preserved. The original IDs
-    the extracted edges are stored as the `dgl.EID` feature in the returned graph.
-
-    Parameters
-    ----------
-    g : DGLHeteroGraph
-        Full graph structure.
-    nodes : tensor or dict
-        Node ids to sample neighbors from. The allowed types
-        are dictionary of node types to node id tensors, or simply node id tensor if
-        the given graph g has only one type of nodes.
-
-    Returns
-    -------
-    DGLHeteroGraph
-        The subgraph.
-    """
-    if g.is_block:
-        raise DGLError('Extracting subgraph of a block graph is not allowed.')
-    if not isinstance(nodes, dict):
-        if len(g.ntypes) > 1:
-            raise DGLError("Must specify node type when the graph is not homogeneous.")
-        nodes = {g.ntypes[0] : nodes}
-    nodes = utils.prepare_tensor_dict(g, nodes, 'nodes')
-    nodes_all_types = []
-    for ntype in g.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[g._idtype_str])
-
-    subgidx = _CAPI_DGLOutSubgraph(g._graph, nodes_all_types)
-    induced_edges = subgidx.induced_edges
-    ret = DGLHeteroGraph(subgidx.graph, g.ntypes, g.etypes)
-    for i, etype in enumerate(ret.canonical_etypes):
-        ret.edges[etype].data[EID] = induced_edges[i]
-    return ret
-
 def to_simple(g, return_counts='count', writeback_mapping=False, copy_ndata=True, copy_edata=False):
     r"""Convert a graph to a simple graph without duplicate edges.
 

python/dgl/udf.py

@@ -6,28 +6,26 @@ class EdgeBatch(object):
 
     Parameters
     ----------
-    edges : tuple of utils.Index
-        The edge tuple (u, v, eid). eid can be ALL
-    src_data : dict
-        The src node features, in the form of ``dict``
-        with ``str`` keys and ``tensor`` values
-    edge_data : dict
-        The edge features, in the form of ``dict`` with
-        ``str`` keys and ``tensor`` values
-    dst_data : dict of tensors
-        The dst node features, in the form of ``dict``
-        with ``str`` keys and ``tensor`` values
-    canonical_etype : tuple of (str, str, str), optional
-        Canonical edge type of the edge batch, if UDF is
-        running on a heterograph.
+    graph : DGLGraph
+        Graph object.
+    eid : Tensor
+        Edge IDs.
+    etype : (str, str, str)
+        Edge type.
+    src_data : dict[str, Tensor]
+        Src node features.
+    edge_data : dict[str, Tensor]
+        Edge features.
+    dst_data : dict[str, Tensor]
+        Dst node features.
     """
-    def __init__(self, edges, src_data, edge_data, dst_data,
-                 canonical_etype=(None, None, None)):
-        self._edges = edges
+    def __init__(self, graph, eid, etype, src_data, edge_data, dst_data):
+        self._graph = graph
+        self._eid = eid
+        self._etype = etype
         self._src_data = src_data
         self._edge_data = edge_data
         self._dst_data = dst_data
-        self._canonical_etype = canonical_etype
 
     @property
     def src(self):
@@ -67,14 +65,15 @@ class EdgeBatch(object):
 
         Returns
         -------
-        tuple of three tensors
-            The edge tuple :math:`(src, dst, eid)`. :math:`src[i],
-            dst[i], eid[i]` separately specifies the source node,
-            destination node and the edge id for the ith edge
-            in the batch.
+        Tensor
+            Source node IDs.
+        Tensor
+            Destination node IDs.
+        Tensor
+            Edge IDs.
         """
-        u, v, eid = self._edges
-        return (u.tousertensor(), v.tousertensor(), eid.tousertensor())
+        u, v = self._graph.find_edges(self._eid, etype=self.canonical_etype)
+        return u, v, self._eid
 
     def batch_size(self):
         """Return the number of edges in this edge batch.
@@ -83,7 +82,7 @@ class EdgeBatch(object):
         -------
         int
         """
-        return len(self._edges[0])
+        return len(self._eid)
 
     def __len__(self):
         """Return the number of edges in this edge batch.
@@ -97,31 +96,31 @@ class EdgeBatch(object):
     @property
     def canonical_etype(self):
         """Return the canonical edge type (i.e. triplet of source, edge, and
-        destination node type) for this edge batch, if available."""
-        return self._canonical_etype
+        destination node type) for this edge batch."""
+        return self._etype
 
 class NodeBatch(object):
-    """The class that can represent a batch of nodes.
+    """The class to represent a batch of nodes.
 
     Parameters
     ----------
-    nodes : utils.Index
-        The node ids.
-    data : dict
-        The node features, in the form of ``dict``
-        with ``str`` keys and ``tensor`` values
-    msgs : dict, optional
-        The messages, , in the form of ``dict``
-        with ``str`` keys and ``tensor`` values
+    graph : DGLGraph
+        Graph object.
+    nodes : Tensor
+        Node ids.
     ntype : str, optional
-        The node type of this node batch, if running
-        on a heterograph.
+        The node type of this node batch,
+    data : dict[str, Tensor]
+        Node feature data.
+    msgs : dict[str, Tensor], optional
+        Messages data.
     """
-    def __init__(self, nodes, data, msgs=None, ntype=None):
+    def __init__(self, graph, nodes, ntype, data, msgs=None):
+        self._graph = graph
         self._nodes = nodes
+        self._ntype = ntype
         self._data = data
         self._msgs = msgs
-        self._ntype = ntype
 
     @property
     def data(self):
@@ -156,7 +155,7 @@ class NodeBatch(object):
         tensor
             The nodes.
         """
-        return self._nodes.tousertensor()
+        return self._nodes
 
     def batch_size(self):
         """Return the number of nodes in this batch.

python/dgl/utils/checks.py

@@ -63,6 +63,36 @@ def prepare_tensor_dict(g, data, name):
     return {key : prepare_tensor(g, val, '{}["{}"]'.format(name, key))
             for key, val in data.items()}
 
+def parse_edges_arg_to_eid(g, edges, etid, argname='edges'):
+    """Parse the :attr:`edges` argument and return an edge ID tensor.
+
+    The resulting edge ID tensor has the same ID type and device of :attr:`g`.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        Graph
+    edges : pair of Tensor, Tensor, iterable[int]
+        Argument for specifying edges.
+    etid : int
+        Edge type ID.
+    argname : str, optional
+        Argument name.
+
+    Returns
+    -------
+    Tensor
+        Edge ID tensor
+    """
+    if isinstance(edges, tuple):
+        u, v = edges
+        u = prepare_tensor(g, u, '{}[0]'.format(argname))
+        v = prepare_tensor(g, v, '{}[1]'.format(argname))
+        eid = g.edge_ids(u, v, etype=g.canonical_etypes[etid])
+    else:
+        eid = prepare_tensor(g, edges, argname)
+    return eid
+
 def check_all_same_idtype(glist, name):
     """Check all the graphs have the same idtype."""
     if len(glist) == 0:

python/dgl/utils/internal.py

@@ -6,11 +6,10 @@ from collections import defaultdict
 from functools import wraps
 import numpy as np
 
-from ..base import DGLError, dgl_warning
+from ..base import DGLError, dgl_warning, NID, EID
 from .. import backend as F
 from .. import ndarray as nd
 
-
 class InconsistentDtypeException(DGLError):
     """Exception class for inconsistent dtype between graph and tensor"""
     def __init__(self, msg='', *args, **kwargs): #pylint: disable=W1113
@@ -708,3 +707,89 @@ def compensate(ids, origin_ids):
                          ids,
                          F.full_1d(len(ids), 0, F.dtype(ids), F.context(ids)))
     return F.tensor(F.nonzero_1d(mask), dtype=F.dtype(ids))
+
+def relabel(x):
+    """Relabel the input ids to continuous ids that starts from zero.
+
+    Ids are assigned new ids according to their ascending order.
+
+    Examples
+    --------
+    >>> x = [1, 5, 3, 6]
+    >>> n2o, o2n = build_relabel_map(x)
+    >>> n2o
+    [1, 3, 5, 6]
+    >>> o2n
+    [n/a, 0, n/a, 1, n/a, 2, 3]
+
+    "n/a" will be filled with 0
+
+    Parameters
+    ----------
+    x : Tensor
+        ID tensor.
+
+    Returns
+    -------
+    new_to_old : Tensor
+        The mapping from new id to old id.
+    old_to_new : Tensor
+        The mapping from old id to new id. It is a vector of length MAX(x).
+        One can use advanced indexing to convert an old id tensor to a
+        new id tensor: new_id = old_to_new[old_id]
+    """
+    unique_x = F.unique(x)
+    map_len = F.as_scalar(F.max(unique_x, dim=0)) + 1
+    ctx = F.context(x)
+    dtype = F.dtype(x)
+    old_to_new = F.zeros((map_len,), dtype=dtype, ctx=ctx)
+    old_to_new = F.scatter_row(old_to_new, unique_x,
+                               F.copy_to(F.arange(0, len(unique_x), dtype), ctx))
+    return unique_x, old_to_new
+
+def extract_subframes(graph, nodes, edges):
+    """Extract node/edge features of the given nodes and edges from :attr:`graph`
+    and return them in frames.
+
+    Note that this function does not perform actual tensor memory copy but using `Frame.subframe`
+    to get the features. If :attr:`nodes` is None, it performs a shallow copy of the
+    original node frames that only copies the dictionary structure but not the tensor
+    contents.
+
+    Parameters
+    ----------
+    graph : DGLGraph
+        The graph to extract features from.
+    nodes : list[Tensor] or None
+        Node IDs. If not None, the list length must be equal to the number of node types
+        in the graph. The returned frames store the node IDs in the ``dgl.NID`` field
+        unless it is None, which means the whole frame is shallow-copied.
+    edges : list[Tensor] or None
+        Edge IDs. If not None, the list length must be equal to the number of edge types
+        in the graph. The returned frames store the edge IDs in the ``dgl.NID`` field
+        unless it is None, which means the whole frame is shallow-copied.
+
+    Returns
+    -------
+    list[Frame]
+        Extracted node frames.
+    list[Frame]
+        Extracted edge frames.
+    """
+    if nodes is None:
+        node_frames = [nf.clone() for nf in graph._node_frames]
+    else:
+        node_frames = []
+        for i, ind_nodes in enumerate(nodes):
+            subf = graph._node_frames[i].subframe(ind_nodes)
+            subf[NID] = ind_nodes
+            node_frames.append(subf)
+    if edges is None:
+        edge_frames = [nf.clone() for nf in graph._edge_frames]
+    else:
+        edge_frames = []
+        for i, ind_edges in enumerate(edges):
+            subf = graph._edge_frames[i].subframe(ind_edges)
+            subf[EID] = ind_edges
+            edge_frames.append(subf)
+    return node_frames, edge_frames

src/graph/heterograph_capi.cc

@@ -648,7 +648,7 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroGetFormatGraph")
     *rv = HeteroGraphRef(hgptr);
 });
 
-DGL_REGISTER_GLOBAL("transform._CAPI_DGLInSubgraph")
+DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLInSubgraph")
 .set_body([] (DGLArgs args, DGLRetValue *rv) {
     HeteroGraphRef hg = args[0];
     const auto& nodes = ListValueToVector<IdArray>(args[1]);
@@ -657,7 +657,7 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLInSubgraph")
     *rv = HeteroGraphRef(ret);
   });
 
-DGL_REGISTER_GLOBAL("transform._CAPI_DGLOutSubgraph")
+DGL_REGISTER_GLOBAL("subgraph._CAPI_DGLOutSubgraph")
 .set_body([] (DGLArgs args, DGLRetValue *rv) {
     HeteroGraphRef hg = args[0];
     const auto& nodes = ListValueToVector<IdArray>(args[1]);

tests/compute/test_basics.py

@@ -6,6 +6,7 @@ import networkx as nx
 from dgl import DGLGraph
 from collections import defaultdict as ddict
 import unittest
+from test_utils import parametrize_dtype
 
 D = 5
 reduce_msg_shapes = set()
@@ -25,7 +26,7 @@ def reduce_func(nodes):
 def apply_node_func(nodes):
     return {'h' : nodes.data['h'] + nodes.data['accum']}
 
-def generate_graph(grad=False):
+def generate_graph_old(grad=False):
     g = DGLGraph()
     g.add_nodes(10) # 10 nodes
     # create a graph where 0 is the source and 9 is the sink
@@ -48,10 +49,51 @@ def generate_graph(grad=False):
     g.set_e_initializer(dgl.init.zero_initializer)
     return g
 
-def test_batch_setter_getter():
+def generate_graph(idtype, grad=False):
+    '''
+    s, d, eid
+    0, 1, 0
+    1, 9, 1
+    0, 2, 2
+    2, 9, 3
+    0, 3, 4
+    3, 9, 5
+    0, 4, 6
+    4, 9, 7
+    0, 5, 8
+    5, 9, 9
+    0, 6, 10
+    6, 9, 11
+    0, 7, 12
+    7, 9, 13
+    0, 8, 14
+    8, 9, 15
+    9, 0, 16
+    '''
+    u = F.tensor([0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 9])
+    v = F.tensor([1, 9, 2, 9, 3, 9, 4, 9, 5, 9, 6, 9, 7, 9, 8, 9, 0])
+    g = dgl.graph((u, v), idtype=idtype)
+    assert g.device == F.ctx()
+    ncol = F.randn((10, D))
+    ecol = F.randn((17, D))
+    if grad:
+        ncol = F.attach_grad(ncol)
+        ecol = F.attach_grad(ecol)
+
+    g.ndata['h'] = ncol
+    g.edata['w'] = ecol
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
+    return g
+
+def test_compatible():
+    g = generate_graph_old()
+
+@parametrize_dtype
+def test_batch_setter_getter(idtype):
     def _pfc(x):
         return list(F.zerocopy_to_numpy(x)[:,0])
-    g = generate_graph()
+    g = generate_graph(idtype)
     # set all nodes
     g.ndata['h'] = F.zeros((10, D))
     assert F.allclose(g.ndata['h'], F.zeros((10, D)))
@@ -128,9 +170,9 @@ def test_batch_setter_getter():
     v = F.tensor([3, 4, 5], g.idtype)
     assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
 
-
-def test_batch_setter_autograd():
-    g = generate_graph(grad=True)
+@parametrize_dtype
+def test_batch_setter_autograd(idtype):
+    g = generate_graph(idtype, grad=True)
     h1 = g.ndata['h']
     # partial set
     v = F.tensor([1, 2, 8], g.idtype)
@@ -261,10 +303,11 @@ def _test_nx_conversion():
     assert F.allclose(g.edata['h'], F.tensor([[1., 2.], [1., 2.],
                                               [2., 3.], [2., 3.]]))
 
-def test_apply_nodes():
+@parametrize_dtype
+def test_apply_nodes(idtype):
     def _upd(nodes):
         return {'h' : nodes.data['h'] * 2}
-    g = generate_graph()
+    g = generate_graph(idtype)
     old = g.ndata['h']
     g.apply_nodes(_upd)
     assert F.allclose(old * 2, g.ndata['h'])
@@ -272,10 +315,11 @@ def test_apply_nodes():
     g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
     assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
 
-def test_apply_edges():
+@parametrize_dtype
+def test_apply_edges(idtype):
     def _upd(edges):
         return {'w' : edges.data['w'] * 2}
-    g = generate_graph()
+    g = generate_graph(idtype)
     old = g.edata['w']
     g.apply_edges(_upd)
     assert F.allclose(old * 2, g.edata['w'])
@@ -285,8 +329,9 @@ def test_apply_edges():
     eid = F.tensor(g.edge_ids(u, v))
     assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
 
-def test_update_routines():
-    g = generate_graph()
+@parametrize_dtype
+def test_update_routines(idtype):
+    g = generate_graph(idtype)
 
     # send_and_recv
     reduce_msg_shapes.clear()
@@ -321,28 +366,23 @@ def test_update_routines():
     assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
     reduce_msg_shapes.clear()
 
-def test_update_all_0deg():
+@parametrize_dtype
+def test_update_all_0deg(idtype):
     # test#1
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(5)
-    g.add_edge(1, 0)
-    g.add_edge(2, 0)
-    g.add_edge(3, 0)
-    g.add_edge(4, 0)
+    g = dgl.graph([(1,0), (2,0), (3,0), (4,0)], idtype=idtype, device=F.ctx())
     def _message(edges):
         return {'m' : edges.src['h']}
     def _reduce(nodes):
-        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
+        return {'x' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
     def _apply(nodes):
-        return {'h' : nodes.data['h'] * 2}
+        return {'x' : nodes.data['x'] * 2}
     def _init2(shape, dtype, ctx, ids):
         return 2 + F.zeros(shape, dtype, ctx)
-    g.set_n_initializer(_init2, 'h')
+    g.set_n_initializer(_init2, 'x')
     old_repr = F.randn((5, 5))
     g.ndata['h'] = old_repr
     g.update_all(_message, _reduce, _apply)
-    new_repr = g.ndata['h']
+    new_repr = g.ndata['x']
     # the first row of the new_repr should be the sum of all the node
     # features; while the 0-deg nodes should be initialized by the
     # initializer and applied with UDF.
@@ -350,35 +390,30 @@ def test_update_all_0deg():
     assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
 
     # test#2: graph with no edge
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(5)
-    g.set_n_initializer(_init2, 'h')
+    g = dgl.graph([], num_nodes=5, idtype=idtype, device=F.ctx())
     g.ndata['h'] = old_repr
-    g.update_all(_message, _reduce, _apply)
+    g.update_all(_message, _reduce, lambda nodes : {'h' : nodes.data['h'] * 2})
     new_repr = g.ndata['h']
     # should fallback to apply
     assert F.allclose(new_repr, 2*old_repr)
 
-def test_pull_0deg():
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(2)
-    g.add_edge(0, 1)
+@parametrize_dtype
+def test_pull_0deg(idtype):
+    g = dgl.graph([(0,1)], idtype=idtype, device=F.ctx())
     def _message(edges):
         return {'m' : edges.src['h']}
     def _reduce(nodes):
-        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
+        return {'x' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
     def _apply(nodes):
-        return {'h' : nodes.data['h'] * 2}
+        return {'x' : nodes.data['x'] * 2}
     def _init2(shape, dtype, ctx, ids):
         return 2 + F.zeros(shape, dtype, ctx)
-    g.set_n_initializer(_init2, 'h')
+    g.set_n_initializer(_init2, 'x')
     # test#1: pull both 0deg and non-0deg nodes
     old = F.randn((2, 5))
     g.ndata['h'] = old
     g.pull([0, 1], _message, _reduce, _apply)
-    new = g.ndata.pop('h')
+    new = g.ndata['x']
     # 0deg check: initialized with the func and got applied
     assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
     # non-0deg check
@@ -387,8 +422,8 @@ def test_pull_0deg():
     # test#2: pull only 0deg node
     old = F.randn((2, 5))
     g.ndata['h'] = old
-    g.pull(0, _message, _reduce, _apply)
-    new = g.ndata.pop('h')
+    g.pull(0, _message, _reduce, lambda nodes : {'h' : nodes.data['h'] * 2})
+    new = g.ndata['h']
     # 0deg check: fallback to apply
     assert F.allclose(new[0], 2*old[0])
     # non-0deg check: not touched
@@ -428,81 +463,19 @@ def test_dynamic_addition():
     assert len(g.edata['h1']) == len(g.edata['h2'])
 
 
-def test_repr():
-    g = dgl.DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(10)
-    g.add_edge(0, 1)
+@parametrize_dtype
+def test_repr(idtype):
+    g = dgl.graph([(0,1), (0,2), (1,2)], num_nodes=10, idtype=idtype, device=F.ctx())
     repr_string = g.__repr__()
     print(repr_string)
     g.ndata['x'] = F.zeros((10, 5))
-    g.add_edges([0, 1], 2)
     g.edata['y'] = F.zeros((3, 4))
     repr_string = g.__repr__()
     print(repr_string)
 
-
-def test_group_apply_edges():
-    def edge_udf(edges):
-        h = F.sum(edges.data['feat'] * (edges.src['h'] + edges.dst['h']), dim=2)
-        normalized_feat = F.softmax(h, dim=1)
-        return {"norm_feat": normalized_feat}
-
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(10)
-    g.add_edges(0, [1, 2, 3, 4, 5, 6, 7, 8])
-    g.add_edges(1, [2, 3, 4, 6, 7, 8])
-    g.add_edges(2, [2, 3, 4, 5, 6, 7, 8])
-
-    g.ndata['h'] = F.randn((g.number_of_nodes(), D))
-    g.edata['feat'] = F.randn((g.number_of_edges(), D))
-
-    def _test(group_by):
-        g.group_apply_edges(group_by=group_by, func=edge_udf)
-        if group_by == 'src':
-            u, v, eid = g.out_edges(1, form='all')
-        else:
-            u, v, eid = g.in_edges(5, form='all')
-        out_feat = g.edges[eid].data['norm_feat']
-        result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
-        result = F.softmax(F.sum(result, dim=1), dim=0)
-        assert F.allclose(out_feat, result)
-
-    # test group by source nodes
-    _test('src')
-
-    # test group by destination nodes
-    _test('dst')
-
-
-# GitHub issue #1036
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_group_apply_edges2():
-    m = ssp.random(10, 10, 0.2)
-    g = DGLGraph(m, readonly=True)
-    g = g.to(F.ctx())
-    g.ndata['deg'] = g.in_degrees()
-    g.ndata['id'] = F.arange(0, g.number_of_nodes(), g.idtype)
-    g.edata['id'] = F.arange(0, g.number_of_edges(), g.idtype)
-
-    def apply(edges):
-        w = edges.data['id']
-        n_nodes, deg = w.shape
-
-        dst = edges.dst['id'][:, 0]
-        # TODO: tmp hack
-        eid1 = F.asnumpy(g.in_edges(dst, 'eid')).reshape(n_nodes, deg).sort(1)
-        eid2 = F.asnumpy(edges.data['id']).sort(1)
-        assert np.array_equal(eid1, eid2)
-        return {'id2': w}
-
-    g.group_apply_edges('dst', apply, inplace=True)
-
-
-def test_local_var():
-    g = DGLGraph(nx.path_graph(5))
-    g = g.to(F.ctx())
+@parametrize_dtype
+def test_local_var(idtype):
+    g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
     g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
     g.edata['w'] = F.zeros((g.number_of_edges(), 4))
     # test override
@@ -539,10 +512,7 @@ def test_local_var():
     assert 'ww' not in g.edata
 
     # test initializer1
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(2)
-    g.add_edges([0, 1], [1, 1])
+    g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
     g.set_n_initializer(dgl.init.zero_initializer)
     def foo(g):
         g = g.local_var()
@@ -561,9 +531,9 @@ def test_local_var():
         assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
     foo(g)
 
-def test_local_scope():
-    g = DGLGraph(nx.path_graph(5))
-    g = g.to(F.ctx())
+@parametrize_dtype
+def test_local_scope(idtype):
+    g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
     g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
     g.edata['w'] = F.zeros((g.number_of_edges(), 4))
     # test override
@@ -614,10 +584,7 @@ def test_local_scope():
     assert 'ww' not in g.edata
 
     # test initializer1
-    g = DGLGraph()
-    g = g.to(F.ctx())
-    g.add_nodes(2)
-    g.add_edges([0, 1], [1, 1])
+    g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
     g.set_n_initializer(dgl.init.zero_initializer)
     def foo(g):
         with g.local_scope():
@@ -636,23 +603,59 @@ def test_local_scope():
             assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
     foo(g)
 
-if __name__ == '__main__':
-    #test_nx_conversion()
-    test_batch_setter_getter()
-    test_batch_setter_autograd()
-    test_batch_send()
-    test_batch_recv()
-    test_apply_nodes()
-    test_apply_edges()
-    test_update_routines()
-    test_recv_0deg()
-    test_recv_0deg_newfld()
-    test_update_all_0deg()
-    test_pull_0deg()
-    test_send_multigraph()
-    test_dynamic_addition()
-    test_repr()
-    test_group_apply_edges()
-    test_group_apply_edges2()
-    test_local_var()
-    test_local_scope()
+@parametrize_dtype
+def test_isolated_nodes(idtype):
+    g = dgl.graph([(0, 1), (1, 2)], num_nodes=5, idtype=idtype, device=F.ctx())
+    assert g.number_of_nodes() == 5
+
+    # Test backward compatibility
+    g = dgl.graph([(0, 1), (1, 2)], card=5, idtype=idtype, device=F.ctx())
+    assert g.number_of_nodes() == 5
+
+    g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
+                      'game', num_nodes=(5, 7), idtype=idtype, device=F.ctx())
+    assert g.idtype == idtype
+    assert g.number_of_nodes('user') == 5
+    assert g.number_of_nodes('game') == 7
+
+    # Test backward compatibility
+    g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
+                      'game', card=(5, 7), idtype=idtype, device=F.ctx())
+    assert g.idtype == idtype
+    assert g.number_of_nodes('user') == 5
+    assert g.number_of_nodes('game') == 7
+
+@parametrize_dtype
+def test_send_multigraph(idtype):
+    g = dgl.graph([(0,1), (0,1), (0,1), (2,1)], idtype=idtype, device=F.ctx())
+
+    def _message_a(edges):
+        return {'a': edges.data['a']}
+    def _message_b(edges):
+        return {'a': edges.data['a'] * 3}
+    def _reduce(nodes):
+        return {'a': F.max(nodes.mailbox['a'], 1)}
+
+    def answer(*args):
+        return F.max(F.stack(args, 0), 0)
+
+    assert g.is_multigraph
+
+    # send by eid
+    old_repr = F.randn((4, 5))
+    # send_and_recv_on
+    g.ndata['a'] = F.zeros((3, 5))
+    g.edata['a'] = old_repr
+    g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
+    new_repr = g.ndata['a']
+    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
+    assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
+
+@parametrize_dtype
+def test_issue_1088(idtype):
+    # This test ensures that message passing on a heterograph with one edge type
+    # would not crash (GitHub issue #1088).
+    import dgl.function as fn
+    g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])}, idtype=idtype, device=F.ctx())
+    g.nodes['U'].data['x'] = F.randn((3, 3))
+    g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))

tests/compute/test_frame.py.bak

-import numpy as np
-from dgl.frame import Frame, FrameRef
-from dgl.utils import Index, toindex
-import backend as F
-import dgl
-import unittest
-import pickle
-import pytest
-import io
-
-N = 10
-D = 5
-
-def check_fail(fn):
-    try:
-        fn()
-        return False
-    except:
-        return True
-
-def create_test_data(grad=False, dtype=F.float32):
-    c1 = F.astype(F.randn((N, D)), dtype)
-    c2 = F.astype(F.randn((N, D)), dtype)
-    c3 = F.astype(F.randn((N, D)), dtype)
-    if grad:
-        c1 = F.attach_grad(c1)
-        c2 = F.attach_grad(c2)
-        c3 = F.attach_grad(c3)
-    return {'a1' : c1, 'a2' : c2, 'a3' : c3}
-
-def test_create():
-    data = create_test_data()
-    f1 = Frame(num_rows=N)
-    for k, v in data.items():
-        f1.update_column(k, v)
-    print(f1.schemes)
-    assert f1.keys() == set(data.keys())
-    assert f1.num_columns == 3
-    assert f1.num_rows == N
-    f2 = Frame(data)
-    assert f2.keys() == set(data.keys())
-    assert f2.num_columns == 3
-    assert f2.num_rows == N
-    f1.clear()
-    assert len(f1.schemes) == 0
-    assert f1.num_rows == 0
-
-def test_column1():
-    # Test frame column getter/setter
-    data = create_test_data()
-    f = Frame(data)
-    assert f.num_rows == N
-    assert len(f) == 3
-    assert F.allclose(f['a1'].data, data['a1'])
-    f['a1'] = data['a2']
-    assert F.allclose(f['a2'].data, data['a2'])
-    # add a different length column should fail
-    def failed_add_col():
-        f['a4'] = F.zeros([N+1, D])
-    assert check_fail(failed_add_col)
-    # delete all the columns
-    del f['a1']
-    del f['a2']
-    assert len(f) == 1
-    del f['a3']
-    assert len(f) == 0
-
-def test_column2():
-    # Test frameref column getter/setter
-    data = Frame(create_test_data())
-    f = FrameRef(data, toindex([3, 4, 5, 6, 7]))
-    assert f.num_rows == 5
-    assert len(f) == 3
-    assert F.allclose(f['a1'], F.narrow_row(data['a1'].data, 3, 8))
-    # set column should reflect on the referenced data
-    f['a1'] = F.zeros([5, D])
-    assert F.allclose(F.narrow_row(data['a1'].data, 3, 8), F.zeros([5, D]))
-    # add new partial column should fail with error initializer
-    f.set_initializer(lambda shape, dtype : assert_(False))
-    def failed_add_col():
-        f['a4'] = F.ones([5, D])
-    assert check_fail(failed_add_col)
-
-def test_append1():
-    # test append API on Frame
-    data = create_test_data()
-    f1 = Frame()
-    f2 = Frame(data)
-    f1.append(data)
-    assert f1.num_rows == N
-    f1.append(f2)
-    assert f1.num_rows == 2 * N
-    c1 = f1['a1']
-    assert tuple(F.shape(c1.data)) == (2 * N, D)
-    truth = F.cat([data['a1'], data['a1']], 0)
-    assert F.allclose(truth, c1.data)
-    # append dict of different length columns should fail
-    f3 = {'a1' : F.zeros((3, D)), 'a2' : F.zeros((3, D)), 'a3' : F.zeros((2, D))}
-    def failed_append():
-        f1.append(f3)
-    assert check_fail(failed_append)
-
-def test_append2():
-    # test append on FrameRef
-    data = Frame(create_test_data())
-    f = FrameRef(data)
-    assert f.is_contiguous()
-    assert f.is_span_whole_column()
-    assert f.num_rows == N
-    # append on the underlying frame should not reflect on the ref
-    data.append(data)
-    assert f.is_contiguous()
-    assert not f.is_span_whole_column()
-    assert f.num_rows == N
-    # append on the FrameRef should work
-    f.append(data)
-    assert not f.is_contiguous()
-    assert not f.is_span_whole_column()
-    assert f.num_rows == 3 * N
-    new_idx = list(range(N)) + list(range(2*N, 4*N))
-    assert F.array_equal(f._index.tousertensor(), F.copy_to(F.tensor(new_idx, dtype=F.int64), F.cpu()))
-    assert data.num_rows == 4 * N
-
-def test_append3():
-    # test append on empty frame
-    f = Frame(num_rows=5)
-    data = {'h' : F.ones((3, 2))}
-    f.append(data)
-    assert f.num_rows == 8
-    ans = F.cat([F.zeros((5, 2)), F.ones((3, 2))], 0)
-    assert F.allclose(f['h'].data, ans)
-    # test append with new column
-    data = {'h' : 2 * F.ones((3, 2)), 'w' : 2 * F.ones((3, 2))}
-    f.append(data)
-    assert f.num_rows == 11
-    ans1 = F.cat([ans, 2 * F.ones((3, 2))], 0)
-    ans2 = F.cat([F.zeros((8, 2)), 2 * F.ones((3, 2))], 0)
-    assert F.allclose(f['h'].data, ans1)
-    assert F.allclose(f['w'].data, ans2)
-
-def test_row1():
-    # test row getter/setter
-    data = create_test_data()
-    f = FrameRef(Frame(data))
-
-    # getter
-    # test non-duplicate keys
-    rowid = Index(F.tensor([0, 2]))
-    rows = f[rowid]
-    for k, v in rows.items():
-        assert tuple(F.shape(v)) == (len(rowid), D)
-        assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
-    # test duplicate keys
-    rowid = Index(F.tensor([8, 2, 2, 1]))
-    rows = f[rowid]
-    for k, v in rows.items():
-        assert tuple(F.shape(v)) == (len(rowid), D)
-        assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
-
-    # setter
-    rowid = Index(F.tensor([0, 2, 4]))
-    vals = {'a1' : F.zeros((len(rowid), D)),
-            'a2' : F.zeros((len(rowid), D)),
-            'a3' : F.zeros((len(rowid), D)),
-            }
-    f[rowid] = vals
-    for k, v in f[rowid].items():
-        assert F.allclose(v, F.zeros((len(rowid), D)))
-
-    # setting rows with new column should raise error with error initializer
-    f.set_initializer(lambda shape, dtype : assert_(False))
-    def failed_update_rows():
-        vals['a4'] = F.ones((len(rowid), D))
-        f[rowid] = vals
-    assert check_fail(failed_update_rows)
-
-def test_row2():
-    # test row getter/setter autograd compatibility
-    data = create_test_data(grad=True)
-    f = FrameRef(Frame(data))
-
-    with F.record_grad():
-        # getter
-        c1 = f['a1']
-        # test non-duplicate keys
-        rowid = Index(F.tensor([0, 2]))
-        rows = f[rowid]
-        y = rows['a1']
-        F.backward(y, F.ones((len(rowid), D)))
-    assert F.allclose(F.grad(c1)[:,0], F.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
-
-    f['a1'] = F.attach_grad(f['a1'])
-    with F.record_grad():
-        c1 = f['a1']
-        # test duplicate keys
-        rowid = Index(F.tensor([8, 2, 2, 1]))
-        rows = f[rowid]
-        y = rows['a1']
-        F.backward(y, F.ones((len(rowid), D)))
-    assert F.allclose(F.grad(c1)[:,0], F.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
-
-    f['a1'] = F.attach_grad(f['a1'])
-    with F.record_grad():
-        # setter
-        c1 = f['a1']
-        rowid = Index(F.tensor([0, 2, 4]))
-        vals = {'a1' : F.attach_grad(F.zeros((len(rowid), D))),
-                'a2' : F.attach_grad(F.zeros((len(rowid), D))),
-                'a3' : F.attach_grad(F.zeros((len(rowid), D))),
-                }
-        f[rowid] = vals
-        c11 = f['a1']
-        F.backward(c11, F.ones((N, D)))
-    assert F.allclose(F.grad(c1)[:,0], F.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
-    assert F.allclose(F.grad(vals['a1']), F.ones((len(rowid), D)))
-    assert F.is_no_grad(vals['a2'])
-
-def test_row3():
-    # test row delete
-    data = Frame(create_test_data())
-    f = FrameRef(data)
-    assert f.is_contiguous()
-    assert f.is_span_whole_column()
-    assert f.num_rows == N
-    del f[toindex(F.tensor([2, 3]))]
-    assert not f.is_contiguous()
-    assert not f.is_span_whole_column()
-    # delete is lazy: only reflect on the ref while the
-    # underlying storage should not be touched
-    assert f.num_rows == N - 2
-    assert data.num_rows == N
-    newidx = list(range(N))
-    newidx.pop(2)
-    newidx.pop(2)
-    newidx = toindex(newidx)
-    for k, v in f.items():
-        assert F.allclose(v, data[k][newidx])
-
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_row4():
-    # test updating row with empty frame but has preset num_rows
-    f = FrameRef(Frame(num_rows=5))
-    rowid = Index(F.tensor([0, 2, 4]))
-    f[rowid] = {'h' : F.ones((3, 2))}
-    ans = F.zeros((5, 2))
-    ans[F.tensor([0, 2, 4])] = F.ones((3, 2))
-    assert F.allclose(f['h'], ans)
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_sharing():
-    data = Frame(create_test_data())
-    f1 = FrameRef(data, index=toindex([0, 1, 2, 3]))
-    f2 = FrameRef(data, index=toindex([2, 3, 4, 5, 6]))
-    # test read
-    for k, v in f1.items():
-        assert F.allclose(F.narrow_row(data[k].data, 0, 4), v)
-    for k, v in f2.items():
-        assert F.allclose(F.narrow_row(data[k].data, 2, 7), v)
-    f2_a1 = f2['a1']
-    # test write
-    # update own ref should not been seen by the other.
-    f1[Index(F.tensor([0, 1]))] = {
-            'a1' : F.zeros([2, D]),
-            'a2' : F.zeros([2, D]),
-            'a3' : F.zeros([2, D]),
-            }
-    assert F.allclose(f2['a1'], f2_a1)
-    # update shared space should been seen by the other.
-    f1[Index(F.tensor([2, 3]))] = {
-            'a1' : F.ones([2, D]),
-            'a2' : F.ones([2, D]),
-            'a3' : F.ones([2, D]),
-            }
-    F.narrow_row_set(f2_a1, 0, 2, F.ones([2, D]))
-    assert F.allclose(f2['a1'], f2_a1)
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_slicing():
-    data = Frame(create_test_data(grad=True))
-    f1 = FrameRef(data, index=toindex(slice(1, 5)))
-    f2 = FrameRef(data, index=toindex(slice(3, 8)))
-    # test read
-    for k, v in f1.items():
-        assert F.allclose(F.narrow_row(data[k].data, 1, 5), v)
-    f2_a1 = f2['a1']    # is a tensor
-    # test write
-    f1[Index(F.tensor([0, 1]))] = {
-            'a1': F.zeros([2, D]),
-            'a2': F.zeros([2, D]),
-            'a3': F.zeros([2, D]),
-            }
-    assert F.allclose(f2['a1'], f2_a1)
-
-    f1[Index(F.tensor([2, 3]))] = {
-            'a1': F.ones([2, D]),
-            'a2': F.ones([2, D]),
-            'a3': F.ones([2, D]),
-            }
-    F.narrow_row_set(f2_a1, 0, 2, 1)
-    assert F.allclose(f2['a1'], f2_a1)
-
-    f1[toindex(slice(2, 4))] = {
-            'a1': F.zeros([2, D]),
-            'a2': F.zeros([2, D]),
-            'a3': F.zeros([2, D]),
-            }
-    F.narrow_row_set(f2_a1, 0, 2, 0)
-    assert F.allclose(f2['a1'], f2_a1)
-
-def test_add_rows():
-    data = Frame()
-    f1 = FrameRef(data)
-    f1.add_rows(4)
-    x = F.randn((1, 4))
-    f1[Index(F.tensor([0]))] = {'x': x}
-    ans = F.cat([x, F.zeros((3, 4))], 0)
-    assert F.allclose(f1['x'], ans)
-    f1.add_rows(4)
-    f1[toindex(slice(4, 8))] = {'x': F.ones((4, 4)), 'y': F.ones((4, 5))}
-    ans = F.cat([ans, F.ones((4, 4))], 0)
-    assert F.allclose(f1['x'], ans)
-    ans = F.cat([F.zeros((4, 5)), F.ones((4, 5))], 0)
-    assert F.allclose(f1['y'], ans)
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_inplace():
-    f = FrameRef(Frame(create_test_data()))
-    print(f.schemes)
-    a1addr = id(f['a1'])
-    a2addr = id(f['a2'])
-    a3addr = id(f['a3'])
-
-    # column updates are always out-of-place
-    f['a1'] = F.ones((N, D))
-    newa1addr = id(f['a1'])
-    assert a1addr != newa1addr
-    a1addr = newa1addr
-    # full row update that becomes column update
-    f[toindex(slice(0, N))] = {'a1' : F.ones((N, D))}
-    assert id(f['a1']) != a1addr
-
-    # row update (outplace) w/ slice
-    f[toindex(slice(1, 4))] = {'a2' : F.ones((3, D))}
-    newa2addr = id(f['a2'])
-    assert a2addr != newa2addr
-    a2addr = newa2addr
-    # row update (outplace) w/ list
-    f[toindex([1, 3, 5])] = {'a2' : F.ones((3, D))}
-    newa2addr = id(f['a2'])
-    assert a2addr != newa2addr
-    a2addr = newa2addr
-
-    # row update (inplace) w/ slice
-    f.update_data(toindex(slice(1, 4)), {'a2' : F.ones((3, D))}, True)
-    newa2addr = id(f['a2'])
-    assert a2addr == newa2addr
-    # row update (inplace) w/ list
-    f.update_data(toindex([1, 3, 5]), {'a2' : F.ones((3, D))}, True)
-    newa2addr = id(f['a2'])
-    assert a2addr == newa2addr
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_clone():
-    f = FrameRef(Frame(create_test_data()))
-    f1 = f.clone()
-    f2 = f.deepclone()
-
-    f1['b'] = F.randn((N, D))
-    f2['c'] = F.randn((N, D))
-    assert 'b' not in f
-    assert 'c' not in f
-
-    f1['a1'][0, 0] = -10.
-    assert float(F.asnumpy(f['a1'][0, 0])) == -10.
-    x = float(F.asnumpy(f['a2'][0, 0]))
-    f2['a2'][0, 0] = -10.
-    assert float(F.asnumpy(f['a2'][0, 0])) == x
-
-def _reconstruct_pickle(obj):
-    f = io.BytesIO()
-    pickle.dump(obj, f)
-    f.seek(0)
-    obj = pickle.load(f)
-    f.close()
-    return obj
-
-@pytest.mark.parametrize('dtype',
-        [F.float32, F.int32] if dgl.backend.backend_name == "mxnet" else [F.float32, F.int32, F.bool])
-def test_pickle(dtype):
-    f = create_test_data(dtype=dtype)
-    newf = _reconstruct_pickle(f)
-    assert F.array_equal(f['a1'], newf['a1'])
-    assert F.array_equal(f['a2'], newf['a2'])
-    assert F.array_equal(f['a3'], newf['a3'])
-
-if __name__ == '__main__':
-    test_create()
-    test_column1()
-    test_column2()
-    test_append1()
-    test_append2()
-    test_append3()
-    test_row1()
-    test_row2()
-    test_row3()
-    test_row4()
-    test_sharing()
-    test_slicing()
-    test_add_rows()
-    test_inplace()

tests/compute/test_hetero_basics.py

-"""Test from `test_basics.py` but for heterograph. Merge this
-with `test_basics.py` once DGLHeteroGraph is compatible with DGLGraph.
-"""
-import backend as F
-import dgl
-import networkx as nx
-from collections import defaultdict as ddict
-import unittest
-import pytest
-import inspect
-from utils import parametrize_dtype
-
-D = 5
-reduce_msg_shapes = set()
-
-def message_func(edges):
-    assert F.ndim(edges.src['h']) == 2
-    assert F.shape(edges.src['h'])[1] == D
-    return {'m' : edges.src['h']}
-
-def reduce_func(nodes):
-    msgs = nodes.mailbox['m']
-    reduce_msg_shapes.add(tuple(msgs.shape))
-    assert F.ndim(msgs) == 3
-    assert F.shape(msgs)[2] == D
-    return {'accum' : F.sum(msgs, 1)}
-
-def apply_node_func(nodes):
-    return {'h' : nodes.data['h'] + nodes.data['accum']}
-
-def generate_graph(idtype=F.int64, grad=False):
-    '''
-    s, d, eid
-    0, 1, 0
-    1, 9, 1
-    0, 2, 2
-    2, 9, 3
-    0, 3, 4
-    3, 9, 5
-    0, 4, 6
-    4, 9, 7
-    0, 5, 8
-    5, 9, 9
-    0, 6, 10
-    6, 9, 11
-    0, 7, 12
-    7, 9, 13
-    0, 8, 14
-    8, 9, 15
-    9, 0, 16
-    '''
-    u = F.tensor([0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 9])
-    v = F.tensor([1, 9, 2, 9, 3, 9, 4, 9, 5, 9, 6, 9, 7, 9, 8, 9, 0])
-    g = dgl.graph((u, v), idtype=idtype)
-    assert g.device == F.ctx()
-    ncol = F.randn((10, D))
-    ecol = F.randn((17, D))
-    if grad:
-        ncol = F.attach_grad(ncol)
-        ecol = F.attach_grad(ecol)
-
-    g.ndata['h'] = ncol
-    g.edata['w'] = ecol
-    g.set_n_initializer(dgl.init.zero_initializer)
-    g.set_e_initializer(dgl.init.zero_initializer)
-    return g
-
-
-@parametrize_dtype
-def test_isolated_nodes(idtype):
-    g = dgl.graph([(0, 1), (1, 2)], num_nodes=5, idtype=idtype, device=F.ctx())
-    assert g.number_of_nodes() == 5
-
-    # Test backward compatibility
-    g = dgl.graph([(0, 1), (1, 2)], card=5, idtype=idtype, device=F.ctx())
-    assert g.number_of_nodes() == 5
-
-    g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
-                      'game', num_nodes=(5, 7), idtype=idtype, device=F.ctx())
-    assert g.idtype == idtype
-    assert g.number_of_nodes('user') == 5
-    assert g.number_of_nodes('game') == 7
-
-    # Test backward compatibility
-    g = dgl.bipartite([(0, 2), (0, 3), (1, 2)], 'user', 'plays',
-                      'game', card=(5, 7), idtype=idtype, device=F.ctx())
-    assert g.idtype == idtype
-    assert g.number_of_nodes('user') == 5
-    assert g.number_of_nodes('game') == 7
-
-@parametrize_dtype
-def test_batch_setter_getter(idtype):
-    def _pfc(x):
-        return list(F.zerocopy_to_numpy(x)[:,0])
-    g = generate_graph(idtype)
-    # set all nodes
-    g.ndata['h'] = F.zeros((10, D))
-    assert F.allclose(g.ndata['h'], F.zeros((10, D)))
-    # pop nodes
-    old_len = len(g.ndata)
-    assert _pfc(g.ndata.pop('h')) == [0.] * 10
-    assert len(g.ndata) == old_len - 1
-    g.ndata['h'] = F.zeros((10, D))
-    # set partial nodes
-    u = F.tensor([1, 3, 5], idtype)
-    g.nodes[u].data['h'] = F.ones((3, D))
-    assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
-    # get partial nodes
-    u = F.tensor([1, 2, 3], idtype)
-    assert _pfc(g.nodes[u].data['h']) == [1., 0., 1.]
-
-    '''
-    s, d, eid
-    0, 1, 0
-    1, 9, 1
-    0, 2, 2
-    2, 9, 3
-    0, 3, 4
-    3, 9, 5
-    0, 4, 6
-    4, 9, 7
-    0, 5, 8
-    5, 9, 9
-    0, 6, 10
-    6, 9, 11
-    0, 7, 12
-    7, 9, 13
-    0, 8, 14
-    8, 9, 15
-    9, 0, 16
-    '''
-    # set all edges
-    g.edata['l'] = F.zeros((17, D))
-    assert _pfc(g.edata['l']) == [0.] * 17
-    # pop edges
-    old_len = len(g.edata)
-    assert _pfc(g.edata.pop('l')) == [0.] * 17
-    assert len(g.edata) == old_len - 1
-    g.edata['l'] = F.zeros((17, D))
-    # set partial edges (many-many)
-    u = F.tensor([0, 0, 2, 5, 9], dtype=idtype)
-    v = F.tensor([1, 3, 9, 9, 0], dtype=idtype)
-    g.edges[u, v].data['l'] = F.ones((5, D))
-    truth = [0.] * 17
-    truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
-    assert _pfc(g.edata['l']) == truth
-    # set partial edges (many-one)
-    u = F.tensor([3, 4, 6], dtype=idtype)
-    v = F.tensor([9], dtype=idtype)
-    g.edges[u, v].data['l'] = F.ones((3, D))
-    truth[5] = truth[7] = truth[11] = 1.
-    assert _pfc(g.edata['l']) == truth
-    # set partial edges (one-many)
-    u = F.tensor([0], dtype=idtype)
-    v = F.tensor([4, 5, 6], dtype=idtype)
-    g.edges[u, v].data['l'] = F.ones((3, D))
-    truth[6] = truth[8] = truth[10] = 1.
-    assert _pfc(g.edata['l']) == truth
-    # get partial edges (many-many)
-    u = F.tensor([0, 6, 0], dtype=idtype)
-    v = F.tensor([6, 9, 7], dtype=idtype)
-    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
-    # get partial edges (many-one)
-    u = F.tensor([5, 6, 7], dtype=idtype)
-    v = F.tensor([9], dtype=idtype)
-    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
-    # get partial edges (one-many)
-    u = F.tensor([0], dtype=idtype)
-    v = F.tensor([3, 4, 5], dtype=idtype)
-    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
-
-
-@parametrize_dtype
-def test_batch_setter_autograd(idtype):
-    g = generate_graph(idtype=idtype, grad=True)
-    h1 = g.ndata['h']
-    # partial set
-    v = F.tensor([1, 2, 8], idtype)
-    hh = F.attach_grad(F.zeros((len(v), D)))
-    with F.record_grad():
-        g.nodes[v].data['h'] = hh
-        h2 = g.ndata['h']
-        F.backward(h2, F.ones((10, D)) * 2)
-    assert F.array_equal(F.grad(h1)[:,0], F.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
-    assert F.array_equal(F.grad(hh)[:,0], F.tensor([2., 2., 2.]))
-
-
-@parametrize_dtype
-def test_nx_conversion(idtype):
-    # check conversion between networkx and DGLGraph
-
-    def _check_nx_feature(nxg, nf, ef):
-        # check node and edge feature of nxg
-        # this is used to check to_networkx
-        num_nodes = len(nxg)
-        num_edges = nxg.size()
-        if num_nodes > 0:
-            node_feat = ddict(list)
-            for nid, attr in nxg.nodes(data=True):
-                assert len(attr) == len(nf)
-                for k in nxg.nodes[nid]:
-                    node_feat[k].append(F.unsqueeze(attr[k], 0))
-            for k in node_feat:
-                feat = F.cat(node_feat[k], 0)
-                assert F.allclose(feat, nf[k])
-        else:
-            assert len(nf) == 0
-        if num_edges > 0:
-            edge_feat = ddict(lambda: [0] * num_edges)
-            for u, v, attr in nxg.edges(data=True):
-                assert len(attr) == len(ef) + 1 # extra id
-                eid = attr['id']
-                for k in ef:
-                    edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
-            for k in edge_feat:
-                feat = F.cat(edge_feat[k], 0)
-                assert F.allclose(feat, ef[k])
-        else:
-            assert len(ef) == 0
-
-    n1 = F.randn((5, 3))
-    n2 = F.randn((5, 10))
-    n3 = F.randn((5, 4))
-    e1 = F.randn((4, 5))
-    e2 = F.randn((4, 7))
-    g = dgl.graph([(0,2),(1,4),(3,0),(4,3)], idtype=idtype, device=F.ctx())
-    g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
-    g.edata.update({'e1': e1, 'e2': e2})
-
-    # convert to networkx
-    nxg = dgl.to_networkx(g.cpu(), node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
-    assert len(nxg) == 5
-    assert nxg.size() == 4
-    _check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})
-
-    # convert to DGLGraph, nx graph has id in edge feature
-    # use id feature to test non-tensor copy
-    g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'], idtype=idtype)
-    assert g.idtype == idtype
-    assert g.device == F.cpu()
-    g = g.to(F.ctx())
-    # check graph size
-    assert g.number_of_nodes() == 5
-    assert g.number_of_edges() == 4
-    # check number of features
-    # test with existing dglgraph (so existing features should be cleared)
-    assert len(g.ndata) == 1
-    assert len(g.edata) == 2
-    # check feature values
-    assert F.allclose(g.ndata['n1'], n1)
-    # with id in nx edge feature, e1 should follow original order
-    assert F.allclose(g.edata['e1'], e1)
-    assert F.array_equal(g.edata['id'], F.arange(0, 4, F.dtype(g.edata['id'])))
-
-    # test conversion after modifying DGLGraph
-    # TODO(minjie): enable after mutation is supported
-    #g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
-    #new_n = F.randn((2, 3))
-    #new_e = F.randn((3, 5))
-    #g.add_nodes(2, data={'n1': new_n})
-    ## add three edges, one is a multi-edge
-    #g.add_edges([3, 6, 0], [4, 5, 2], data={'e1': new_e})
-    #n1 = F.cat((n1, new_n), 0)
-    #e1 = F.cat((e1, new_e), 0)
-    ## convert to networkx again
-    #nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
-    #assert len(nxg) == 7
-    #assert nxg.size() == 7
-    #_check_nx_feature(nxg, {'n1': n1}, {'e1': e1})
-
-    # now test convert from networkx without id in edge feature
-    # first pop id in edge feature
-    for _, _, attr in nxg.edges(data=True):
-        attr.pop('id')
-    # test with a new graph
-    g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'], idtype=idtype)
-    # check graph size
-    assert g.number_of_nodes() == 5
-    assert g.number_of_edges() == 4
-    # check number of features
-    assert len(g.ndata) == 1
-    assert len(g.edata) == 1
-    # check feature values
-    assert F.allclose(g.ndata['n1'], n1)
-    # edge feature order follows nxg.edges()
-    edge_feat = []
-    for _, _, attr in nxg.edges(data=True):
-        edge_feat.append(F.unsqueeze(attr['e1'], 0))
-    edge_feat = F.cat(edge_feat, 0)
-    assert F.allclose(g.edata['e1'], edge_feat)
-
-@parametrize_dtype
-def test_apply_nodes(idtype):
-    def _upd(nodes):
-        return {'h' : nodes.data['h'] * 2}
-    g = generate_graph(idtype=idtype)
-    old = g.ndata['h']
-    g.apply_nodes(_upd)
-    assert F.allclose(old * 2, g.ndata['h'])
-    u = F.tensor([0, 3, 4, 6], idtype)
-    g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
-    assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
-
-@parametrize_dtype
-def test_apply_edges(idtype):
-    def _upd(edges):
-        return {'w' : edges.data['w'] * 2}
-    g = generate_graph(idtype=idtype)
-    old = g.edata['w']
-    g.apply_edges(_upd)
-    assert F.allclose(old * 2, g.edata['w'])
-    u = F.tensor([0, 0, 0, 4, 5, 6], idtype)
-    v = F.tensor([1, 2, 3, 9, 9, 9], idtype)
-    g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
-    eid = F.tensor(g.edge_ids(u, v), idtype)
-    assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
-
-@parametrize_dtype
-def test_update_routines(idtype):
-    g = generate_graph(idtype=idtype)
-
-    # send_and_recv
-    reduce_msg_shapes.clear()
-    u = [0, 0, 0, 4, 5, 6]
-    v = [1, 2, 3, 9, 9, 9]
-    g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
-    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
-    reduce_msg_shapes.clear()
-    try:
-        g.send_and_recv([u, v], message_func, reduce_func, apply_node_func)
-        assert False
-    except dgl.DGLError:
-        pass
-
-    # pull
-    v = F.tensor([1, 2, 3, 9], idtype)
-    reduce_msg_shapes.clear()
-    g.pull(v, message_func, reduce_func, apply_node_func)
-    assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
-    reduce_msg_shapes.clear()
-
-    # push
-    v = F.tensor([0, 1, 2, 3], idtype)
-    reduce_msg_shapes.clear()
-    g.push(v, message_func, reduce_func, apply_node_func)
-    assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
-    reduce_msg_shapes.clear()
-
-    # update_all
-    reduce_msg_shapes.clear()
-    g.update_all(message_func, reduce_func, apply_node_func)
-    assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
-    reduce_msg_shapes.clear()
-
-@parametrize_dtype
-def test_update_all_0deg(idtype):
-    # test#1
-    g = dgl.graph([(1,0), (2,0), (3,0), (4,0)], idtype=idtype, device=F.ctx())
-    def _message(edges):
-        return {'m' : edges.src['h']}
-    def _reduce(nodes):
-        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
-    def _apply(nodes):
-        return {'h' : nodes.data['h'] * 2}
-    def _init2(shape, dtype, ctx, ids):
-        return 2 + F.zeros(shape, dtype, ctx)
-    g.set_n_initializer(_init2, 'h')
-    old_repr = F.randn((5, 5))
-    g.ndata['h'] = old_repr
-    g.update_all(_message, _reduce, _apply)
-    new_repr = g.ndata['h']
-    # the first row of the new_repr should be the sum of all the node
-    # features; while the 0-deg nodes should be initialized by the
-    # initializer and applied with UDF.
-    assert F.allclose(new_repr[1:], 2*(2+F.zeros((4,5))))
-    assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
-
-    # test#2:
-    g = dgl.graph([], num_nodes=5, idtype=idtype, device=F.ctx())
-    g.set_n_initializer(_init2, 'h')
-    g.ndata['h'] = old_repr
-    g.update_all(_message, _reduce, _apply)
-    new_repr = g.ndata['h']
-    # should fallback to apply
-    assert F.allclose(new_repr, 2*old_repr)
-
-@parametrize_dtype
-def test_pull_0deg(idtype):
-    g = dgl.graph([(0,1)], idtype=idtype, device=F.ctx())
-    def _message(edges):
-        return {'m' : edges.src['h']}
-    def _reduce(nodes):
-        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
-    def _apply(nodes):
-        return {'h' : nodes.data['h'] * 2}
-    def _init2(shape, dtype, ctx, ids):
-        return 2 + F.zeros(shape, dtype, ctx)
-    g.set_n_initializer(_init2, 'h')
-    # test#1: pull both 0deg and non-0deg nodes
-    old = F.randn((2, 5))
-    g.ndata['h'] = old
-    g.pull([0, 1], _message, _reduce, _apply)
-    new = g.ndata.pop('h')
-    # 0deg check: initialized with the func and got applied
-    assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
-    # non-0deg check
-    assert F.allclose(new[1], F.sum(old, 0) * 2)
-
-    # test#2: pull only 0deg node
-    old = F.randn((2, 5))
-    g.ndata['h'] = old
-    g.pull(0, _message, _reduce, _apply)
-    new = g.ndata.pop('h')
-    # 0deg check: fallback to apply
-    assert F.allclose(new[0], 2*old[0])
-    # non-0deg check: not touched
-    assert F.allclose(new[1], old[1])
-
-@parametrize_dtype
-def test_send_multigraph(idtype):
-    g = dgl.graph([(0,1), (0,1), (0,1), (2,1)], idtype=idtype, device=F.ctx())
-
-    def _message_a(edges):
-        return {'a': edges.data['a']}
-    def _message_b(edges):
-        return {'a': edges.data['a'] * 3}
-    def _reduce(nodes):
-        return {'a': F.max(nodes.mailbox['a'], 1)}
-
-    def answer(*args):
-        return F.max(F.stack(args, 0), 0)
-
-    assert g.is_multigraph
-
-    # send by eid
-    old_repr = F.randn((4, 5))
-    # send_and_recv_on
-    g.ndata['a'] = F.zeros((3, 5))
-    g.edata['a'] = old_repr
-    g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
-    new_repr = g.ndata['a']
-    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
-    assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
-
-# Disabled - Heterograph doesn't support mutation
-def _test_dynamic_addition():
-    N = 3
-    D = 1
-
-    g = dgl.DGLGraph()
-
-    # Test node addition
-    g.add_nodes(N)
-    g.ndata.update({'h1': F.randn((N, D)),
-                    'h2': F.randn((N, D))})
-    g.add_nodes(3)
-    assert g.ndata['h1'].shape[0] == g.ndata['h2'].shape[0] == N + 3
-
-    # Test edge addition
-    g.add_edge(0, 1)
-    g.add_edge(1, 0)
-    g.edata.update({'h1': F.randn((2, D)),
-                    'h2': F.randn((2, D))})
-    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 2
-
-    g.add_edges([0, 2], [2, 0])
-    g.edata['h1'] = F.randn((4, D))
-    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 4
-
-    g.add_edge(1, 2)
-    g.edges[4].data['h1'] = F.randn((1, D))
-    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 5
-
-    # test add edge with part of the features
-    g.add_edge(2, 1, {'h1': F.randn((1, D))})
-    assert len(g.edata['h1']) == len(g.edata['h2'])
-
-@parametrize_dtype
-def test_repr(idtype):
-    G = dgl.graph([(0,1), (0,2), (1,2)], num_nodes=10, idtype=idtype, device=F.ctx())
-    repr_string = G.__repr__()
-    print(repr_string)
-    G.ndata['x'] = F.zeros((10, 5))
-    G.edata['y'] = F.zeros((3, 4))
-    repr_string = G.__repr__()
-    print(repr_string)
-
-
-@parametrize_dtype
-def test_group_apply_edges(idtype):
-    def edge_udf(edges):
-        h = F.sum(edges.data['feat'] * (edges.src['h'] + edges.dst['h']), dim=2)
-        normalized_feat = F.softmax(h, dim=1)
-        return {"norm_feat": normalized_feat}
-
-    elist = []
-    for v in [1, 2, 3, 4, 5, 6, 7, 8]:
-        elist.append((0, v))
-    for v in [2, 3, 4, 6, 7, 8]:
-        elist.append((1, v))
-    for v in [2, 3, 4, 5, 6, 7, 8]:
-        elist.append((2, v))
-    g = dgl.graph(elist, idtype=idtype, device=F.ctx())
-
-    g.ndata['h'] = F.randn((g.number_of_nodes(), D))
-    g.edata['feat'] = F.randn((g.number_of_edges(), D))
-
-    def _test(group_by):
-        g.group_apply_edges(group_by=group_by, func=edge_udf)
-        if group_by == 'src':
-            u, v, eid = g.out_edges(1, form='all')
-        else:
-            u, v, eid = g.in_edges(5, form='all')
-        out_feat = g.edges[eid].data['norm_feat']
-        result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
-        result = F.softmax(F.sum(result, dim=1), dim=0)
-        assert F.allclose(out_feat, result)
-
-    # test group by source nodes
-    _test('src')
-
-    # test group by destination nodes
-    _test('dst')
-
-@parametrize_dtype
-def test_local_var(idtype):
-    g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
-    g = g.to(F.ctx())
-    g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
-    g.edata['w'] = F.zeros((g.number_of_edges(), 4))
-    # test override
-    def foo(g):
-        g = g.local_var()
-        g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
-        g.edata['w'] = F.ones((g.number_of_edges(), 4))
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test out-place update
-    def foo(g):
-        g = g.local_var()
-        g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
-        g.edges[[2, 3]].data['w'] = F.ones((2, 4))
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test out-place update 2
-    def foo(g):
-        g = g.local_var()
-        g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
-        g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test auto-pop
-    def foo(g):
-        g = g.local_var()
-        g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
-        g.edata['ww'] = F.ones((g.number_of_edges(), 4))
-    foo(g)
-    assert 'hh' not in g.ndata
-    assert 'ww' not in g.edata
-
-    # test initializer1
-    g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
-    g = g.to(F.ctx())
-    g.set_n_initializer(dgl.init.zero_initializer)
-    def foo(g):
-        g = g.local_var()
-        g.nodes[0].data['h'] = F.ones((1, 1))
-        assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
-    foo(g)
-    # test initializer2
-    def foo_e_initializer(shape, dtype, ctx, id_range):
-        return F.ones(shape)
-    g.set_e_initializer(foo_e_initializer, field='h')
-    def foo(g):
-        g = g.local_var()
-        g.edges[0, 1].data['h'] = F.ones((1, 1))
-        assert F.allclose(g.edata['h'], F.ones((2, 1)))
-        g.edges[0, 1].data['w'] = F.ones((1, 1))
-        assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
-    foo(g)
-
-@parametrize_dtype
-def test_local_scope(idtype):
-    g = dgl.graph([(0,1), (1,2), (2,3), (3,4)], idtype=idtype, device=F.ctx())
-    g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
-    g.edata['w'] = F.zeros((g.number_of_edges(), 4))
-    # test override
-    def foo(g):
-        with g.local_scope():
-            g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
-            g.edata['w'] = F.ones((g.number_of_edges(), 4))
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test out-place update
-    def foo(g):
-        with g.local_scope():
-            g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
-            g.edges[[2, 3]].data['w'] = F.ones((2, 4))
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test out-place update 2
-    def foo(g):
-        with g.local_scope():
-            g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
-            g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
-    foo(g)
-    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
-    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
-    # test auto-pop
-    def foo(g):
-        with g.local_scope():
-            g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
-            g.edata['ww'] = F.ones((g.number_of_edges(), 4))
-    foo(g)
-    assert 'hh' not in g.ndata
-    assert 'ww' not in g.edata
-
-    # test nested scope
-    def foo(g):
-        with g.local_scope():
-            g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
-            g.edata['ww'] = F.ones((g.number_of_edges(), 4))
-            with g.local_scope():
-                g.ndata['hhh'] = F.ones((g.number_of_nodes(), 3))
-                g.edata['www'] = F.ones((g.number_of_edges(), 4))
-            assert 'hhh' not in g.ndata
-            assert 'www' not in g.edata
-    foo(g)
-    assert 'hh' not in g.ndata
-    assert 'ww' not in g.edata
-
-    # test initializer1
-    g = dgl.graph([(0,1), (1,1)], idtype=idtype, device=F.ctx())
-    g.set_n_initializer(dgl.init.zero_initializer)
-    def foo(g):
-        with g.local_scope():
-            g.nodes[0].data['h'] = F.ones((1, 1))
-            assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
-    foo(g)
-    # test initializer2
-    def foo_e_initializer(shape, dtype, ctx, id_range):
-        return F.ones(shape)
-    g.set_e_initializer(foo_e_initializer, field='h')
-    def foo(g):
-        with g.local_scope():
-            g.edges[0, 1].data['h'] = F.ones((1, 1))
-            assert F.allclose(g.edata['h'], F.ones((2, 1)))
-            g.edges[0, 1].data['w'] = F.ones((1, 1))
-            assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
-    foo(g)
-
-@parametrize_dtype
-def test_issue_1088(idtype):
-    # This test ensures that message passing on a heterograph with one edge type
-    # would not crash (GitHub issue #1088).
-    import dgl.function as fn
-    g = dgl.heterograph({('U', 'E', 'V'): ([0, 1, 2], [1, 2, 3])}, idtype=idtype, device=F.ctx())
-    g.nodes['U'].data['x'] = F.randn((3, 3))
-    g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'y'))
-
-if __name__ == '__main__':
-    #test_isolated_nodes("int32")
-    test_batch_setter_getter(F.int32)
-    # test_batch_recv("int64")
-    # test_apply_edges("int32")
-    # test_batch_setter_autograd()
-    # test_batch_send()
-    # test_batch_recv()
-    # test_apply_nodes()
-    # test_apply_edges()
-    #test_update_routines(F.int32)
-    # test_recv_0deg()
-    # test_recv_0deg_newfld()
-    # test_update_all_0deg()
-    # test_pull_0deg()
-    # test_send_multigraph()
-    # test_dynamic_addition()
-    # test_repr()
-    # test_group_apply_edges()
-    # test_local_var()
-    # test_local_scope()
-    #test_issue_1088('int64')
-    pass

tests/compute/test_heterograph.py

@@ -1410,44 +1410,6 @@ def test_level2(idtype):
     with pytest.raises(DGLError):
         g.send_and_recv([2, 3], mfunc, rfunc)
 
-    # 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
-    with pytest.raises(DGLError):
-        g.multi_send_and_recv(
-            {'plays' : (g.edges(etype='plays'), mfunc, rfunc),
-             'follows': (g.edges(etype='follows'), mfunc, rfunc2)},
-            'sum')
-
     g.nodes['game'].data.clear()
 
     #############################################################
@@ -1468,49 +1430,6 @@ def test_level2(idtype):
     with pytest.raises(DGLError):
         g.pull(1, mfunc, rfunc)
 
-    # 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
-    with pytest.raises(DGLError):
-        g.multi_pull(
-            1,
-            {'plays' : (mfunc, rfunc),
-             'follows': (mfunc, rfunc2)},
-            'sum')
-
     g.nodes['game'].data.clear()
 
     #############################################################
@@ -1672,8 +1591,8 @@ def test_empty_heterograph(idtype):
     assert g.number_of_edges('develops') == 2
     assert g.number_of_nodes('developer') == 2
 
-
-def test_types_in_function():
+@parametrize_dtype
+def test_types_in_function(idtype):
     def mfunc1(edges):
         assert edges.canonical_etype == ('user', 'follow', 'user')
         return {}
@@ -1706,7 +1625,7 @@ def test_types_in_function():
         assert edges.canonical_etype == ('user', 'plays', 'game')
         return F.zeros((2,))
 
-    g = dgl.graph([(0, 1), (1, 2)], 'user', 'follow')
+    g = dgl.graph([(0, 1), (1, 2)], 'user', 'follow', idtype=idtype, device=F.ctx())
     g.apply_nodes(rfunc1)
     g.apply_edges(mfunc1)
     g.update_all(mfunc1, rfunc1)
@@ -1716,7 +1635,7 @@ def test_types_in_function():
     g.filter_nodes(filter_nodes1)
     g.filter_edges(filter_edges1)
 
-    g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
+    g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
     g.apply_nodes(rfunc2, ntype='game')
     g.apply_edges(mfunc2)
     g.update_all(mfunc2, rfunc2)

tests/compute/test_inplace_update.py.bak

-import numpy as np
-import scipy.sparse as sp
-import dgl
-import dgl.function as fn
-import backend as F
-import unittest
-
-D = 5
-
-def generate_graph():
-    g = dgl.DGLGraph()
-    g.add_nodes(10)
-    # create a graph where 0 is the source and 9 is the sink
-    for i in range(1, 9):
-        g.add_edge(0, i)
-        g.add_edge(i, 9)
-    # add a back flow from 9 to 0
-    g.add_edge(9, 0)
-    g.ndata['f'] = F.randn((10, D))
-    g.edata['e'] = F.randn((17, D))
-    return g
-
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_inplace_snr():
-    u = F.tensor([0, 0, 0, 3, 4, 9])
-    v = F.tensor([1, 2, 3, 9, 9, 0])
-
-    def message_func(edges):
-        return {'m' : edges.src['f']}
-
-    def reduce_func(nodes):
-        return {'f' : F.sum(nodes.mailbox['m'], 1)}
-
-    def apply_func(nodes):
-        return {'f' : 2 * nodes.data['f']}
-
-    def _test(apply_func):
-        g = generate_graph()
-        f = g.ndata['f']
-
-        # an out place run to get result
-        g.send_and_recv((u, v), fn.copy_src(src='f', out='m'),
-                fn.sum(msg='m', out='f'), apply_func)
-        result = g.ndata['f']
-
-        # inplace deg bucket
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.send_and_recv((u, v), message_func, reduce_func, apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace v2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.send_and_recv((u, v), fn.copy_src(src='f', out='m'),
-                        fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace e2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.send_and_recv((u, v), message_func,
-                        fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-    # test send_and_recv with apply_func
-    _test(apply_func)
-    # test send_and_recv without apply_func
-    _test(None)
-
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_inplace_push():
-    nodes = F.tensor([0, 3, 4, 9])
-
-    def message_func(edges):
-        return {'m' : edges.src['f']}
-
-    def reduce_func(nodes):
-        return {'f' : F.sum(nodes.mailbox['m'], 1)}
-
-    def apply_func(nodes):
-        return {'f' : 2 * nodes.data['f']}
-
-    def _test(apply_func):
-        g = generate_graph()
-        f = g.ndata['f']
-
-        # an out place run to get result
-        g.push(nodes,
-               fn.copy_src(src='f', out='m'), fn.sum(msg='m', out='f'), apply_func)
-        result = g.ndata['f']
-
-        # inplace deg bucket
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.push(nodes, message_func, reduce_func, apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace v2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.push(nodes, fn.copy_src(src='f', out='m'),
-               fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace e2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.push(nodes,
-               message_func, fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-    # test send_and_recv with apply_func
-    _test(apply_func)
-    # test send_and_recv without apply_func
-    _test(None)
-
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_inplace_pull():
-    nodes = F.tensor([1, 2, 3, 9])
-
-    def message_func(edges):
-        return {'m' : edges.src['f']}
-
-    def reduce_func(nodes):
-        return {'f' : F.sum(nodes.mailbox['m'], 1)}
-
-    def apply_func(nodes):
-        return {'f' : 2 * nodes.data['f']}
-
-    def _test(apply_func):
-        g = generate_graph()
-        f = g.ndata['f']
-
-        # an out place run to get result
-        g.pull(nodes,
-               fn.copy_src(src='f', out='m'), fn.sum(msg='m', out='f'), apply_func)
-        result = g.ndata['f']
-
-        # inplace deg bucket
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.pull(nodes, message_func, reduce_func, apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace v2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.pull(nodes, fn.copy_src(src='f', out='m'),
-               fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-        # inplace e2v spmv
-        v1 = F.clone(f)
-        g.ndata['f'] = v1
-        g.pull(nodes,
-               message_func, fn.sum(msg='m', out='f'), apply_func, inplace=True)
-        r1 = g.ndata['f']
-        # check result
-        assert F.allclose(r1, result)
-        # check inplace
-        assert F.allclose(v1, r1)
-
-    # test send_and_recv with apply_func
-    _test(apply_func)
-    # test send_and_recv without apply_func
-    _test(None)
-
-
-@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
-def test_inplace_apply():
-    def apply_node_func(nodes):
-        return {'f': nodes.data['f'] * 2}
-
-    def apply_edge_func(edges):
-        return {'e': edges.data['e'] * 2}
-
-    g = generate_graph()
-    nodes = [1, 2, 3, 9]
-    nf = g.ndata['f']
-    # out place run
-    g.apply_nodes(apply_node_func, nodes)
-    new_nf = g.ndata['f']
-    # in place run
-    g.ndata['f'] = nf
-    g.apply_nodes(apply_node_func, nodes, inplace=True)
-    # check results correct and in place
-    assert F.allclose(nf, new_nf)
-    # test apply all nodes, should not be done in place
-    g.ndata['f'] = nf
-    g.apply_nodes(apply_node_func, inplace=True)
-    assert (F.allclose(nf, g.ndata['f']) == False)
-
-    edges = [3, 5, 7, 10]
-    ef = g.edata['e']
-    # out place run
-    g.apply_edges(apply_edge_func, edges)
-    new_ef = g.edata['e']
-    # in place run
-    g.edata['e'] = ef
-    g.apply_edges(apply_edge_func, edges, inplace=True)
-    g.edata['e'] = ef
-    assert F.allclose(ef, new_ef)
-    # test apply all edges, should not be done in place
-    g.edata['e'] == ef
-    g.apply_edges(apply_edge_func, inplace=True)
-    assert F.allclose(ef, g.edata['e']) == False
-
-if __name__ == '__main__':
-    test_inplace_snr()
-    test_inplace_push()
-    test_inplace_pull()
-    test_inplace_apply()

tests/compute/test_kernel.py

@@ -4,6 +4,8 @@ import networkx as nx
 import numpy as np
 import backend as F
 from itertools import product
+from test_utils import parametrize_dtype, get_cases
+import pytest
 
 def udf_copy_src(edges):
     return {'m': edges.src['u']}
@@ -179,6 +181,7 @@ def test_copy_edge_reduce():
         def _print_error(a, b):
             print("ERROR: Test copy_edge_{} partial: {}".
                   format(red, partial))
+            return
             for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
                 if not np.allclose(x, y):
                     print('@{} {} v.s. {}'.format(i, x, y))
@@ -291,19 +294,13 @@ def test_all_binary_builtins():
             lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
             rhs_grad_2 = F.grad(target_feature_switch(g, rhs))
 
-        if reducer == 'prod':
-            # increase tolerance for prod reducer
-            # NOTE(zihao) as far as I know prod reducer has never
-            # been used in any gnn models.
-            rtol = 1e-2
-            atol = 1e-2
-        else:
-            rtol = 1e-4
-            atol = 1e-4
+        rtol = 1e-4
+        atol = 1e-4
 
         def _print_error(a, b):
             print("ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}".
                   format(lhs, binary_op, rhs, reducer, broadcast, partial))
+            return
             if lhs == 'u':
                 lhs_data = hu
             elif lhs == 'v':
@@ -357,15 +354,25 @@ def test_all_binary_builtins():
     for lhs, rhs in product(target, target):
         if lhs == rhs:
             continue
-        for binary_op in ["add", "sub", "mul", "div", "dot"]:
-            for reducer in ["sum", "max", "min", "prod", "mean"]:
+        for binary_op in ["add", "sub", "mul", "div"]:
+            for reducer in ["sum", "max", "min", "mean"]:
                 for broadcast in ["none", lhs, rhs]:
                     for partial in [False, True]:
+                        print(lhs, rhs, binary_op, reducer, broadcast, partial)
                         _test(g, lhs, rhs, binary_op, reducer, partial, nid,
                               broadcast=broadcast)
 
+@parametrize_dtype
+@pytest.mark.parametrize('g', get_cases(['homo-zero-degree']))
+def test_mean_zero_degree(g, idtype):
+    g = g.astype(idtype).to(F.ctx())
+    g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
+    g.update_all(fn.copy_u('h', 'm'), fn.mean('m', 'x'))
+    deg = F.asnumpy(g.in_degrees())
+    v = F.tensor(np.where(deg == 0)[0])
+    assert F.allclose(F.gather_row(g.ndata['x'], v), F.zeros((len(v), 3)))
+
 if __name__ == '__main__':
     test_copy_src_reduce()
     test_copy_edge_reduce()
     test_all_binary_builtins()
-

tests/compute/test_pickle.py

@@ -2,7 +2,6 @@ import networkx as nx
 import scipy.sparse as ssp
 import dgl
 import dgl.contrib as contrib
-from dgl.frame import Frame, FrameRef, Column
 from dgl.graph_index import create_graph_index
 from dgl.utils import toindex
 import backend as F
@@ -128,24 +127,6 @@ def test_pickling_graph_index():
     assert F.array_equal(dst_idx.tousertensor(), dst_idx2.tousertensor())
 
 
-def test_pickling_frame():
-    x = F.randn((3, 7))
-    y = F.randn((3, 5))
-
-    c = Column(x)
-
-    c2 = _reconstruct_pickle(c)
-    assert F.allclose(c.data, c2.data)
-
-    fr = Frame({'x': x, 'y': y})
-
-    fr2 = _reconstruct_pickle(fr)
-    assert F.allclose(fr2['x'].data, x)
-    assert F.allclose(fr2['y'].data, y)
-
-    fr = Frame()
-
-
 def _global_message_func(nodes):
     return {'x': nodes.data['x']}
 

tests/compute/test_shared_mem.py

@@ -2,7 +2,6 @@ import networkx as nx
 import scipy.sparse as ssp
 import dgl
 import dgl.contrib as contrib
-from dgl.frame import Frame, FrameRef, Column
 from dgl.graph_index import create_graph_index
 from dgl.utils import toindex
 import backend as F

tests/compute/test_specialization.py

@@ -147,171 +147,6 @@ def test_v2v_pull(idtype):
     # test 2d node features
     _test('f2')
 
-@parametrize_dtype
-def test_v2v_update_all_multi_fn(idtype):
-    def message_func(edges):
-        return {'m2': edges.src['f2']}
-
-    def message_func_edge(edges):
-        return {'m2': edges.src['f2'] * edges.data['e2']}
-
-    def reduce_func(nodes):
-        return {'v1': F.sum(nodes.mailbox['m2'], 1)}
-
-    g = generate_graph(idtype)
-    g.ndata.update({'v1' : F.zeros((10,)), 'v2' : F.zeros((10,))})
-    fld = 'f2'
-
-    g.update_all(message_func, reduce_func)
-    v1 = g.ndata['v1']
-
-    # 1 message, 2 reduces
-    g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')])
-    v2 = g.ndata['v2']
-    v3 = g.ndata['v3']
-    assert F.allclose(v1, v2)
-    assert F.allclose(v1, v3)
-
-    # update all with edge weights, 2 message, 3 reduces
-    g.update_all([fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
-                 [fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
-                 None)
-    v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    v3 = g.ndata['v3']
-    assert F.allclose(v1, v2)
-    assert F.allclose(v1, v3)
-
-    # run UDF with single message and reduce
-    g.update_all(message_func_edge, reduce_func, None)
-    v2 = g.ndata['v2']
-    assert F.allclose(v1, v2)
-
-@parametrize_dtype
-def test_v2v_snr_multi_fn(idtype):
-    u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
-    v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
-
-    def message_func(edges):
-        return {'m2': edges.src['f2']}
-
-    def message_func_edge(edges):
-        return {'m2': edges.src['f2'] * edges.data['e2']}
-
-    def reduce_func(nodes):
-        return {'v1' : F.sum(nodes.mailbox['m2'], 1)}
-
-    g = generate_graph(idtype)
-    g.ndata.update({'v1' : F.zeros((10, D)), 'v2' : F.zeros((10, D)),
-        'v3' : F.zeros((10, D))})
-    fld = 'f2'
-
-    g.send_and_recv((u, v), message_func, reduce_func)
-    v1 = g.ndata['v1']
-
-    # 1 message, 2 reduces
-    g.send_and_recv((u, v),
-            fn.copy_src(src=fld, out='m'),
-            [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')],
-            None)
-    v2 = g.ndata['v2']
-    v3 = g.ndata['v3']
-    assert F.allclose(v1, v2)
-    assert F.allclose(v1, v3)
-
-    # send and recv with edge weights, 2 message, 3 reduces
-    g.send_and_recv((u, v),
-                    [fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
-                    [fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
-                    None)
-    v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    v3 = g.ndata['v3']
-    assert F.allclose(v1, v2)
-    assert F.allclose(v1, v3)
-
-    # run UDF with single message and reduce
-    g.send_and_recv((u, v), message_func_edge,
-            reduce_func, None)
-    v2 = g.ndata['v2']
-    assert F.allclose(v1, v2)
-
-@parametrize_dtype
-def test_e2v_update_all_multi_fn(idtype):
-    def _test(fld):
-        def message_func(edges):
-            return {'m1' : edges.src[fld] + edges.dst[fld],
-                    'm2' : edges.src[fld] * edges.dst[fld]}
-
-        def reduce_func(nodes):
-            return {fld : F.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
-
-        def apply_func(nodes):
-            return {fld : 2 * nodes.data[fld]}
-
-        def apply_func_2(nodes):
-            return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
-
-        g = generate_graph(idtype)
-        # update all
-        v1 = g.ndata[fld]
-        # no specialization
-        g.update_all(message_func, reduce_func, apply_func)
-        v2 = g.ndata[fld]
-
-        # user break reduce func into 2 builtin
-        g.ndata.update({fld : v1})
-        g.update_all(message_func,
-                     [fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
-                     apply_func_2)
-        v3 = g.ndata[fld]
-
-        assert F.allclose(v2, v3)
-
-    # test 1d node features
-    _test('f1')
-    # test 2d node features
-    _test('f2')
-
-@parametrize_dtype
-def test_e2v_snr_multi_fn(idtype):
-    u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
-    v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
-    def _test(fld):
-        def message_func(edges):
-            return {'m1' : edges.src[fld] + edges.dst[fld],
-                    'm2' : edges.src[fld] * edges.dst[fld]}
-
-        def reduce_func(nodes):
-            return {fld : F.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
-
-        def apply_func(nodes):
-            return {fld : 2 * nodes.data[fld]}
-
-        def apply_func_2(nodes):
-            return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
-
-        g = generate_graph(idtype)
-        # send_and_recv
-        v1 = g.ndata[fld]
-        # no specialization
-        g.send_and_recv((u, v), message_func, reduce_func, apply_func)
-        v2 = g.ndata[fld]
-
-        # user break reduce func into 2 builtin
-        g.ndata.update({fld : v1})
-        g.send_and_recv((u, v), message_func,
-                        [fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
-                        apply_func_2)
-        v3 = g.ndata[fld]
-
-        assert F.allclose(v2, v3)
-
-    # test 1d node features
-    _test('f1')
-    # test 2d node features
-    _test('f2')
-
 @parametrize_dtype
 def test_update_all_multi_fallback(idtype):
     # create a graph with zero in degree nodes
@@ -357,18 +192,6 @@ def test_update_all_multi_fallback(idtype):
                  fn.sum(msg='m2', out='o2'),
                  _afunc)
     assert F.allclose(o2, g.ndata.pop('o2'))
-    # multi builtins, both v2v spmv
-    g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
-                 [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
-                 _afunc)
-    assert F.allclose(o1, g.ndata.pop('o1'))
-    assert F.allclose(o1, g.ndata.pop('o2'))
-    # multi builtins, one v2v spmv, one fallback to e2v
-    g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w2', out='m2')],
-                 [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
-                 _afunc)
-    assert F.allclose(o1, g.ndata.pop('o1'))
-    assert F.allclose(o2, g.ndata.pop('o2'))
 
 @parametrize_dtype
 def test_pull_multi_fallback(idtype):
@@ -417,20 +240,6 @@ def test_pull_multi_fallback(idtype):
                      fn.sum(msg='m2', out='o2'),
                      _afunc)
         assert F.allclose(o2, g.ndata.pop('o2'))
-        # multi builtins, both v2v spmv
-        g.pull(nodes,
-               [fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
-               [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
-               _afunc)
-        assert F.allclose(o1, g.ndata.pop('o1'))
-        assert F.allclose(o1, g.ndata.pop('o2'))
-        # multi builtins, one v2v spmv, one fallback to e2v
-        g.pull(nodes,
-               [fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w2', out='m2')],
-               [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
-               _afunc)
-        assert F.allclose(o1, g.ndata.pop('o1'))
-        assert F.allclose(o2, g.ndata.pop('o2'))
     # test#1: non-0deg nodes
     nodes = [1, 2, 9]
     _pull_nodes(nodes)

tests/compute/test_subgraph.py

 import numpy as np
+import networkx as nx
+import unittest
+import scipy.sparse as ssp
+
 import dgl
 import backend as F
-import unittest
+from test_utils import parametrize_dtype
 
 D = 5
 
@@ -80,3 +84,293 @@ def _test_map_to_subgraph():
     h = g.subgraph([0, 1, 2, 5, 8])
     v = h.map_to_subgraph_nid([0, 8, 2])
     assert np.array_equal(F.asnumpy(v), np.array([0, 4, 2]))
+
+def create_test_heterograph(idtype):
+    # test heterograph from the docstring, plus a user -- wishes -- game relation
+    # 3 users, 2 games, 2 developers
+    # metagraph:
+    #    ('user', 'follows', 'user'),
+    #    ('user', 'plays', 'game'),
+    #    ('user', 'wishes', 'game'),
+    #    ('developer', 'develops', 'game')])
+
+    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', idtype=idtype, device=F.ctx())
+    plays_g = dgl.bipartite(plays_spmat, 'user', 'plays', 'game', idtype=idtype, device=F.ctx())
+    wishes_g = dgl.bipartite(wishes_nx, 'user', 'wishes', 'game', idtype=idtype, device=F.ctx())
+    develops_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game', idtype=idtype, device=F.ctx())
+    assert follows_g.idtype == idtype
+    assert plays_g.idtype == idtype
+    assert wishes_g.idtype == idtype
+    assert develops_g.idtype == idtype
+    g = dgl.hetero_from_relations([follows_g, plays_g, wishes_g, develops_g])
+    assert g.idtype == idtype
+    assert g.device == F.ctx()
+    return g
+
+@unittest.skipIf(dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support bool tensor")
+@parametrize_dtype
+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_dtype
+def test_subgraph1(idtype):
+    g = create_test_heterograph(idtype)
+    g_graph = g['follows']
+    g_bipartite = g['plays']
+
+    x = F.randn((3, 5))
+    y = F.randn((2, 4))
+    g.nodes['user'].data['h'] = x
+    g.edges['follows'].data['h'] = y
+
+    def _check_subgraph(g, sg):
+        assert sg.idtype == g.idtype
+        assert sg.device == g.device
+        assert sg.ntypes == g.ntypes
+        assert sg.etypes == g.etypes
+        assert sg.canonical_etypes == g.canonical_etypes
+        assert F.array_equal(F.tensor(sg.nodes['user'].data[dgl.NID]),
+                             F.tensor([1, 2], g.idtype))
+        assert F.array_equal(F.tensor(sg.nodes['game'].data[dgl.NID]),
+                             F.tensor([0], g.idtype))
+        assert F.array_equal(F.tensor(sg.edges['follows'].data[dgl.EID]),
+                             F.tensor([1], g.idtype))
+        assert F.array_equal(F.tensor(sg.edges['plays'].data[dgl.EID]),
+                             F.tensor([1], g.idtype))
+        assert F.array_equal(F.tensor(sg.edges['wishes'].data[dgl.EID]),
+                             F.tensor([1], g.idtype))
+        assert sg.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], 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 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)
+
+    # Test for restricted format
+    if F._default_context_str != 'gpu':
+        # TODO(minjie): enable this later
+        for fmt in ['csr', 'csc', 'coo']:
+            g = dgl.graph([(0, 1), (1, 2)]).formats(fmt)
+            sg = g.subgraph({g.ntypes[0]: [1, 0]})
+            nids = F.asnumpy(sg.ndata[dgl.NID])
+            assert np.array_equal(nids, np.array([1, 0]))
+            src, dst = sg.edges(order='eid')
+            src = F.asnumpy(src)
+            dst = F.asnumpy(dst)
+            assert np.array_equal(src, np.array([1]))
+
+
+@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
+@parametrize_dtype
+def test_in_subgraph(idtype):
+    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
+    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
+    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
+    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
+    hg = dgl.hetero_from_relations([g1, g2, g3, g4])
+    subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0})
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+    u, v = subg['follow'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(hg['follow'].edge_ids(u, v), subg['follow'].edata[dgl.EID])
+    assert edge_set == {(1,0),(2,0),(3,0),(0,1),(2,1),(3,1)}
+    u, v = subg['play'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(hg['play'].edge_ids(u, v), subg['play'].edata[dgl.EID])
+    assert edge_set == {(0,0)}
+    u, v = subg['liked-by'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert F.array_equal(hg['liked-by'].edge_ids(u, v), subg['liked-by'].edata[dgl.EID])
+    assert edge_set == {(2,0),(2,1),(1,0),(0,0)}
+    assert subg['flips'].number_of_edges() == 0
+
+@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
+@parametrize_dtype
+def test_out_subgraph(idtype):
+    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
+    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
+    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
+    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
+    hg = dgl.hetero_from_relations([g1, g2, g3, g4])
+    subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0})
+    assert subg.idtype == idtype
+    assert len(subg.ntypes) == 3
+    assert len(subg.etypes) == 4
+    u, v = subg['follow'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(1,0),(0,1),(0,2)}
+    assert F.array_equal(hg['follow'].edge_ids(u, v), subg['follow'].edata[dgl.EID])
+    u, v = subg['play'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0,0),(0,1),(1,2)}
+    assert F.array_equal(hg['play'].edge_ids(u, v), subg['play'].edata[dgl.EID])
+    u, v = subg['liked-by'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0,0)}
+    assert F.array_equal(hg['liked-by'].edge_ids(u, v), subg['liked-by'].edata[dgl.EID])
+    u, v = subg['flips'].edges()
+    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
+    assert edge_set == {(0,0),(1,0)}
+    assert F.array_equal(hg['flips'].edge_ids(u, v), subg['flips'].edata[dgl.EID])

tests/compute/test_transform.py

@@ -668,61 +668,6 @@ def test_reorder_nodes():
         old_neighs2 = g.predecessors(old_nid)
         assert np.all(np.sort(old_neighs1) == np.sort(F.asnumpy(old_neighs2)))
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
-@parametrize_dtype
-def test_in_subgraph(idtype):
-    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
-    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
-    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
-    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
-    hg = dgl.hetero_from_relations([g1, g2, g3, g4])
-    subg = dgl.in_subgraph(hg, {'user' : [0,1], 'game' : 0})
-    assert subg.idtype == idtype
-    assert len(subg.ntypes) == 3
-    assert len(subg.etypes) == 4
-    u, v = subg['follow'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert F.array_equal(hg['follow'].edge_ids(u, v), subg['follow'].edata[dgl.EID])
-    assert edge_set == {(1,0),(2,0),(3,0),(0,1),(2,1),(3,1)}
-    u, v = subg['play'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert F.array_equal(hg['play'].edge_ids(u, v), subg['play'].edata[dgl.EID])
-    assert edge_set == {(0,0)}
-    u, v = subg['liked-by'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert F.array_equal(hg['liked-by'].edge_ids(u, v), subg['liked-by'].edata[dgl.EID])
-    assert edge_set == {(2,0),(2,1),(1,0),(0,0)}
-    assert subg['flips'].number_of_edges() == 0
-
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
-@parametrize_dtype
-def test_out_subgraph(idtype):
-    g1 = dgl.graph([(1,0),(2,0),(3,0),(0,1),(2,1),(3,1),(0,2)], 'user', 'follow', idtype=idtype)
-    g2 = dgl.bipartite([(0,0),(0,1),(1,2),(3,2)], 'user', 'play', 'game', idtype=idtype)
-    g3 = dgl.bipartite([(2,0),(2,1),(2,2),(1,0),(1,3),(0,0)], 'game', 'liked-by', 'user', idtype=idtype)
-    g4 = dgl.bipartite([(0,0),(1,0),(2,0),(3,0)], 'user', 'flips', 'coin', idtype=idtype)
-    hg = dgl.hetero_from_relations([g1, g2, g3, g4])
-    subg = dgl.out_subgraph(hg, {'user' : [0,1], 'game' : 0})
-    assert subg.idtype == idtype
-    assert len(subg.ntypes) == 3
-    assert len(subg.etypes) == 4
-    u, v = subg['follow'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert edge_set == {(1,0),(0,1),(0,2)}
-    assert F.array_equal(hg['follow'].edge_ids(u, v), subg['follow'].edata[dgl.EID])
-    u, v = subg['play'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert edge_set == {(0,0),(0,1),(1,2)}
-    assert F.array_equal(hg['play'].edge_ids(u, v), subg['play'].edata[dgl.EID])
-    u, v = subg['liked-by'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert edge_set == {(0,0)}
-    assert F.array_equal(hg['liked-by'].edge_ids(u, v), subg['liked-by'].edata[dgl.EID])
-    u, v = subg['flips'].edges()
-    edge_set = set(zip(list(F.asnumpy(u)), list(F.asnumpy(v))))
-    assert edge_set == {(0,0),(1,0)}
-    assert F.array_equal(hg['flips'].edge_ids(u, v), subg['flips'].edata[dgl.EID])
-
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU compaction not implemented")
 @parametrize_dtype
 def test_compact(idtype):