Repository: dmlc/dgl

Issue:
relgraphconv basis_message_func bug
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce
Use a small graph where #nodes * num_classes < hidden dimension of the graph convolutional layer. And use nodeID as the input feature. The bug lies in the basis_message_func function.

Steps to reproduce the behavior:

1. Construct a small graph (example given below)
2. Use node ID as the input feature, i.e., the same setup as the examples provided for the RGCN.
3. Run RGCN

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
Any small graph with #nodes * num_classes < hiddden_dim will trigger the bug. An example graph is given below:
# load graph data
    hg = dgl.heterograph({
         ('pax', 'ios_wifi_bssis', 'pax'): ([0, 1, 0, 1, 2], [1, 0, 2, 3, 3]),
         ('pax', 'adr_mac_addr', 'pax'): ([0, 2, 2, 3], [2, 3, 4, 4])
         })
    hg = dgl.to_bidirected(hg)
    num_rels = len(hg.canonical_etypes)
    
    # calculate norm for each edge type and store in edge
    for canonical_etype in hg.canonical_etypes:
        u, v, eid = hg.all_edges(form='all', etype=canonical_etype)
        _, inverse_index, count = torch.unique(v, return_inverse=True, return_counts=True)
        degrees = count[inverse_index]
        norm = torch.ones(eid.shape[0]).float() / degrees.float()
        norm = norm.unsqueeze(1)
        hg.edges[canonical_etype].data['norm'] = norm
    g = dgl.to_homogeneous(hg, edata=['norm'])
    
    train_mask = torch.tensor([True, True, True, True, False])
    test_mask = torch.tensor([False, False, False, True, True])
    labels = torch.tensor([1, 1, 1, 1, 1])
    train_idx = torch.nonzero(train_mask, as_tuple=False).squeeze()
    test_idx = torch.nonzero(test_mask, as_tuple=False).squeeze()

    # split dataset into train, validate, test
    if args.validation:
        val_idx = train_idx[:len(train_idx) // 5]
        train_idx = train_idx[len(train_idx) // 5:]
    else:
        val_idx = train_idx

    num_classes = 2
    num_nodes = g.number_of_nodes()
    edge_norm = g.edata['norm']
    edge_type = g.edata[dgl.ETYPE].long()
    in_feat = 3

    # since the nodes are featureless, the input feature is then the node id.
    feats = torch.arange(num_nodes)
    in_feat = num_nodes

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
The code failed by running on a small graph because of a bug in the code. 
Suggestion for code change in the basis_message_func function:

#Original code snippet: 
if h.dtype == th.int64 and h.ndim == 1:
            if isinstance(etypes, list):
                etypes = th.repeat_interleave(th.arange(len(etypes), device=device),
                                              th.tensor(etypes, device=device))
            weight = weight.view(-1, weight.shape[2])
            flatidx = etypes * weight.shape[1] + h
            msg = weight.index_select(0, flatidx)

#Suggested code snippet:
if h.dtype == th.int64 and h.ndim == 1:
            if isinstance(etypes, list):
                etypes = th.repeat_interleave(th.arange(len(etypes), device=device),
                                              th.tensor(etypes, device=device))
            **idim = weight.shape[1]**
            weight = weight.view(-1, weight.shape[2])
            **flatidx = etypes * idim + h**
            msg = weight.index_select(0, flatidx)

## Environment

 - DGL Version (e.g., 1.0): 0.6.1
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.9.0 CPU
 - OS (e.g., Linux): Linux
 - How you installed DGL (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): 
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
This is embarrassing.  Would you like to submit a PR or you want us to do it?
Hi @BarclayII , I do not have permission to do a PR. Please go ahead and fix it. 

Resolving patch:
diff --git a/python/dgl/nn/pytorch/conv/relgraphconv.py b/python/dgl/nn/pytorch/conv/relgraphconv.py
--- a/python/dgl/nn/pytorch/conv/relgraphconv.py
+++ b/python/dgl/nn/pytorch/conv/relgraphconv.py
@@ -218,8 +218,9 @@ def basis_message_func(self, edges, etypes):
             if isinstance(etypes, list):
                 etypes = th.repeat_interleave(th.arange(len(etypes), device=device),
                                               th.tensor(etypes, device=device))
+            idim = weight.shape[1]
             weight = weight.view(-1, weight.shape[2])
-            flatidx = etypes * weight.shape[1] + h
+            flatidx = etypes * idim + h
             msg = weight.index_select(0, flatidx)
         elif self.low_mem:
             # A more memory-friendly implementation.


Test patch: