Repository: dmlc/dgl

Issue:
canonical_etypes should be used in HeteroGraphConv
## 🐛 Bug

Trying to train this kind of graphs:
```
metagraph=[('tweet', 'url', 'links'), ('tweet', 'user', 'mentions'), ('tweet', 'hashtag', 'tags'), ('user', 'url', 'links'), ('user', 'user', 'mentions'), ('user', 'hashtag', 'tags'), ('user', 'tweet', 'tweets')])
```
leads to error:

dgl._ffi.base.DGLError: Edge type "links" is ambiguous. Please use canonical edge type in the form of (srctype, etype, dsttype)

but HeteroGraphConv has as input:

mods : dict[str, nn.Module]
        Modules associated with every edge types. 

So canonical types cannot be used.

It seems to be impossible to use HeteroGraphConv in case the same relation is used between different types of nodes.

see also #3435
## To Reproduce

Steps to reproduce the behavior:

1. Create a graph with the same relation name between more than one couple of nodes, such as ('tweet', 'url', 'links'), ('user', 'url', 'links'),
2. Create a heterograph with something like
```
HeteroGraphConv({
            name: GraphConv(in_size, out_size) for name in etypes
        })
```
3. Call the forward function of HeteroGraphConv

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

It should be possible to specify canonical etypes as an input to HeteroGraphConv, and therefore remove the ambiguity

## Environment

 - DGL Version (e.g., 1.0): dgl==0.6.0
## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
I also encounterd this error just a few minutes ago. 
I think there are two use cases:
* Use the same NN module for `('tweet', 'url', 'links')` and `('user', 'url', 'links')` as a way of parameter sharing/regularization.
* Use different NN modules for `('tweet', 'url', 'links')` and `('user', 'url', 'links')`.

I believe both use cases are valid and we should extend `HeteroGraphConv` to support them. I've marked this as a feature request to be implemented.

Resolving patch:
diff --git a/python/dgl/nn/pytorch/hetero.py b/python/dgl/nn/pytorch/hetero.py
--- a/python/dgl/nn/pytorch/hetero.py
+++ b/python/dgl/nn/pytorch/hetero.py
@@ -122,7 +122,13 @@ def my_agg_func(tensors, dsttype):
     """
     def __init__(self, mods, aggregate='sum'):
         super(HeteroGraphConv, self).__init__()
+        self.mod_dict = mods
+        mods = {str(k): v for k, v in mods.items()}
+        # Register as child modules
         self.mods = nn.ModuleDict(mods)
+        # PyTorch ModuleDict doesn't have get() method, so I have to store two
+        # dictionaries so that I can index with both canonical edge type and
+        # edge type with the get() method.
         # Do not break if graph has 0-in-degree nodes.
         # Because there is no general rule to add self-loop for heterograph.
         for _, v in self.mods.items():
@@ -134,6 +140,16 @@ def __init__(self, mods, aggregate='sum'):
         else:
             self.agg_fn = aggregate
 
+    def _get_module(self, etype):
+        mod = self.mod_dict.get(etype, None)
+        if mod is not None:
+            return mod
+        if isinstance(etype, tuple):
+            # etype is canonical
+            _, etype, _ = etype
+            return self.mod_dict[etype]
+        raise KeyError('Cannot find module with edge type %s' % etype)
+
     def forward(self, g, inputs, mod_args=None, mod_kwargs=None):
         """Forward computation
 
@@ -171,7 +187,7 @@ def forward(self, g, inputs, mod_args=None, mod_kwargs=None):
                 rel_graph = g[stype, etype, dtype]
                 if stype not in src_inputs or dtype not in dst_inputs:
                     continue
-                dstdata = self.mods[etype](
+                dstdata = self._get_module((stype, etype, dtype))(
                     rel_graph,
                     (src_inputs[stype], dst_inputs[dtype]),
                     *mod_args.get(etype, ()),
@@ -182,7 +198,7 @@ def forward(self, g, inputs, mod_args=None, mod_kwargs=None):
                 rel_graph = g[stype, etype, dtype]
                 if stype not in inputs:
                     continue
-                dstdata = self.mods[etype](
+                dstdata = self._get_module((stype, etype, dtype))(
                     rel_graph,
                     (inputs[stype], inputs[dtype]),
                     *mod_args.get(etype, ()),


Test patch:
diff --git a/tests/pytorch/test_nn.py b/tests/pytorch/test_nn.py
--- a/tests/pytorch/test_nn.py
+++ b/tests/pytorch/test_nn.py
@@ -1146,17 +1146,26 @@ def myagg(alist, dsttype):
 
 @parametrize_idtype
 @pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg])
-def test_hetero_conv(agg, idtype):
+@pytest.mark.parametrize('canonical_keys', [False, True])
+def test_hetero_conv(agg, idtype, canonical_keys):
     g = dgl.heterograph({
         ('user', 'follows', 'user'): ([0, 0, 2, 1], [1, 2, 1, 3]),
         ('user', 'plays', 'game'): ([0, 0, 0, 1, 2], [0, 2, 3, 0, 2]),
         ('store', 'sells', 'game'): ([0, 0, 1, 1], [0, 3, 1, 2])},
         idtype=idtype, device=F.ctx())
-    conv = nn.HeteroGraphConv({
-        'follows': nn.GraphConv(2, 3, allow_zero_in_degree=True),
-        'plays': nn.GraphConv(2, 4, allow_zero_in_degree=True),
-        'sells': nn.GraphConv(3, 4, allow_zero_in_degree=True)},
-        agg)
+    if not canonical_keys:
+        conv = nn.HeteroGraphConv({
+            'follows': nn.GraphConv(2, 3, allow_zero_in_degree=True),
+            'plays': nn.GraphConv(2, 4, allow_zero_in_degree=True),
+            'sells': nn.GraphConv(3, 4, allow_zero_in_degree=True)},
+            agg)
+    else:
+        conv = nn.HeteroGraphConv({
+            ('user', 'follows', 'user'): nn.GraphConv(2, 3, allow_zero_in_degree=True),
+            ('user', 'plays', 'game'): nn.GraphConv(2, 4, allow_zero_in_degree=True),
+            ('store', 'sells', 'game'): nn.GraphConv(3, 4, allow_zero_in_degree=True)},
+            agg)
+
     conv = conv.to(F.ctx())
 
     # test pickle
@@ -1621,4 +1630,4 @@ def test_dgn_conv(in_size, out_size, aggregators, scalers, delta,
     aggregators_non_eig = [aggr for aggr in aggregators if not aggr.startswith('dir')]
     model = nn.DGNConv(in_size, out_size, aggregators_non_eig, scalers, delta, dropout,
         num_towers, edge_feat_size, residual).to(dev)
-    model(g, h, edge_feat=e)
\ No newline at end of file
+    model(g, h, edge_feat=e)