Pytorch_geometric: RGCN edge_type argument

[cgoliver]

Hello,

Just wondering if there are any additional details on the RGCN class, namely the edge_type argument of the forward() method (https://rusty1s.github.io/pytorch_geometric/build/html/modules/nn.html#torch_geometric.nn.conv.RGCNConv).

I'd like to know what the shape of the tensor should be and what it should reflect. My guess is that it should be an 1-hot index over the edge labels for each edge. However, when using that I get the following error:

e.g. For 18 edge features and 8 edges I have

edge_index (torch.Size([2, 8])): 

tensor([[0, 0, 1, 1, 2, 2, 3, 3],                                                                                                   
        [1, 2, 0, 3, 0, 3, 1, 2]])

edge_type (torch.Size([8, 18])): 

tensor([[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

File "/anaconda3/lib/python3.6/site-packages/torch_geometric/nn/conv/message_passing.py", line 55, in propagate
    out = self.message(*message_args)
  File "/anaconda3/lib/python3.6/site-packages/torch_geometric/nn/conv/rgcn_conv.py", line 82, in message
    out = torch.bmm(x_j.unsqueeze(1), w).squeeze(-2)
RuntimeError: Expected 3-dimensional tensor, but got 4-dimensional tensor for argument #2 'batch2' (while checking arguments for bmm)

Thanks very much!

[rusty1s]

Hi,
The edge_type needs to be a one-dimensional long tensor which holds the relation id for each edge. You can verify this in this example. I will add a note to its documentation. Thanks!