[NN] Support Unidirectional Bipartite Graphs in CFConv (dmlc#2674)

* Update * update * Update Co-authored-by: Jinjing Zhou <[email protected]>
feixian15 · Mar 3, 2021 · 7380d61 · 7380d61
1 parent 97bdae9
commit 7380d61
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/python/dgl/nn/pytorch/conv/cfconv.py b/python/dgl/nn/pytorch/conv/cfconv.py
@@ -123,15 +123,23 @@ def forward(self, g, node_feats, edge_feats):
         ----------
         g : DGLGraph
             The graph.
-        node_feats : float32 tensor of shape (V, node_in_feats)
-            Input node features, V for the number of nodes.
-        edge_feats : float32 tensor of shape (E, edge_in_feats)
-            Input edge features, E for the number of edges.
+        node_feats : torch.Tensor or pair of torch.Tensor
+            The input node features. If a torch.Tensor is given, it represents the input
+            node feature of shape :math:`(N, D_{in})` where :math:`D_{in}` is size of
+            input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, which is the case for bipartite graph,
+            the pair must contain two tensors of shape :math:`(N_{src}, D_{in_{src}})` and
+            :math:`(N_{dst}, D_{in_{dst}})` separately for the source and destination nodes.
+
+        edge_feats : torch.Tensor
+            The input edge feature of shape :math:`(E, edge_in_feats)`
+            where :math:`E` is the number of edges.
 
         Returns
         -------
-        float32 tensor of shape (V, out_feats)
-            Updated node representations.
+        torch.Tensor
+            The output node feature of shape :math:`(N_{out}, out_feats)`
+            where :math:`N_{out}` is the number of destination nodes.
         """
         with g.local_scope():
             if isinstance(node_feats, tuple):

diff --git a/tests/pytorch/test_nn.py b/tests/pytorch/test_nn.py
@@ -923,7 +923,7 @@ def test_atomic_conv(g, idtype):
     assert h.shape[-1] == 4
 
 @parametrize_dtype
-@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
+@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite'], exclude=['zero-degree']))
 @pytest.mark.parametrize('out_dim', [1, 3])
 def test_cf_conv(g, idtype, out_dim):
     g = g.astype(idtype).to(F.ctx())
@@ -936,9 +936,15 @@ def test_cf_conv(g, idtype, out_dim):
     if F.gpu_ctx():
         cfconv = cfconv.to(ctx)
 
-    node_feats = F.randn((g.number_of_nodes(), 2))
+    src_feats = F.randn((g.number_of_src_nodes(), 2))
     edge_feats = F.randn((g.number_of_edges(), 3))
-    h = cfconv(g, node_feats, edge_feats)
+    h = cfconv(g, src_feats, edge_feats)
+    # current we only do shape check
+    assert h.shape[-1] == out_dim
+
+    # case for bipartite graphs
+    dst_feats = F.randn((g.number_of_dst_nodes(), 3))
+    h = cfconv(g, (src_feats, dst_feats), edge_feats)
     # current we only do shape check
     assert h.shape[-1] == out_dim