TransformerConv + to_hetero + varying shape edge_attr

[wgeul]

Hi,

I'd like implement a heterogeneous network with TransformerConv layers, but with edges that have varying feature dimensions. I found a code example and added my specification below, changing the edge dimensions from 3 to 4 for one edge type:

from torch_geometric.nn import TransformerConv
from torch_geometric.nn import to_hetero
import torch
import torch.nn as nn
from torch_geometric.datasets import FakeHeteroDataset


dataset = FakeHeteroDataset(1, edge_dim=3)
data = dataset[0]

# Change shape from (12039, 3) to (12039, 4)
data[("v1", "e0", "v2")].edge_attr = torch.rand((12039, 4))


class TransformerModel(nn.Module):
    def __init__(self, num_layers, edge_dim):
        super(TransformerModel, self).__init__()
        self.convlist = nn.ModuleList()
        for i in range(0, num_layers):
            conv = TransformerConv(
                (-1, -1), 16, edge_dim=edge_dim, beta=True, dropout=0.1
            )
            self.convlist.append(conv)
        self.linear_layer = nn.Linear(16, 1)

    def forward(self, x, edge_index, edge_attr):
        for i, _ in enumerate(self.convlist):
            x = self.convlist[i](x, edge_index, edge_attr)
        x = self.linear_layer(x)
        return x


transformer_model = TransformerModel(2, 3)
transformer_model = to_hetero(transformer_model, data.metadata(), aggr="sum")

out_transformer = transformer_model(
    data.x_dict, data.edge_index_dict, data.edge_attr_dict
)
print("finished")

I received this error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (12039x4 and 3x16)

Is it possible to infer edge_dim for each edge type?

[wgeul]

After some browsing in the discussions, I found out that this can be achieved with edge_dim=-1.

Another requirement is that some edges contain no edge_attr. If I apply edge_dim=-1 I receive an Assertion error:

--> 210             assert edge_attr is not None
    211             edge_attr = self.lin_edge(edge_attr).view(-1, self.heads,
    212                                                       self.out_channels)

AssertionError: 

Is the only solution to pass e.g. edge_attr filled with zeros for such cases?

[rusty1s]

to_hetero currently assumes that all edge types have the same information. Your workaround should definitely work.