[Bugfix] Fix UVA sampling with partially specified node types (#3897)

* fix uva with partial node types

* lint

* skip tensorflow unit test

python/dgl/subgraph.py

@@ -10,7 +10,6 @@ from .base import DGLError, dgl_warning
 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
 from .utils import recursive_apply, context_of
@@ -142,12 +141,12 @@ def node_subgraph(graph, nodes, *, relabel_nodes=True, store_ids=True, output_de
             return F.astype(F.nonzero_1d(F.copy_to(v, graph.device)), graph.idtype)
         else:
             return utils.prepare_tensor(graph, v, 'nodes["{}"]'.format(ntype))
+    nodes = {ntype: _process_nodes(ntype, v) for ntype, v in nodes.items()}
+    device = context_of(nodes)
 
-    induced_nodes = []
-    for ntype in graph.ntypes:
-        nids = nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), graph.device))
-        induced_nodes.append(_process_nodes(ntype, nids))
-    device = context_of(induced_nodes)
+    induced_nodes = [
+        nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device))
+        for ntype in graph.ntypes]
     sgi = graph._graph.node_subgraph(induced_nodes, relabel_nodes)
     induced_edges = sgi.induced_edges
     # (BarclayII) should not write induced_nodes = sgi.induced_nodes due to the same
@@ -301,13 +300,13 @@ def edge_subgraph(graph, edges, *, relabel_nodes=True, store_ids=True, output_de
             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 = []
-    for cetype in graph.canonical_etypes:
-        eids = edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), graph.device))
-        induced_edges.append(_process_edges(cetype, eids))
-    device = context_of(induced_edges)
+    edges = {etype: _process_edges(etype, e) for etype, e in edges.items()}
+    device = context_of(edges)
+    induced_edges = [
+        edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), device))
+        for cetype in graph.canonical_etypes]
+
     sgi = graph._graph.edge_subgraph(induced_edges, not relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
     subg = _create_hetero_subgraph(
@@ -430,12 +429,9 @@ def in_subgraph(graph, nodes, *, relabel_nodes=False, store_ids=True, output_dev
         nodes = {graph.ntypes[0] : nodes}
     nodes = utils.prepare_tensor_dict(graph, nodes, 'nodes')
     device = context_of(nodes)
-    nodes_all_types = []
-    for ntype in graph.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[graph._idtype_str])
+    nodes_all_types = [
+        F.to_dgl_nd(nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device)))
+        for ntype in graph.ntypes]
 
     sgi = _CAPI_DGLInSubgraph(graph._graph, nodes_all_types, relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
@@ -560,12 +556,9 @@ def out_subgraph(graph, nodes, *, relabel_nodes=False, store_ids=True, output_de
         nodes = {graph.ntypes[0] : nodes}
     nodes = utils.prepare_tensor_dict(graph, nodes, 'nodes')
     device = context_of(nodes)
-    nodes_all_types = []
-    for ntype in graph.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[graph._idtype_str])
+    nodes_all_types = [
+        F.to_dgl_nd(nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device)))
+        for ntype in graph.ntypes]
 
     sgi = _CAPI_DGLOutSubgraph(graph._graph, nodes_all_types, relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
@@ -693,7 +686,8 @@ def khop_in_subgraph(graph, nodes, k, *, relabel_nodes=True, store_ids=True, out
 
     last_hop_nodes = nodes
     k_hop_nodes_ = [last_hop_nodes]
-    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), graph.device)
+    device = context_of(nodes)
+    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), device)
     for _ in range(k):
         current_hop_nodes = {nty: [] for nty in graph.ntypes}
         for cetype in graph.canonical_etypes:
@@ -853,7 +847,8 @@ def khop_out_subgraph(graph, nodes, k, *, relabel_nodes=True, store_ids=True, ou
 
     last_hop_nodes = nodes
     k_hop_nodes_ = [last_hop_nodes]
-    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), graph.device)
+    device = context_of(nodes)
+    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), device)
     for _ in range(k):
         current_hop_nodes = {nty: [] for nty in graph.ntypes}
         for cetype in graph.canonical_etypes:

src/graph/subgraph.cc

@@ -13,11 +13,12 @@ HeteroSubgraph InEdgeGraphRelabelNodes(
   CHECK_EQ(vids.size(), graph->NumVertexTypes())
     << "Invalid input: the input list size must be the same as the number of vertex types.";
   std::vector<IdArray> eids(graph->NumEdgeTypes());
+  DLContext ctx = aten::GetContextOf(vids);
   for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
     auto pair = graph->meta_graph()->FindEdge(etype);
     const dgl_type_t dst_vtype = pair.second;
     if (aten::IsNullArray(vids[dst_vtype])) {
-      eids[etype] = IdArray::Empty({0}, graph->DataType(), graph->Context());
+      eids[etype] = IdArray::Empty({0}, graph->DataType(), ctx);
     } else {
       const auto& earr = graph->InEdges(etype, {vids[dst_vtype]});
       eids[etype] = earr.id;
@@ -33,6 +34,7 @@ HeteroSubgraph InEdgeGraphNoRelabelNodes(
     << "Invalid input: the input list size must be the same as the number of vertex types.";
   std::vector<HeteroGraphPtr> subrels(graph->NumEdgeTypes());
   std::vector<IdArray> induced_edges(graph->NumEdgeTypes());
+  DLContext ctx = aten::GetContextOf(vids);
   for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
     auto pair = graph->meta_graph()->FindEdge(etype);
     const dgl_type_t src_vtype = pair.first;
@@ -44,7 +46,7 @@ HeteroSubgraph InEdgeGraphNoRelabelNodes(
         relgraph->NumVertexTypes(),
         graph->NumVertices(src_vtype),
         graph->NumVertices(dst_vtype),
-        graph->DataType(), graph->Context());
+        graph->DataType(), ctx);
       induced_edges[etype] = IdArray::Empty({0}, graph->DataType(), graph->Context());
     } else {
       const auto& earr = graph->InEdges(etype, {vids[dst_vtype]});
@@ -77,11 +79,12 @@ HeteroSubgraph OutEdgeGraphRelabelNodes(
   CHECK_EQ(vids.size(), graph->NumVertexTypes())
     << "Invalid input: the input list size must be the same as the number of vertex types.";
   std::vector<IdArray> eids(graph->NumEdgeTypes());
+  DLContext ctx = aten::GetContextOf(vids);
   for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
     auto pair = graph->meta_graph()->FindEdge(etype);
     const dgl_type_t src_vtype = pair.first;
     if (aten::IsNullArray(vids[src_vtype])) {
-      eids[etype] = IdArray::Empty({0}, graph->DataType(), graph->Context());
+      eids[etype] = IdArray::Empty({0}, graph->DataType(), ctx);
     } else {
       const auto& earr = graph->OutEdges(etype, {vids[src_vtype]});
       eids[etype] = earr.id;
@@ -97,6 +100,7 @@ HeteroSubgraph OutEdgeGraphNoRelabelNodes(
     << "Invalid input: the input list size must be the same as the number of vertex types.";
   std::vector<HeteroGraphPtr> subrels(graph->NumEdgeTypes());
   std::vector<IdArray> induced_edges(graph->NumEdgeTypes());
+  DLContext ctx = aten::GetContextOf(vids);
   for (dgl_type_t etype = 0; etype < graph->NumEdgeTypes(); ++etype) {
     auto pair = graph->meta_graph()->FindEdge(etype);
     const dgl_type_t src_vtype = pair.first;
@@ -108,7 +112,7 @@ HeteroSubgraph OutEdgeGraphNoRelabelNodes(
         relgraph->NumVertexTypes(),
         graph->NumVertices(src_vtype),
         graph->NumVertices(dst_vtype),
-        graph->DataType(), graph->Context());
+        graph->DataType(), ctx);
       induced_edges[etype] = IdArray::Empty({0}, graph->DataType(), graph->Context());
     } else {
       const auto& earr = graph->OutEdges(etype, {vids[src_vtype]});

src/graph/unit_graph.cc

@@ -381,7 +381,8 @@ class UnitGraph::COO : public BaseHeteroGraph {
     CHECK(aten::IsValidIdArray(dstvids)) << "Invalid vertex id array.";
     HeteroSubgraph subg;
     const auto& submat = aten::COOSliceMatrix(adj_, srcvids, dstvids);
-    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
+    DLContext ctx = aten::GetContextOf(vids);
+    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), ctx);
     subg.graph = std::make_shared<COO>(meta_graph(), submat.num_rows, submat.num_cols,
         submat.row, submat.col);
     subg.induced_vertices = vids;
@@ -801,7 +802,8 @@ class UnitGraph::CSR : public BaseHeteroGraph {
     CHECK(aten::IsValidIdArray(dstvids)) << "Invalid vertex id array.";
     HeteroSubgraph subg;
     const auto& submat = aten::CSRSliceMatrix(adj_, srcvids, dstvids);
-    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), Context());
+    DLContext ctx = aten::GetContextOf(vids);
+    IdArray sub_eids = aten::Range(0, submat.data->shape[0], NumBits(), ctx);
     subg.graph = std::make_shared<CSR>(meta_graph(), submat.num_rows, submat.num_cols,
         submat.indptr, submat.indices, sub_eids);
     subg.induced_vertices = vids;

tests/compute/test_subgraph.py

@@ -100,6 +100,8 @@ def create_test_heterograph(idtype):
         ('user', 'wishes', 'game'): ([0, 2], [1, 0]),
         ('developer', 'develops', 'game'): ([0, 1], [0, 1])
     }, idtype=idtype, device=F.ctx())
+    for etype in g.etypes:
+        g.edges[etype].data['weight'] = F.randn((g.num_edges(etype),))
     assert g.idtype == idtype
     assert g.device == F.ctx()
     return g
@@ -629,6 +631,28 @@ def test_subframes(parent_idx_device, child_device):
     if parent_device == 'uva':
         g.unpin_memory_()
 
+@unittest.skipIf(F._default_context_str != "gpu", reason="UVA only available on GPU")
+@pytest.mark.parametrize('device', [F.cpu(), F.cuda()])
+@parametrize_dtype
+def test_uva_subgraph(idtype, device):
+    g = create_test_heterograph(idtype)
+    g = g.to(F.cpu())
+    g.create_formats_()
+    g.pin_memory_()
+    indices = {'user': F.copy_to(F.tensor([0], idtype), device)}
+    edge_indices = {'follows': F.copy_to(F.tensor([0], idtype), device)}
+    assert g.subgraph(indices).device == device
+    assert g.edge_subgraph(edge_indices).device == device
+    assert g.in_subgraph(indices).device == device
+    assert g.out_subgraph(indices).device == device
+    if dgl.backend.backend_name != 'tensorflow':
+        # (BarclayII) Most of Tensorflow functions somehow do not preserve device: a CPU tensor
+        # becomes a GPU tensor after operations such as concat(), unique() or even sin().
+        # Not sure what should be the best fix.
+        assert g.khop_in_subgraph(indices, 1)[0].device == device
+        assert g.khop_out_subgraph(indices, 1)[0].device == device
+    assert g.sample_neighbors(indices, 1).device == device
+    g.unpin_memory_()
+
 if __name__ == '__main__':
-    test_khop_out_subgraph(F.int64)
-    test_subframes(('cpu', F.cpu()), F.cuda())
+    test_uva_subgraph(F.int64, F.cpu())